Integrate with Authentication Extension

pyproject.toml

@@ -6,8 +6,10 @@ classifiers = [
   "License :: OSI Approved :: MIT License",
 ]
 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

@@ -17,6 +17,7 @@
 from .middleware import (
     AddProcessTimeHeaderMiddleware,
     ApplyCql2FilterMiddleware,
+    AuthenticationExtensionMiddleware,
     BuildCql2FilterMiddleware,
     EnforceAuthMiddleware,
     OpenApiMiddleware,
@@ -86,6 +87,13 @@ async def lifespan(app: FastAPI):
     #
     # Middleware (order is important, last added = first to run)
     #
+    app.add_middleware(
+        AuthenticationExtensionMiddleware,
+        default_public=settings.default_public,
+        public_endpoints=settings.public_endpoints,
+        private_endpoints=settings.private_endpoints,
+    )
+
     if settings.openapi_spec_endpoint:
         app.add_middleware(
             OpenApiMiddleware,

src/stac_auth_proxy/middleware/AuthenticationExtensionMiddleware.py

@@ -0,0 +1,125 @@
+"""Middleware to add auth information to item response served by upstream API."""
+
+import logging
+import re
+from dataclasses import dataclass, field
+from itertools import chain
+from typing import Any
+from urllib.parse import urlparse
+
+from starlette.datastructures import Headers
+from starlette.requests import Request
+from starlette.types import ASGIApp
+
+from ..config import EndpointMethods
+from ..utils.middleware import JsonResponseMiddleware
+from ..utils.requests import find_match
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AuthenticationExtensionMiddleware(JsonResponseMiddleware):
+    """Middleware to add the authentication extension to the response."""
+
+    app: ASGIApp
+
+    default_public: bool
+    private_endpoints: EndpointMethods
+    public_endpoints: EndpointMethods
+
+    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/(geo\+)?json"
+
+    state_key: str = "oidc_metadata"
+
+    def should_transform_response(
+        self, request: Request, response_headers: Headers
+    ) -> bool:
+        """Determine if the response should be transformed."""
+        # Match STAC catalog, collection, or item URLs with a single regex
+        return all(
+            [
+                re.match(
+                    # catalog, collections, collection, items, item, search
+                    r"^(/|/collections(/[^/]+(/items(/[^/]+)?)?)?|/search)$",
+                    request.url.path,
+                ),
+                re.match(
+                    self.json_content_type_expr,
+                    response_headers.get("content-type", ""),
+                ),
+            ]
+        )
+
+    def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, Any]:
+        """Augment the STAC Item with auth information."""
+        extensions = data.setdefault("stac_extensions", [])
+        if self.extension_url not in extensions:
+            extensions.append(self.extension_url)
+
+        # 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
+
+        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] = {
+            "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", []))
+                    },
+                },
+            },
+        }
+
+        # auth:refs
+        # ---
+        # Annotate links with "auth:refs": [auth_scheme]
+        links = chain(
+            # Item/Collection
+            data.get("links", []),
+            # Collections/Items/Search
+            (
+                link
+                for prop in ["features", "collections"]
+                for object_with_links in data.get(prop, [])
+                for link in object_with_links.get("links", [])
+            ),
+        )
+        for link in links:
+            if "href" not in link:
+                logger.warning("Link %s has no href", link)
+                continue
+            match = find_match(
+                path=urlparse(link["href"]).path,
+                method="GET",
+                private_endpoints=self.private_endpoints,
+                public_endpoints=self.public_endpoints,
+                default_public=self.default_public,
+            )
+            if match.is_private:
+                link.setdefault("auth:refs", []).append(self.auth_scheme_name)
+
+        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/middleware/__init__.py

@@ -2,17 +2,16 @@
 
 from .AddProcessTimeHeaderMiddleware import AddProcessTimeHeaderMiddleware
 from .ApplyCql2FilterMiddleware import ApplyCql2FilterMiddleware
+from .AuthenticationExtensionMiddleware import AuthenticationExtensionMiddleware
 from .BuildCql2FilterMiddleware import BuildCql2FilterMiddleware
 from .EnforceAuthMiddleware import EnforceAuthMiddleware
 from .UpdateOpenApiMiddleware import OpenApiMiddleware
 
 __all__ = [
-    x.__name__
-    for x in [
-        OpenApiMiddleware,
-        AddProcessTimeHeaderMiddleware,
-        EnforceAuthMiddleware,
-        BuildCql2FilterMiddleware,
-        ApplyCql2FilterMiddleware,
-    ]
+    "AddProcessTimeHeaderMiddleware",
+    "ApplyCql2FilterMiddleware",
+    "AuthenticationExtensionMiddleware",
+    "BuildCql2FilterMiddleware",
+    "EnforceAuthMiddleware",
+    "OpenApiMiddleware",
 ]

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