FIX: Add destination allowlist before forwarding Bearer credential to Graph in EntraAuthMiddleware (microsoft#2183)

Author: jsong468

pyrit/backend/middleware/auth.py

@@ -19,6 +19,7 @@
 import os
 from dataclasses import dataclass
 from typing import Any, ClassVar
+from urllib.parse import urlparse
 
 import httpx
 import jwt
@@ -51,6 +52,13 @@ class EntraAuthMiddleware(BaseHTTPMiddleware):
         "/api/media",
     }
 
+    # Hosts the user's Bearer token may be forwarded to
+    _GRAPH_HOSTS: ClassVar[set[str]] = {"graph.microsoft.com"}
+
+    # Trusted URL for resolving group membership; built from a constant (not the
+    # token-supplied endpoint) so token data cannot control the request destination.
+    _GRAPH_MEMBER_OBJECTS_URL: ClassVar[str] = "https://graph.microsoft.com/v1.0/me/getMemberObjects"
+
     def __init__(self, app: ASGIApp) -> None:
         """Initialize the middleware with Entra ID configuration from environment variables."""
         super().__init__(app)
@@ -167,34 +175,34 @@ async def _resolve_excess_groups_async(self, claims: dict[str, Any], token: str)
         """
         Resolve group membership via Microsoft Graph when user is in >200 groups.
 
-        When a user is in >200 groups, Entra ID replaces the `groups` claim with
-        `_claim_sources` containing a Graph API endpoint. This method calls the
-        Microsoft Graph `getMemberObjects` endpoint to retrieve transitive group
-        memberships, using the user's access token.
+        When a user is in >200 groups, Entra ID replaces the `groups` claim with a
+        `_claim_names` / `_claim_sources` overage pointer. This method confirms the
+        overage is present, then calls the Microsoft Graph `getMemberObjects` endpoint
+        to retrieve transitive group memberships, using the user's access token.
+
+        The Graph URL is built from a trusted constant (``_GRAPH_MEMBER_OBJECTS_URL``)
+        rather than the token-supplied endpoint, so token data cannot control the
+        request host, path, or port. Pagination links from the Graph response are
+        additionally checked against the Graph allowlist (see ``_is_trusted_graph_url``).
 
         Args:
-            claims: The decoded JWT claims containing _claim_sources.
+            claims: The decoded JWT claims containing the group overage pointer.
             token: The raw Bearer token to forward to Graph API.
 
         Returns:
-            List of group IDs the user belongs to, or empty list on failure.
+            List of group IDs the user belongs to, or empty list on failure or
+            when no group overage is present.
         """
         try:
-            claim_sources = claims.get("_claim_sources", {})
-            src = claim_sources.get("src1", {})
-            endpoint = src.get("endpoint", "")
-
-            if not endpoint:
-                logger.debug("No group resolution endpoint found in _claim_sources")
+            # Entra signals a groups overage via `_claim_names` -> source key -> `_claim_sources`.
+            # We only confirm the overage is present; the endpoint it supplies is intentionally
+            # ignored in favor of a trusted constant so token data cannot control the request.
+            claim_source_key = claims.get("_claim_names", {}).get("groups")
+            if not claim_source_key or claim_source_key not in claims.get("_claim_sources", {}):
+                logger.debug("No group overage claim source found")
                 return []
 
-            # The _claim_sources endpoint may be a legacy graph.windows.net URL.
-            # Rewrite to Microsoft Graph (graph.microsoft.com) which is the
-            # supported API. The legacy Azure AD Graph was retired in 2023.
-            if "graph.windows.net" in endpoint:
-                # Legacy format: https://graph.windows.net/{tenant}/users/{oid}/getMemberObjects
-                # Graph format:  https://graph.microsoft.com/v1.0/me/getMemberObjects
-                endpoint = "https://graph.microsoft.com/v1.0/me/getMemberObjects"
+            endpoint = self._GRAPH_MEMBER_OBJECTS_URL
 
             all_group_ids: list[str] = []
             async with httpx.AsyncClient() as client:
@@ -222,6 +230,9 @@ async def _resolve_excess_groups_async(self, claims: dict[str, Any], token: str)
                 # Handle pagination — Graph may return @odata.nextLink for large results
                 next_link = data.get("@odata.nextLink")
                 while next_link:
+                    if not self._is_trusted_graph_url(next_link):
+                        logger.warning("Refusing to follow untrusted pagination link: %s", next_link)
+                        break
                     response = await client.get(
                         next_link,
                         headers={"Authorization": f"Bearer {token}"},
@@ -241,6 +252,16 @@ async def _resolve_excess_groups_async(self, claims: dict[str, Any], token: str)
             logger.warning("Failed to resolve group memberships: %s", e)
             return []
 
+    def _is_trusted_graph_url(self, url: str) -> bool:
+        """
+        Whether *url* is an HTTPS Microsoft Graph URL safe to forward the token to.
+
+        Returns:
+            True if *url* uses HTTPS and its host is in the Graph allowlist, False otherwise.
+        """
+        parsed = urlparse(url)
+        return parsed.scheme == "https" and parsed.hostname in self._GRAPH_HOSTS
+
     def _validate_token(self, token: str) -> tuple[AuthenticatedUser | None, dict[str, Any]]:
         """
         Validate a JWT against Entra ID JWKS.

tests/unit/backend/test_auth_middleware.py

@@ -5,11 +5,18 @@
 Tests for the Entra ID auth middleware.
 """
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
 
 from pyrit.backend.middleware.auth import EntraAuthMiddleware
 
 
+def _make_middleware() -> EntraAuthMiddleware:
+    with patch.dict("os.environ", {"ENTRA_TENANT_ID": "", "ENTRA_CLIENT_ID": ""}, clear=False):
+        return EntraAuthMiddleware(MagicMock())
+
+
 def test_validate_token_returns_none_when_jwks_client_is_none():
     """Test that _validate_token returns (None, {}) when _jwks_client is None."""
     mock_app = MagicMock()
@@ -27,3 +34,84 @@ def test_validate_token_returns_none_when_jwks_client_is_none():
 
     assert user is None
     assert claims == {}
+
+
+@pytest.mark.parametrize(
+    "url, expected",
+    [
+        ("https://graph.microsoft.com/v1.0/me/getMemberObjects", True),
+        ("https://graph.microsoft.com/v1.0/me/memberOf?$skiptoken=abc", True),
+        ("http://graph.microsoft.com/v1.0/me/getMemberObjects", False),  # not https
+        ("https://evil.com/v1.0/me/getMemberObjects", False),  # wrong host
+        ("https://graph.microsoft.com.evil.com/x", False),  # suffix spoof
+        ("https://graph.microsoft.com@evil.com/x", False),  # userinfo spoof
+        ("", False),
+    ],
+)
+def test_is_trusted_graph_url(url, expected):
+    """Only HTTPS Microsoft Graph hosts are trusted to receive the forwarded token."""
+    assert _make_middleware()._is_trusted_graph_url(url) is expected
+
+
+async def test_resolve_excess_groups_no_overage_returns_empty():
+    """Claims without a groups overage pointer return [] and make no HTTP request."""
+    middleware = _make_middleware()
+    claims = {"_claim_sources": {"src1": {"endpoint": "https://graph.microsoft.com/x"}}}  # no _claim_names.groups
+
+    with patch("pyrit.backend.middleware.auth.httpx.AsyncClient") as mock_client:
+        result = await middleware._resolve_excess_groups_async(claims, "the-token")
+
+    assert result == []
+    mock_client.assert_not_called()
+
+
+async def test_resolve_excess_groups_ignores_token_endpoint():
+    """The Graph request uses the trusted constant URL, never the token-supplied endpoint."""
+    middleware = _make_middleware()
+    claims = {
+        "_claim_names": {"groups": "src1"},
+        "_claim_sources": {"src1": {"endpoint": "https://evil.com/steal"}},
+    }
+
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"value": ["group-1", "group-2"]}
+
+    mock_client = AsyncMock()
+    mock_client.post.return_value = mock_response
+    mock_client_cm = MagicMock()
+    mock_client_cm.__aenter__ = AsyncMock(return_value=mock_client)
+    mock_client_cm.__aexit__ = AsyncMock(return_value=False)
+
+    with patch("pyrit.backend.middleware.auth.httpx.AsyncClient", return_value=mock_client_cm):
+        result = await middleware._resolve_excess_groups_async(claims, "the-token")
+
+    assert result == ["group-1", "group-2"]
+    posted_url = mock_client.post.call_args.args[0]
+    assert posted_url == EntraAuthMiddleware._GRAPH_MEMBER_OBJECTS_URL
+    assert "evil.com" not in posted_url
+
+
+async def test_resolve_excess_groups_stops_on_untrusted_pagination_link():
+    """An untrusted @odata.nextLink halts pagination without issuing the follow-up GET."""
+    middleware = _make_middleware()
+    claims = {
+        "_claim_names": {"groups": "src1"},
+        "_claim_sources": {"src1": {"endpoint": "https://graph.microsoft.com/x"}},
+    }
+
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"value": ["group-1"], "@odata.nextLink": "https://evil.com/next"}
+
+    mock_client = AsyncMock()
+    mock_client.post.return_value = mock_response
+    mock_client_cm = MagicMock()
+    mock_client_cm.__aenter__ = AsyncMock(return_value=mock_client)
+    mock_client_cm.__aexit__ = AsyncMock(return_value=False)
+
+    with patch("pyrit.backend.middleware.auth.httpx.AsyncClient", return_value=mock_client_cm):
+        result = await middleware._resolve_excess_groups_async(claims, "the-token")
+
+    assert result == ["group-1"]
+    mock_client.get.assert_not_called()