Process links

Author: alukach

src/stac_auth_proxy/app.py

@@ -18,10 +18,10 @@
     AddProcessTimeHeaderMiddleware,
     ApplyCql2FilterMiddleware,
     AuthenticationExtensionMiddleware,
+    BasePathMiddleware,
     BuildCql2FilterMiddleware,
     EnforceAuthMiddleware,
     OpenApiMiddleware,
-    RemoveRootPathMiddleware,
 )
 from .utils.lifespan import check_conformance, check_server_health
 
@@ -121,11 +121,6 @@ async def lifespan(app: FastAPI):
             items_filter=settings.items_filter(),
         )
 
-    if settings.enable_compression:
-        app.add_middleware(
-            CompressionMiddleware,
-        )
-
     app.add_middleware(
         AddProcessTimeHeaderMiddleware,
     )
@@ -140,8 +135,13 @@ async def lifespan(app: FastAPI):
 
     if settings.root_path:
         app.add_middleware(
-            RemoveRootPathMiddleware,
+            BasePathMiddleware,
             base_path=settings.root_path,
         )
 
+    if settings.enable_compression:
+        app.add_middleware(
+            CompressionMiddleware,
+        )
+
     return app

src/stac_auth_proxy/middleware/AuthenticationExtensionMiddleware.py

@@ -3,7 +3,6 @@
 import logging
 import re
 from dataclasses import dataclass, field
-from itertools import chain
 from typing import Any
 from urllib.parse import urlparse
 
@@ -14,6 +13,7 @@
 from ..config import EndpointMethods
 from ..utils.middleware import JsonResponseMiddleware
 from ..utils.requests import find_match
+from ..utils.stac import get_links
 
 logger = logging.getLogger(__name__)
 
@@ -101,18 +101,7 @@ def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, An
         # 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:
+        for link in get_links(data):
             if "href" not in link:
                 logger.warning("Link %s has no href", link)
                 continue

src/stac_auth_proxy/middleware/BasePathMiddleware.py

@@ -1,14 +1,25 @@
-"""Middleware to remove BASE_PATH from incoming requests."""
+"""Middleware to remove BASE_PATH from incoming requests and update links in responses."""
 
+import logging
+import re
 from dataclasses import dataclass
+from typing import Any
+from urllib.parse import urlparse, urlunparse
 
+from starlette.datastructures import Headers
+from starlette.requests import Request
 from starlette.types import ASGIApp, Receive, Scope, Send
 
+from ..utils.middleware import JsonResponseMiddleware
+from ..utils.stac import get_links
+
+logger = logging.getLogger(__name__)
+
 
 @dataclass
-class RemoveRootPathMiddleware:
+class BasePathMiddleware(JsonResponseMiddleware):
     """
-    Middleware to remove BASE_PATH from incoming requests.
+    Middleware to handle the base path of the request and update links in responses.
 
     IMPORTANT: This middleware must be the first middleware in the chain (ie last in the
     order of declaration) so that it trims the base_path from the request path before
@@ -17,6 +28,9 @@ class RemoveRootPathMiddleware:
 
     app: ASGIApp
     base_path: str
+    transform_links: bool = True
+
+    json_content_type_expr: str = r"application/(geo\+)?json"
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Remove BASE_PATH from the request path if it exists."""
@@ -30,4 +44,36 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             scope["raw_path"] = scope["path"].encode()
             scope["path"] = path[len(self.base_path) :] or "/"
 
-        return await self.app(scope, receive, send)
+        return await super().__call__(scope, receive, send)
+
+    def should_transform_response(self, request: Request, scope: Scope) -> bool:
+        """Only transform responses with JSON content type."""
+        return bool(
+            re.match(
+                self.json_content_type_expr,
+                Headers(scope=scope).get("content-type", ""),
+            )
+        )
+
+    def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, Any]:
+        """Update links in the response to include base_path."""
+        for link in get_links(data):
+            href = link.get("href")
+            if not href:
+                continue
+
+            try:
+                parsed_link = urlparse(href)
+
+                # Ignore links that are not for this proxy
+                if parsed_link.netloc != request.headers.get("host"):
+                    continue
+
+                parsed_link = parsed_link._replace(
+                    path=f"{self.base_path}{parsed_link.path}"
+                )
+                link["href"] = urlunparse(parsed_link)
+            except Exception as e:
+                logger.warning("Failed to parse link href %s: %s", href, str(e))
+
+        return data

src/stac_auth_proxy/middleware/__init__.py

@@ -3,7 +3,7 @@
 from .AddProcessTimeHeaderMiddleware import AddProcessTimeHeaderMiddleware
 from .ApplyCql2FilterMiddleware import ApplyCql2FilterMiddleware
 from .AuthenticationExtensionMiddleware import AuthenticationExtensionMiddleware
-from .BasePathMiddleware import RemoveRootPathMiddleware
+from .BasePathMiddleware import BasePathMiddleware
 from .BuildCql2FilterMiddleware import BuildCql2FilterMiddleware
 from .EnforceAuthMiddleware import EnforceAuthMiddleware
 from .UpdateOpenApiMiddleware import OpenApiMiddleware
@@ -12,7 +12,7 @@
     "AddProcessTimeHeaderMiddleware",
     "ApplyCql2FilterMiddleware",
     "AuthenticationExtensionMiddleware",
-    "RemoveRootPathMiddleware",
+    "BasePathMiddleware",
     "BuildCql2FilterMiddleware",
     "EnforceAuthMiddleware",
     "OpenApiMiddleware",

src/stac_auth_proxy/utils/stac.py

@@ -0,0 +1,18 @@
+"""STAC-specific utilities."""
+
+from itertools import chain
+
+
+def get_links(data: dict) -> chain[dict]:  # type: ignore
+    """Get all links from a STAC response."""
+    return 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", [])
+        ),
+    )