Add token validation (#1186)

* update requirements

* add token validation

* fix typchecker

* fixing tests

* fix comment linke 404

---------

Co-authored-by: Matt Gotteiner <[email protected]>

Author: mattgotteiner

app/backend/core/authentication.py

@@ -7,8 +7,15 @@
 import aiohttp
 from azure.search.documents.aio import SearchClient
 from azure.search.documents.indexes.models import SearchIndex
+from jose import jwt
 from msal import ConfidentialClientApplication
 from msal.token_cache import TokenCache
+from tenacity import (
+    AsyncRetrying,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_random_exponential,
+)
 
 
 # AuthError is raised when the authentication token sent by the client UI cannot be parsed or there is an authentication error accessing the graph API
@@ -40,6 +47,15 @@ def __init__(
         self.client_app_id = client_app_id
         self.tenant_id = tenant_id
         self.authority = f"https://login.microsoftonline.com/{tenant_id}"
+        # Depending on if requestedAccessTokenVersion is 1 or 2, the issuer and audience of the token may be different
+        # See https://learn.microsoft.com/graph/api/resources/apiapplication
+        self.valid_issuers = [
+            f"https://sts.windows.net/{tenant_id}/",
+            f"https://login.microsoftonline.com/{tenant_id}/v2.0",
+        ]
+        self.valid_audiences = [f"api://{server_app_id}", str(server_app_id)]
+        # See https://learn.microsoft.com/entra/identity-platform/access-tokens#validate-the-issuer for more information on token validation
+        self.key_url = f"{self.authority}/discovery/v2.0/keys"
 
         if self.use_authentication:
             field_names = [field.name for field in search_index.fields] if search_index else []
@@ -182,6 +198,11 @@ async def get_auth_claims_if_enabled(self, headers: dict) -> dict[str, Any]:
             # The scope is set to the Microsoft Graph API, which may need to be called for more authorization information
             # https://learn.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-on-behalf-of-flow
             auth_token = AuthenticationHelper.get_token_auth_header(headers)
+            # Validate the token before use
+            await self.validate_access_token(auth_token)
+
+            # Use the on-behalf-of-flow to acquire another token for use with Microsoft Graph
+            # See https://learn.microsoft.com/entra/identity-platform/v2-oauth2-on-behalf-of-flow for more information
             graph_resource_access_token = self.confidential_client.acquire_token_on_behalf_of(
                 user_assertion=auth_token, scopes=["https://graph.microsoft.com/.default"]
             )
@@ -207,7 +228,6 @@ async def get_auth_claims_if_enabled(self, headers: dict) -> dict[str, Any]:
                 auth_claims["groups"] = await AuthenticationHelper.list_groups(graph_resource_access_token)
             return auth_claims
         except AuthError as e:
-            print(e.error)
             logging.exception("Exception getting authorization information - " + json.dumps(e.error))
             if self.require_access_control:
                 raise
@@ -237,3 +257,74 @@ async def check_path_auth(self, path: str, auth_claims: dict[str, Any], search_c
             break
 
         return allowed
+
+    # See https://github.com/Azure-Samples/ms-identity-python-on-behalf-of/blob/939be02b11f1604814532fdacc2c2eccd198b755/FlaskAPI/helpers/authorization.py#L44
+    async def validate_access_token(self, token: str):
+        """
+        Validate an access token is issued by Entra
+        """
+        jwks = None
+        async for attempt in AsyncRetrying(
+            retry=retry_if_exception_type(AuthError),
+            wait=wait_random_exponential(min=15, max=60),
+            stop=stop_after_attempt(5),
+        ):
+            with attempt:
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(url=self.key_url) as resp:
+                        resp_status = resp.status
+                        if resp_status in [500, 502, 503, 504]:
+                            raise AuthError(
+                                error=f"Failed to get keys info: {await resp.text()}", status_code=resp_status
+                            )
+                        jwks = await resp.json()
+
+        if not jwks or "keys" not in jwks:
+            raise AuthError({"code": "invalid_keys", "description": "Unable to get keys to validate auth token."}, 401)
+
+        rsa_key = None
+        issuer = None
+        audience = None
+        try:
+            unverified_header = jwt.get_unverified_header(token)
+            unverified_claims = jwt.get_unverified_claims(token)
+            issuer = unverified_claims.get("iss")
+            audience = unverified_claims.get("aud")
+            for key in jwks["keys"]:
+                if key["kid"] == unverified_header["kid"]:
+                    rsa_key = {"kty": key["kty"], "kid": key["kid"], "use": key["use"], "n": key["n"], "e": key["e"]}
+                    break
+        except Exception as exc:
+            raise AuthError(
+                {"code": "invalid_header", "description": "Unable to parse authorization token."}, 401
+            ) from exc
+        if not rsa_key:
+            raise AuthError({"code": "invalid_header", "description": "Unable to find appropriate key"}, 401)
+
+        if issuer not in self.valid_issuers:
+            raise AuthError(
+                {"code": "invalid_header", "description": f"Issuer {issuer} not in {','.join(self.valid_issuers)}"}, 401
+            )
+
+        if audience not in self.valid_audiences:
+            raise AuthError(
+                {
+                    "code": "invalid_header",
+                    "description": f"Audience {audience} not in {','.join(self.valid_audiences)}",
+                },
+                401,
+            )
+
+        try:
+            jwt.decode(token, rsa_key, algorithms=["RS256"], audience=audience, issuer=issuer)
+        except jwt.ExpiredSignatureError as jwt_expired_exc:
+            raise AuthError({"code": "token_expired", "description": "token is expired"}, 401) from jwt_expired_exc
+        except jwt.JWTClaimsError as jwt_claims_exc:
+            raise AuthError(
+                {"code": "invalid_claims", "description": "incorrect claims," "please check the audience and issuer"},
+                401,
+            ) from jwt_claims_exc
+        except Exception as exc:
+            raise AuthError(
+                {"code": "invalid_header", "description": "Unable to parse authorization token."}, 401
+            ) from exc

app/backend/requirements.in

@@ -3,6 +3,7 @@ quart
 quart-cors
 openai[datalib]>=1.3.7
 tiktoken
+tenacity
 azure-search-documents==11.4.0b11
 azure-storage-blob
 uvicorn
@@ -14,3 +15,5 @@ opentelemetry-instrumentation-requests
 opentelemetry-instrumentation-aiohttp-client
 msal
 azure-keyvault-secrets
+cryptography
+python-jose[cryptography]

app/backend/requirements.txt

@@ -67,16 +67,20 @@ click==8.1.7
     #   flask
     #   quart
     #   uvicorn
-cryptography==41.0.7
+cryptography==42.0.1
     # via
+    #   -r requirements.in
     #   azure-identity
     #   azure-storage-blob
     #   msal
     #   pyjwt
+    #   python-jose
 deprecated==1.2.14
     # via opentelemetry-api
 distro==1.8.0
     # via openai
+ecdsa==0.18.0
+    # via python-jose
 fixedint==0.1.6
     # via azure-monitor-opentelemetry-exporter
 flask==3.0.0
@@ -257,18 +261,22 @@ portalocker==2.8.2
     # via msal-extensions
 priority==2.0.0
     # via hypercorn
+pyasn1==0.5.1
+    # via
+    #   python-jose
+    #   rsa
 pycparser==2.21
     # via cffi
 pydantic==2.5.2
     # via openai
 pydantic-core==2.14.5
     # via pydantic
 pyjwt[crypto]==2.8.0
-    # via
-    #   msal
-    #   pyjwt
+    # via msal
 python-dateutil==2.8.2
     # via pandas
+python-jose[cryptography]==3.3.0
+    # via -r requirements.in
 pytz==2023.3.post1
     # via pandas
 quart==0.19.4
@@ -288,16 +296,21 @@ requests==2.31.0
     #   tiktoken
 requests-oauthlib==1.3.1
     # via msrest
+rsa==4.9
+    # via python-jose
 six==1.16.0
     # via
     #   azure-core
+    #   ecdsa
     #   isodate
     #   python-dateutil
 sniffio==1.3.0
     # via
     #   anyio
     #   httpx
     #   openai
+tenacity==8.2.3
+    # via -r requirements.in
 tiktoken==0.5.2
     # via -r requirements.in
 tqdm==4.66.1

app/mypy.ini

@@ -2,3 +2,6 @@
 
 [mypy-msal.*]
 ignore_missing_imports = True
+
+[mypy-jose.*]
+ignore_missing_imports = True

scripts/auth_init.py

@@ -109,6 +109,9 @@ def create_server_app_permission_setup_payload(server_app_id: str):
                     "type": "User",
                 }
             ],
+            # Required to match v2.0 OpenID configuration for token validation
+            # Learn more at https://learn.microsoft.com/entra/identity-platform/v2-protocols-oidc#find-your-apps-openid-configuration-document-uri
+            "requestedAccessTokenVersion": 2,
         },
         "requiredResourceAccess": [
             {

scripts/requirements.in

@@ -12,3 +12,5 @@ azure-keyvault-secrets
 Pillow
 PyMuPDF
 types-Pillow
+cryptography
+python-jose[cryptography]

scripts/requirements.txt

@@ -54,14 +54,18 @@ cffi==1.16.0
     # via cryptography
 charset-normalizer==3.3.2
     # via requests
-cryptography==41.0.7
+cryptography==42.0.1
     # via
+    #   -r requirements.in
     #   azure-identity
     #   azure-storage-blob
     #   msal
     #   pyjwt
+    #   python-jose
 distro==1.8.0
     # via openai
+ecdsa==0.18.0
+    # via python-jose
 frozenlist==1.4.0
     # via
     #   aiohttp
@@ -116,16 +120,18 @@ pillow==10.2.0
     # via -r requirements.in
 portalocker==2.8.2
     # via msal-extensions
+pyasn1==0.5.1
+    # via
+    #   python-jose
+    #   rsa
 pycparser==2.21
     # via cffi
 pydantic==2.5.2
     # via openai
 pydantic-core==2.14.5
     # via pydantic
 pyjwt[crypto]==2.8.0
-    # via
-    #   msal
-    #   pyjwt
+    # via msal
 pymupdf==1.23.7
     # via -r requirements.in
 pymupdfb==1.23.7
@@ -134,6 +140,8 @@ pypdf==3.17.1
     # via -r requirements.in
 python-dateutil==2.8.2
     # via pandas
+python-jose[cryptography]==3.3.0
+    # via -r requirements.in
 pytz==2023.3.post1
     # via pandas
 regex==2023.10.3
@@ -147,9 +155,12 @@ requests==2.31.0
     #   tiktoken
 requests-oauthlib==1.3.1
     # via msrest
+rsa==4.9
+    # via python-jose
 six==1.16.0
     # via
     #   azure-core
+    #   ecdsa
     #   isodate
     #   python-dateutil
 sniffio==1.3.0

tests/conftest.py

@@ -23,6 +23,7 @@
 from openai.types.create_embedding_response import Usage
 
 import app
+import core
 from core.authentication import AuthenticationHelper
 
 from .mocks import (
@@ -302,6 +303,7 @@ async def auth_client(
     mock_openai_chatcompletion,
     mock_openai_embedding,
     mock_confidential_client_success,
+    mock_validate_token_success,
     mock_list_groups_success,
     mock_acs_search_filter,
     mock_get_secret,
@@ -329,6 +331,14 @@ async def auth_client(
             yield client
 
 
+@pytest.fixture
+def mock_validate_token_success(monkeypatch):
+    async def mock_validate_access_token(self, token):
+        pass
+
+    monkeypatch.setattr(core.authentication.AuthenticationHelper, "validate_access_token", mock_validate_access_token)
+
+
 @pytest.fixture
 def mock_confidential_client_success(monkeypatch):
     def mock_acquire_token_on_behalf_of(self, *args, **kwargs):