Retrieve oidc info from scope

Author: alukach

src/stac_auth_proxy/app.py

@@ -92,7 +92,6 @@ async def lifespan(app: FastAPI):
         default_public=settings.default_public,
         public_endpoints=settings.public_endpoints,
         private_endpoints=settings.private_endpoints,
-        oidc_config_url=settings.oidc_discovery_internal_url,
     )
 
     if settings.openapi_spec_endpoint:

src/stac_auth_proxy/middleware/AuthenticationExtensionMiddleware.py

@@ -4,14 +4,12 @@
 import re
 from dataclasses import dataclass, field
 from itertools import chain
-from typing import Any, Optional
+from typing import Any
 from urllib.parse import urlparse
 
-import httpx
-from pydantic import HttpUrl
 from starlette.datastructures import Headers
 from starlette.requests import Request
-from starlette.types import ASGIApp
+from starlette.types import ASGIApp, Scope
 
 from ..config import EndpointMethods
 from ..utils.middleware import JsonResponseMiddleware
@@ -30,34 +28,13 @@ class AuthenticationExtensionMiddleware(JsonResponseMiddleware):
     private_endpoints: EndpointMethods
     public_endpoints: EndpointMethods
 
-    oidc_config_url: Optional[HttpUrl] = None
     auth_scheme_name: str = "oauth"
     auth_scheme: dict[str, Any] = field(default_factory=dict)
     extension_url: str = (
         "https://stac-extensions.github.io/authentication/v1.1.0/schema.json"
     )
 
-    json_content_type_expr: str = r"application/json|geo\+json)"
-
-    def __post_init__(self):
-        """Load after initialization."""
-        if self.oidc_config_url and not self.auth_scheme:
-            # Retrieve OIDC configuration and extract authorization and token URLs
-            oidc_config = httpx.get(str(self.oidc_config_url)).json()
-            self.auth_scheme = {
-                "type": "oauth2",
-                "description": "requires an authentication token",
-                "flows": {
-                    "authorizationCode": {
-                        "authorizationUrl": oidc_config.get("authorization_endpoint"),
-                        "tokenUrl": oidc_config.get("token_endpoint"),
-                        "scopes": {
-                            k: k
-                            for k in sorted(oidc_config.get("scopes_supported", []))
-                        },
-                    },
-                },
-            }
+    json_content_type_expr: str = r"(application/json|geo\+json)"
 
     def should_transform_response(
         self, request: Request, response_headers: Headers
@@ -78,7 +55,7 @@ def should_transform_response(
             ]
         )
 
-    def transform_json(self, doc: dict[str, Any]) -> dict[str, Any]:
+    def transform_json(self, doc: dict[str, Any], scope: Scope) -> dict[str, Any]:
         """Augment the STAC Item with auth information."""
         extensions = doc.setdefault("stac_extensions", [])
         if self.extension_url not in extensions:
@@ -92,9 +69,19 @@ def transform_json(self, doc: dict[str, Any]) -> dict[str, Any]:
         # - Catalogs
         # - Collections
         # - Item Properties
+
+        if "oidc_metadata" not in scope:
+            logger.error(
+                "OIDC metadata not found in scope. "
+                "Skipping authentication extension."
+            )
+            return doc
+
         scheme_loc = doc["properties"] if "properties" in doc else doc
         schemes = scheme_loc.setdefault("auth:schemes", {})
-        schemes[self.auth_scheme_name] = self.auth_scheme
+        schemes[self.auth_scheme_name] = self.parse_oidc_config(
+            scope.get("oidc_metadata", {})
+        )
 
         # auth:refs
         # ---
@@ -125,3 +112,19 @@ def transform_json(self, doc: dict[str, Any]) -> dict[str, Any]:
                 link.setdefault("auth:refs", []).append(self.auth_scheme_name)
 
         return doc
+
+    def parse_oidc_config(self, oidc_config: dict[str, Any]) -> dict[str, Any]:
+        """Parse the OIDC configuration."""
+        return {
+            "type": "oauth2",
+            "description": "requires an authentication token",
+            "flows": {
+                "authorizationCode": {
+                    "authorizationUrl": oidc_config["authorization_endpoint"],
+                    "tokenUrl": oidc_config.get("token_endpoint"),
+                    "scopes": {
+                        k: k for k in sorted(oidc_config.get("scopes_supported", []))
+                    },
+                },
+            },
+        }

src/stac_auth_proxy/middleware/EnforceAuthMiddleware.py

@@ -1,7 +1,7 @@
 """Middleware to enforce authentication."""
 
 import logging
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Annotated, Any, Optional, Sequence
 from urllib.parse import urlparse, urlunparse
 
@@ -18,6 +18,53 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class OidcService:
+    """OIDC configuration and JWKS client."""
+
+    oidc_config_url: HttpUrl
+    jwks_client: jwt.PyJWKClient = field(init=False)
+    metadata: dict[str, Any] = field(init=False)
+
+    def __post_init__(self) -> None:
+        """Initialize OIDC config and JWKS client."""
+        logger.debug("Requesting OIDC config")
+        origin_url = str(self.oidc_config_url)
+
+        try:
+            response = httpx.get(origin_url)
+            response.raise_for_status()
+            self.metadata = response.json()
+            assert self.metadata, "OIDC metadata is empty"
+
+            # NOTE: We manually replace the origin of the jwks_uri in the event that
+            # the jwks_uri is not available from within the proxy.
+            oidc_url = urlparse(origin_url)
+            jwks_uri = urlunparse(
+                urlparse(self.metadata["jwks_uri"])._replace(
+                    netloc=oidc_url.netloc, scheme=oidc_url.scheme
+                )
+            )
+            if jwks_uri != self.metadata["jwks_uri"]:
+                logger.warning(
+                    "JWKS URI has been rewritten from %s to %s",
+                    self.metadata["jwks_uri"],
+                    jwks_uri,
+                )
+            self.jwks_client = jwt.PyJWKClient(jwks_uri)
+        except httpx.HTTPStatusError as e:
+            logger.error(
+                "Received a non-200 response when fetching OIDC config: %s",
+                e.response.text,
+            )
+            raise OidcFetchError(
+                f"Request for OIDC config failed with status {e.response.status_code}"
+            ) from e
+        except httpx.RequestError as e:
+            logger.error("Error fetching OIDC config from %s: %s", origin_url, str(e))
+            raise OidcFetchError(f"Request for OIDC config failed: {str(e)}") from e
+
+
 @dataclass
 class EnforceAuthMiddleware:
     """Middleware to enforce authentication."""
@@ -26,56 +73,11 @@ class EnforceAuthMiddleware:
     private_endpoints: EndpointMethods
     public_endpoints: EndpointMethods
     default_public: bool
-
     oidc_config_url: HttpUrl
     allowed_jwt_audiences: Optional[Sequence[str]] = None
-
     state_key: str = "payload"
 
-    # Generated attributes
-    _jwks_client: Optional[jwt.PyJWKClient] = None
-
-    @property
-    def jwks_client(self) -> jwt.PyJWKClient:
-        """Get the OIDC configuration URL."""
-        if not self._jwks_client:
-            logger.debug("Requesting OIDC config")
-            origin_url = str(self.oidc_config_url)
-
-            try:
-                response = httpx.get(origin_url)
-                response.raise_for_status()
-                oidc_config = response.json()
-
-                # NOTE: We manually replace the origin of the jwks_uri in the event that
-                # the jwks_uri is not available from within the proxy.
-                oidc_url = urlparse(origin_url)
-                jwks_uri = urlunparse(
-                    urlparse(oidc_config["jwks_uri"])._replace(
-                        netloc=oidc_url.netloc, scheme=oidc_url.scheme
-                    )
-                )
-                if jwks_uri != oidc_config["jwks_uri"]:
-                    logger.warning(
-                        "JWKS URI has been rewritten from %s to %s",
-                        oidc_config["jwks_uri"],
-                        jwks_uri,
-                    )
-                self._jwks_client = jwt.PyJWKClient(jwks_uri)
-            except httpx.HTTPStatusError as e:
-                logger.error(
-                    "Received a non-200 response when fetching OIDC config: %s",
-                    e.response.text,
-                )
-                raise OidcFetchError(
-                    f"Request for OIDC config failed with status {e.response.status_code}"
-                ) from e
-            except httpx.RequestError as e:
-                logger.error(
-                    "Error fetching OIDC config from %s: %s", origin_url, str(e)
-                )
-                raise OidcFetchError(f"Request for OIDC config failed: {str(e)}") from e
-        return self._jwks_client
+    _oidc_config: Optional[OidcService] = None
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Enforce authentication."""
@@ -107,6 +109,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             self.state_key,
             payload,
         )
+        setattr(request.state, "oidc_metadata", self.oidc_config.metadata)
         return await self.app(scope, receive, send)
 
     def validate_token(
@@ -137,7 +140,7 @@ def validate_token(
 
         # Parse & validate token
         try:
-            key = self.jwks_client.get_signing_key_from_jwt(token).key
+            key = self.oidc_config.jwks_client.get_signing_key_from_jwt(token).key
             payload = jwt.decode(
                 token,
                 key,
@@ -163,6 +166,13 @@ def validate_token(
                     )
         return payload
 
+    @property
+    def oidc_config(self) -> OidcService:
+        """Get the OIDC configuration."""
+        if not self._oidc_config:
+            self._oidc_config = OidcService(oidc_config_url=self.oidc_config_url)
+        return self._oidc_config
+
 
 class OidcFetchError(Exception):
     """Error fetching OIDC configuration."""

src/stac_auth_proxy/middleware/UpdateOpenApiMiddleware.py

@@ -6,7 +6,7 @@
 
 from starlette.datastructures import Headers
 from starlette.requests import Request
-from starlette.types import ASGIApp
+from starlette.types import ASGIApp, Scope
 
 from ..config import EndpointMethods
 from ..utils.middleware import JsonResponseMiddleware
@@ -41,7 +41,9 @@ def should_transform_response(
             ]
         )
 
-    def transform_json(self, openapi_spec: dict[str, Any]) -> dict[str, Any]:
+    def transform_json(
+        self, openapi_spec: dict[str, Any], scope: Scope
+    ) -> dict[str, Any]:
         """Augment the OpenAPI spec with auth information."""
         components = openapi_spec.setdefault("components", {})
         securitySchemes = components.setdefault("securitySchemes", {})

src/stac_auth_proxy/utils/middleware.py

@@ -29,7 +29,7 @@ def should_transform_response(
         ...
 
     @abstractmethod
-    def transform_json(self, data: Any) -> Any:
+    def transform_json(self, data: Any, scope: Scope) -> Any:
         """
         Transform the JSON data.
 
@@ -78,7 +78,7 @@ async def transform_response(message: Message) -> None:
             # Transform the JSON body
             if body:
                 data = json.loads(body)
-                transformed = self.transform_json(data)
+                transformed = self.transform_json(data, scope=scope)
                 body = json.dumps(transformed).encode()
 
             # Update content-length header

tests/test_middleware.py

@@ -6,7 +6,7 @@
 from starlette.datastructures import Headers
 from starlette.requests import Request
 from starlette.testclient import TestClient
-from starlette.types import ASGIApp
+from starlette.types import ASGIApp, Scope
 
 from stac_auth_proxy.utils.middleware import JsonResponseMiddleware
 
@@ -24,7 +24,7 @@ def should_transform_response(
         """Transform JSON responses based on content type."""
         return response_headers.get("content-type", "") == "application/json"
 
-    def transform_json(self, data: Any) -> Any:
+    def transform_json(self, data: Any, scope: Scope) -> Any:
         """Add a test field to the response."""
         if isinstance(data, dict):
             data["transformed"] = True