Merge branch 'main' into example/opa

Author: alukach

README.md

@@ -10,11 +10,13 @@
 
 STAC Auth Proxy is a proxy API that mediates between the client and your internally accessible STAC API to provide flexible authentication, authorization, and content-filtering mechanisms.
 
-## Features
+## ✨Features✨
 
-- 🔐 Authentication: Selectively apply OIDC auth to some or all endpoints & methods
-- 🎟️ Content Filtering: Apply CQL2 filters to client requests, filtering API content based on user context
-- 📖 OpenAPI Augmentation: Update [OpenAPI](https://swagger.io/specification/) with security requirements, keeping auto-generated docs/UIs accurate (e.g. [Swagger UI](https://swagger.io/tools/swagger-ui/))
+- 🔐 Authentication: Selectively apply [OpenID Connect (OIDC)](https://openid.net/developers/how-connect-works/) auth*n token validation & optional scope requirements to some or all endpoints & methods
+- 🛂 Content Filtering: Apply CQL2 filters to client requests, utilizing the [Filter Extension](https://github.com/stac-api-extensions/filter?tab=readme-ov-file) to filter API content based on user context
+- 🧩 Authentication Extension: Integrate the [Authentication Extension](https://github.com/stac-extensions/authentication) into API responses
+- 📘 OpenAPI Augmentation: Update API's [OpenAPI document](https://swagger.io/specification/) with security requirements, keeping auto-generated docs/UIs accurate (e.g. [Swagger UI](https://swagger.io/tools/swagger-ui/))
+- 🗜️ Response compression: Compress API responses via [`starlette-cramjam`](https://github.com/developmentseed/starlette-cramjam/)
 
 ## Usage
 
@@ -75,6 +77,10 @@ The application is configurable via environment variables.
     - **Type:** boolean
     - **Required:** No, defaults to `true`
     - **Example:** `false`, `1`, `True`
+  - **`ENABLE_COMPRESSION`**, enable response compression
+    - **Type:** boolean
+    - **Required:** No, defaults to `true`
+    - **Example:** `false`, `1`, `True`
   - **`HEALTHZ_PREFIX`**, path prefix for health check endpoints
     - **Type:** string
     - **Required:** No, defaults to `/healthz`
@@ -115,6 +121,10 @@ The application is configurable via environment variables.
         "^/healthz": ["GET"]
       }
       ```
+  - **`ENABLE_AUTHENTICATION_EXTENSION`**, enable authentication extension in STAC API responses
+    - **Type:** boolean
+    - **Required:** No, defaults to `true`
+    - **Example:** `false`, `1`, `True`
   - **`OPENAPI_SPEC_ENDPOINT`**, path of OpenAPI specification, used for augmenting spec response with auth configuration
     - **Type:** string or null
     - **Required:** No, defaults to `null` (disabled)

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,14 @@ async def lifespan(app: FastAPI):
     #
     # Middleware (order is important, last added = first to run)
     #
+    if settings.enable_authentication_extension:
+        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,
@@ -105,9 +114,10 @@ async def lifespan(app: FastAPI):
             items_filter=settings.items_filter(),
         )
 
-    app.add_middleware(
-        CompressionMiddleware,
-    )
+    if settings.enable_compression:
+        app.add_middleware(
+            CompressionMiddleware,
+        )
 
     app.add_middleware(
         AddProcessTimeHeaderMiddleware,

src/stac_auth_proxy/config.py

@@ -40,8 +40,8 @@ class Settings(BaseSettings):
 
     wait_for_upstream: bool = True
     check_conformance: bool = True
-
-    # Endpoints
+    enable_compression: bool = True
+    enable_authentication_extension: bool = True
     healthz_prefix: str = Field(pattern=_PREFIX_PATTERN, default="/healthz")
     openapi_spec_endpoint: Optional[str] = Field(pattern=_PREFIX_PATTERN, default=None)
 

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."""