Add middleware for processing links

Author: alukach

src/stac_auth_proxy/app.py

@@ -18,10 +18,11 @@
     AddProcessTimeHeaderMiddleware,
     ApplyCql2FilterMiddleware,
     AuthenticationExtensionMiddleware,
-    BasePathMiddleware,
     BuildCql2FilterMiddleware,
     EnforceAuthMiddleware,
     OpenApiMiddleware,
+    ProcessLinksMiddleware,
+    RemoveRootPathMiddleware,
 )
 from .utils.lifespan import check_conformance, check_server_health
 
@@ -133,9 +134,16 @@ async def lifespan(app: FastAPI):
         oidc_config_url=settings.oidc_discovery_internal_url,
     )
 
+    if settings.root_path or settings.upstream_url.path != "/":
+        app.add_middleware(
+            ProcessLinksMiddleware,
+            upstream_url=str(settings.upstream_url),
+            base_path=settings.root_path,
+        )
+
     if settings.root_path:
         app.add_middleware(
-            BasePathMiddleware,
+            RemoveRootPathMiddleware,
             base_path=settings.root_path,
         )
 

src/stac_auth_proxy/middleware/BasePathMiddleware.py renamed to src/stac_auth_proxy/middleware/ProcessLinksMiddleware.py

@@ -1,14 +1,14 @@
-"""Middleware to remove BASE_PATH from incoming requests and update links in responses."""
+"""Middleware to remove the application root path from incoming requests and update links in responses."""
 
 import logging
 import re
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Optional
 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 starlette.types import ASGIApp, Scope
 
 from ..utils.middleware import JsonResponseMiddleware
 from ..utils.stac import get_links
@@ -17,35 +17,18 @@
 
 
 @dataclass
-class BasePathMiddleware(JsonResponseMiddleware):
+class ProcessLinksMiddleware(JsonResponseMiddleware):
     """
-    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
-    other middleware review the request.
+    Middleware to update links in responses, removing the upstream_url path and adding
+    the root_path if it exists.
     """
 
     app: ASGIApp
-    base_path: str
-    transform_links: bool = True
+    upstream_url: str
+    root_path: Optional[str] = None
 
     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."""
-        if scope["type"] != "http":
-            return await self.app(scope, receive, send)
-
-        path = scope["path"]
-
-        # Remove base_path if it exists at the start of the path
-        if path.startswith(self.base_path):
-            scope["raw_path"] = scope["path"].encode()
-            scope["path"] = path[len(self.base_path) :] or "/"
-
-        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(
@@ -69,11 +52,22 @@ def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, An
                 if parsed_link.netloc != request.headers.get("host"):
                     continue
 
-                parsed_link = parsed_link._replace(
-                    path=f"{self.base_path}{parsed_link.path}"
-                )
+                # Remove the upstream_url path from the link if it exists
+                if urlparse(self.upstream_url).path != "/":
+                    parsed_link = parsed_link._replace(
+                        path=parsed_link.path[len(urlparse(self.upstream_url).path) :]
+                    )
+
+                # Add the root_path to the link if it exists
+                if self.root_path:
+                    parsed_link = parsed_link._replace(
+                        path=f"{self.root_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))
+                logger.error(
+                    "Failed to parse link href %r, (ignoring): %s", href, str(e)
+                )
 
         return data

src/stac_auth_proxy/middleware/RemoveRootPathMiddleware.py

@@ -0,0 +1,40 @@
+"""Middleware to remove BASE_PATH from incoming requests and update links in responses."""
+
+import logging
+from dataclasses import dataclass
+
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RemoveRootPathMiddleware:
+    """
+    Middleware to remove the base path of the request before it is sent to the upstream
+    server.
+
+    IMPORTANT: This middleware must be place early in the middleware chain (ie late in the
+    order of declaration) so that it trims the base_path from the request path before any
+    middleware that may need to use the request path (e.g. EnforceAuthMiddleware).
+    """
+
+    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."""
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        path = scope["path"]
+
+        # Remove base_path if it exists at the start of the path
+        if path.startswith(self.base_path):
+            scope["raw_path"] = scope["path"].encode()
+            scope["path"] = path[len(self.base_path) :] or "/"
+
+        return await self.app(scope, receive, send)

src/stac_auth_proxy/middleware/__init__.py

@@ -3,17 +3,19 @@
 from .AddProcessTimeHeaderMiddleware import AddProcessTimeHeaderMiddleware
 from .ApplyCql2FilterMiddleware import ApplyCql2FilterMiddleware
 from .AuthenticationExtensionMiddleware import AuthenticationExtensionMiddleware
-from .BasePathMiddleware import BasePathMiddleware
 from .BuildCql2FilterMiddleware import BuildCql2FilterMiddleware
 from .EnforceAuthMiddleware import EnforceAuthMiddleware
+from .ProcessLinksMiddleware import ProcessLinksMiddleware
+from .RemoveRootPathMiddleware import RemoveRootPathMiddleware
 from .UpdateOpenApiMiddleware import OpenApiMiddleware
 
 __all__ = [
     "AddProcessTimeHeaderMiddleware",
     "ApplyCql2FilterMiddleware",
     "AuthenticationExtensionMiddleware",
-    "BasePathMiddleware",
     "BuildCql2FilterMiddleware",
     "EnforceAuthMiddleware",
+    "ProcessLinksMiddleware",
+    "RemoveRootPathMiddleware",
     "OpenApiMiddleware",
 ]

src/stac_auth_proxy/utils/stac.py

@@ -3,7 +3,7 @@
 from itertools import chain
 
 
-def get_links(data: dict) -> chain[dict]:  # type: ignore
+def get_links(data: dict) -> chain[dict]:
     """Get all links from a STAC response."""
     return chain(
         # Item/Collection