From 66b2b5fab9b5b87af1d27ceaca76e402064280c1 Mon Sep 17 00:00:00 2001
From: Sara Ghaemi <sara.ghaemi@telus.com>
Date: Tue, 7 May 2024 11:37:04 -0400
Subject: [PATCH] made audience optional and updated docs

---
 docs/my-website/docs/proxy/token_auth.md |  1 +
 litellm/proxy/auth/handle_jwt.py         | 12 ++++++++----
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/docs/my-website/docs/proxy/token_auth.md b/docs/my-website/docs/proxy/token_auth.md
index 81475951fee..e4772d70afa 100644
--- a/docs/my-website/docs/proxy/token_auth.md
+++ b/docs/my-website/docs/proxy/token_auth.md
@@ -17,6 +17,7 @@ This is a new feature, and subject to changes based on feedback.
 ### Step 1. Setup Proxy
 
 - `JWT_PUBLIC_KEY_URL`: This is the public keys endpoint of your OpenID provider. Typically it's `{openid-provider-base-url}/.well-known/openid-configuration/jwks`. For Keycloak it's `{keycloak_base_url}/realms/{your-realm}/protocol/openid-connect/certs`.
+- `JWT_AUDIENCE`: This is the audience used for decoding the JWT. If not set, the decode step will not verify the audience. 
 
 ```bash
 export JWT_PUBLIC_KEY_URL="" # "https://demo.duendesoftware.com/.well-known/openid-configuration/jwks"
diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py
index c12f48cc11f..606ff682813 100644
--- a/litellm/proxy/auth/handle_jwt.py
+++ b/litellm/proxy/auth/handle_jwt.py
@@ -156,6 +156,11 @@ class JWTHandler:
         return public_key
 
     async def auth_jwt(self, token: str) -> dict:
+        audience = os.getenv("JWT_AUDIENCE")
+        decode_options = None
+        if audience is None:
+            decode_options = {"verify_aud": False}
+        
         from jwt.algorithms import RSAAlgorithm
 
         header = jwt.get_unverified_header(token)
@@ -185,7 +190,8 @@ class JWTHandler:
                     token,
                     public_key_rsa,  # type: ignore
                     algorithms=["RS256"],
-                    options={"verify_aud": False},
+                    options=decode_options,
+                    audience=audience,
                 )
                 return payload
 
@@ -195,9 +201,6 @@ class JWTHandler:
             except Exception as e:
                 raise Exception(f"Validation fails: {str(e)}")
         elif public_key is not None and isinstance(public_key, str):
-            audience = os.getenv("JWT_AUDIENCE")
-            if audience is None:
-                raise Exception("Missing JWT Audience from environment.")
             try:
                 cert = x509.load_pem_x509_certificate(public_key.encode(), default_backend())
 
@@ -213,6 +216,7 @@ class JWTHandler:
                     key,
                     algorithms=["RS256"],
                     audience=audience,
+                    options=decode_options
                 )
                 return payload