From ceed3f655d16bce565d868efdbdbcaab629adcd7 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jul 2024 16:36:18 -0700
Subject: [PATCH] fix check if user passed custom header

---
 litellm/proxy/auth/user_api_key_auth.py | 33 ++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index c72c6e353a..dcd2cbb804 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -135,13 +135,14 @@ async def user_api_key_auth(
                         header_key: str = headers.get("litellm_user_api_key", "")
                         if request.headers.get(key=header_key) is not None:
                             api_key = request.headers.get(key=header_key)
+
+        # if user wants to pass LiteLLM_Master_Key as a custom header, example pass litellm keys as X-LiteLLM-Key: Bearer sk-1234
         custom_litellm_key_header_name = general_settings.get("litellm_key_header_name")
         if custom_litellm_key_header_name is not None:
-            # use this as the virtual key passed to litellm proxy
-            passed_api_key = (
-                request.headers.get(key=custom_litellm_key_header_name) or ""  # type: ignore
+            api_key = get_api_key_from_custom_header(
+                request=request,
+                custom_litellm_key_header_name=custom_litellm_key_header_name,
             )
-            api_key = _get_bearer_token(api_key=passed_api_key)
 
         parent_otel_span: Optional[Span] = None
         if open_telemetry_logger is not None:
@@ -1272,3 +1273,27 @@ def _check_valid_ip(allowed_ips: Optional[List[str]], request: Request) -> bool:
         return False
 
     return True
+
+
+def get_api_key_from_custom_header(
+    request: Request, custom_litellm_key_header_name: str
+):
+    # use this as the virtual key passed to litellm proxy
+    custom_litellm_key_header_name = custom_litellm_key_header_name.lower()
+    verbose_proxy_logger.debug(
+        "searching for custom_litellm_key_header_name= %s",
+        custom_litellm_key_header_name,
+    )
+    custom_api_key = request.headers.get(custom_litellm_key_header_name)
+    if custom_api_key:
+        api_key = _get_bearer_token(api_key=custom_api_key)
+        verbose_proxy_logger.debug(
+            "Found custom API key using header: {}, setting api_key={}".format(
+                custom_litellm_key_header_name, api_key
+            )
+        )
+    else:
+        raise ValueError(
+            f"No LiteLLM Virtual Key pass. Please set header={custom_litellm_key_header_name}: Bearer <api_key>"
+        )
+    return api_key