Merge branch 'main' into authentication-ext/asset-signing

Author: alukach

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
   "boto3>=1.37.16",
   "brotli>=1.1.0",
   "cql2>=0.3.6",
+  "cryptography>=44.0.1",
   "fastapi>=0.115.5",
   "httpx[http2]>=0.28.0",
   "jinja2>=3.1.4",

src/stac_auth_proxy/app.py

@@ -102,7 +102,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

@@ -7,8 +7,6 @@
 from typing import Any, Optional
 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
@@ -33,35 +31,16 @@ class AuthenticationExtensionMiddleware(JsonResponseMiddleware):
     private_endpoints: EndpointMethods
     public_endpoints: EndpointMethods
 
-    oidc_config_url: Optional[HttpUrl] = None
     signing_scheme_name: str = "signed_url_auth"
     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)"
+    json_content_type_expr: str = r"application/(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", []))
-                        },
-                    },
-                },
-            }
+    state_key: str = "oidc_metadata"
 
     def should_transform_response(
         self, request: Request, response_headers: Headers
@@ -82,23 +61,42 @@ def should_transform_response(
             ]
         )
 
-    def transform_json(self, doc: dict[str, Any]) -> dict[str, Any]:
+    def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, Any]:
         """Augment the STAC Item with auth information."""
-        extensions = doc.setdefault("stac_extensions", [])
+        extensions = data.setdefault("stac_extensions", [])
         if self.extension_url not in extensions:
             extensions.append(self.extension_url)
 
-        # TODO: Should we add this to items even if the assets don't match the asset expression?
         # auth:schemes
         # ---
         # A property that contains all of the scheme definitions used by Assets and
         # Links in the STAC Item or Collection.
         # - Catalogs
         # - Collections
         # - Item Properties
-        scheme_loc = doc["properties"] if "properties" in doc else doc
+
+        oidc_metadata = getattr(request.state, self.state_key, {})
+        if not oidc_metadata:
+            logger.error(
+                "OIDC metadata not found in scope. Skipping authentication extension."
+            )
+            return data
+
+        scheme_loc = data["properties"] if "properties" in data else data
         schemes = scheme_loc.setdefault("auth:schemes", {})
-        schemes[self.auth_scheme_name] = self.auth_scheme
+        schemes[self.auth_scheme_name] = {
+            "type": "oauth2",
+            "description": "requires an authentication bearertoken",
+            "flows": {
+                "authorizationCode": {
+                    "authorizationUrl": oidc_metadata["authorization_endpoint"],
+                    "tokenUrl": oidc_metadata.get("token_endpoint"),
+                    "scopes": {
+                        k: k for k in sorted(oidc_metadata.get("scopes_supported", []))
+                    },
+                },
+            },
+        }
         if self.signing_endpoint:
             schemes[self.signing_scheme_name] = {
                 "type": "signedUrl",
@@ -138,11 +136,11 @@ def transform_json(self, doc: dict[str, Any]) -> dict[str, Any]:
         if self.signing_endpoint:
             assets = chain(
                 # Item
-                doc.get("assets", {}).values(),
+                data.get("assets", {}).values(),
                 # Items/Search
                 (
                     asset
-                    for item in doc.get("features", [])
+                    for item in data.get("features", [])
                     for asset in item.get("assets", {}).values()
                 ),
             )
@@ -152,16 +150,15 @@ def transform_json(self, doc: dict[str, Any]) -> dict[str, Any]:
                     continue
                 if re.match(self.signed_asset_expression, asset["href"]):
                     asset.setdefault("auth:refs", []).append(self.signing_scheme_name)
-
         # Annotate links with "auth:refs": [auth_scheme]
         links = chain(
             # Item/Collection
-            doc.get("links", []),
+            data.get("links", []),
             # Collections/Items/Search
             (
                 link
                 for prop in ["features", "collections"]
-                for object_with_links in doc.get(prop, [])
+                for object_with_links in data.get(prop, [])
                 for link in object_with_links.get("links", [])
             ),
         )
@@ -179,4 +176,4 @@ def transform_json(self, doc: dict[str, Any]) -> dict[str, Any]:
             if match.is_private:
                 link.setdefault("auth:refs", []).append(self.auth_scheme_name)
 
-        return doc
+        return data

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

@@ -41,15 +41,15 @@ def should_transform_response(
             ]
         )
 
-    def transform_json(self, openapi_spec: dict[str, Any]) -> dict[str, Any]:
+    def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, Any]:
         """Augment the OpenAPI spec with auth information."""
-        components = openapi_spec.setdefault("components", {})
+        components = data.setdefault("components", {})
         securitySchemes = components.setdefault("securitySchemes", {})
         securitySchemes[self.oidc_auth_scheme_name] = {
             "type": "openIdConnect",
             "openIdConnectUrl": self.oidc_config_url,
         }
-        for path, method_config in openapi_spec["paths"].items():
+        for path, method_config in data["paths"].items():
             for method, config in method_config.items():
                 match = find_match(
                     path,
@@ -62,4 +62,4 @@ def transform_json(self, openapi_spec: dict[str, Any]) -> dict[str, Any]:
                     config.setdefault("security", []).append(
                         {self.oidc_auth_scheme_name: match.required_scopes}
                     )
-        return openapi_spec
+        return data

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, request: Request) -> Any:
         """
         Transform the JSON data.
 
@@ -56,9 +56,10 @@ async def transform_response(message: Message) -> None:
 
             start_message = start_message or message
             headers = MutableHeaders(scope=start_message)
+            request = Request(scope)
 
             if not self.should_transform_response(
-                request=Request(scope),
+                request=request,
                 response_headers=headers,
             ):
                 # For non-JSON responses, send the start message immediately
@@ -78,7 +79,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, request=request)
                 body = json.dumps(transformed).encode()
 
             # Update content-length header

tests/conftest.py

@@ -32,7 +32,12 @@ def public_key(test_key: jwk.JWK) -> dict[str, Any]:
 @pytest.fixture(autouse=True)
 def mock_jwks(public_key: dict[str, Any]):
     """Mock JWKS endpoint."""
-    mock_oidc_config = {"jwks_uri": "https://example.com/jwks"}
+    mock_oidc_config = {
+        "jwks_uri": "https://example.com/jwks",
+        "authorization_endpoint": "https://example.com/auth",
+        "token_endpoint": "https://example.com/token",
+        "scopes_supported": ["openid", "profile", "email", "collection:create"],
+    }
 
     mock_jwks = {"keys": [public_key]}