fix: improve link processing (#95)

We expect that upstream STAC APIs will set the origin of URLs contained
in the `href` property of the `links` array to be properly pointed to
the STAC Auth Proxy URL based on the [`Forwarded`
header](https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Forwarded)
we provide. However, in certain configurations, upstream APIs don't
perform this change (e.g. the proxy headers are not properly forwarded
to the upstream STAC API via a reverse proxy). This PR updates how we
rework links to ensure that links that have an origin and path matching
the path of the upstream STAC API will be reworked to instead point to
the origin and potential path of the STAC Auth Proxy.

Author: alukach

src/stac_auth_proxy/__main__.py

@@ -3,16 +3,20 @@
 import uvicorn
 from uvicorn.config import LOGGING_CONFIG
 
-LOGGING_CONFIG["loggers"][__package__] = {
-    "level": "DEBUG",
-    "handlers": ["default"],
-}
-
 uvicorn.run(
     f"{__package__}.app:create_app",
     host="0.0.0.0",
     port=8000,
-    log_config=LOGGING_CONFIG,
+    log_config={
+        **LOGGING_CONFIG,
+        "loggers": {
+            **LOGGING_CONFIG["loggers"],
+            __package__: {
+                "level": "DEBUG",
+                "handlers": ["default"],
+            },
+        },
+    },
     reload=True,
     factory=True,
 )

src/stac_auth_proxy/middleware/ProcessLinksMiddleware.py

@@ -4,13 +4,14 @@
 import re
 from dataclasses import dataclass
 from typing import Any, Optional
-from urllib.parse import urlparse, urlunparse
+from urllib.parse import ParseResult, urlparse, urlunparse
 
 from starlette.datastructures import Headers
 from starlette.requests import Request
 from starlette.types import ASGIApp, Scope
 
 from ..utils.middleware import JsonResponseMiddleware
+from ..utils.requests import get_base_url
 from ..utils.stac import get_links
 
 logger = logging.getLogger(__name__)
@@ -40,37 +41,81 @@ def should_transform_response(self, request: Request, scope: Scope) -> bool:
 
     def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, Any]:
         """Update links in the response to include root_path."""
-        for link in get_links(data):
-            href = link.get("href")
-            if not href:
-                continue
+        # Get the client's actual base URL (accounting for load balancers/proxies)
+        req_base_url = get_base_url(request)
+        parsed_req_url = urlparse(req_base_url)
+        parsed_upstream_url = urlparse(self.upstream_url)
 
+        for link in get_links(data):
             try:
-                parsed_link = urlparse(href)
-
-                # Ignore links that are not for this proxy
-                if parsed_link.netloc != request.headers.get("host"):
-                    continue
-
-                # Remove the upstream_url path from the link if it exists
-                parsed_upstream_url = urlparse(self.upstream_url)
-                if parsed_upstream_url.path != "/" and parsed_link.path.startswith(
-                    parsed_upstream_url.path
-                ):
-                    parsed_link = parsed_link._replace(
-                        path=parsed_link.path[len(parsed_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)
+                self._update_link(link, parsed_req_url, parsed_upstream_url)
             except Exception as e:
                 logger.error(
-                    "Failed to parse link href %r, (ignoring): %s", href, str(e)
+                    "Failed to parse link href %r, (ignoring): %s",
+                    link.get("href"),
+                    str(e),
                 )
-
         return data
+
+    def _update_link(
+        self, link: dict[str, Any], request_url: ParseResult, upstream_url: ParseResult
+    ) -> None:
+        """
+        Ensure that link hrefs that are local to upstream url are rewritten as local to
+        the proxy.
+        """
+        if "href" not in link:
+            logger.warning("Link %r has no href", link)
+            return
+
+        parsed_link = urlparse(link["href"])
+
+        if parsed_link.netloc not in [
+            request_url.netloc,
+            upstream_url.netloc,
+        ]:
+            logger.debug(
+                "Ignoring link %s because it is not for an endpoint behind this proxy (%s or %s)",
+                link["href"],
+                request_url.netloc,
+                upstream_url.netloc,
+            )
+            return
+
+        # If the link path is not a descendant of the upstream path, don't transform it
+        if upstream_url.path != "/" and not parsed_link.path.startswith(
+            upstream_url.path
+        ):
+            logger.debug(
+                "Ignoring link %s because it is not descendant of upstream path (%s)",
+                link["href"],
+                upstream_url.path,
+            )
+            return
+
+        # Replace the upstream host with the client's host
+        if parsed_link.netloc == upstream_url.netloc:
+            parsed_link = parsed_link._replace(netloc=request_url.netloc)._replace(
+                scheme=request_url.scheme
+            )
+
+        # Rewrite the link path
+        if upstream_url.path != "/" and parsed_link.path.startswith(upstream_url.path):
+            parsed_link = parsed_link._replace(
+                path=parsed_link.path[len(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}"
+            )
+
+        logger.debug(
+            "Rewriting %r link %r to %r",
+            link.get("rel"),
+            link["href"],
+            urlunparse(parsed_link),
+        )
+
+        link["href"] = urlunparse(parsed_link)

src/stac_auth_proxy/utils/requests.py

@@ -1,13 +1,18 @@
 """Utility functions for working with HTTP requests."""
 
 import json
+import logging
 import re
 from dataclasses import dataclass, field
-from typing import Optional, Sequence
+from typing import Dict, Optional, Sequence
 from urllib.parse import urlparse
 
+from starlette.requests import Request
+
 from ..config import EndpointMethods
 
+logger = logging.getLogger(__name__)
+
 
 def extract_variables(url: str) -> dict:
     """
@@ -90,3 +95,110 @@ def build_server_timing_header(
     if current_value:
         return f"{current_value}, {metric}"
     return metric
+
+
+def parse_forwarded_header(forwarded_header: str) -> Dict[str, str]:
+    """
+    Parse the Forwarded header according to RFC 7239.
+
+    Args:
+        forwarded_header: The Forwarded header value
+
+    Returns:
+        Dictionary containing parsed forwarded information (proto, host, for, by, etc.)
+
+    Example:
+        >>> parse_forwarded_header("for=192.0.2.43; by=203.0.113.60; proto=https; host=api.example.com")
+        {'for': '192.0.2.43', 'by': '203.0.113.60', 'proto': 'https', 'host': 'api.example.com'}
+
+    """
+    # Forwarded header format: "for=192.0.2.43, for=198.51.100.17; by=203.0.113.60; proto=https; host=example.com"
+    # The format is: for=value1, for=value2; by=value; proto=value; host=value
+    # We need to parse all the key=value pairs, taking the first 'for' value
+    forwarded_info = {}
+
+    try:
+        # Parse all key=value pairs separated by semicolons
+        for pair in forwarded_header.split(";"):
+            pair = pair.strip()
+            if "=" in pair:
+                key, value = pair.split("=", 1)
+                key = key.strip()
+                value = value.strip().strip('"')
+
+                # For 'for' field, only take the first value if there are multiple
+                if key == "for" and key not in forwarded_info:
+                    # Extract the first for value (before comma if present)
+                    first_for_value = value.split(",")[0].strip()
+                    forwarded_info[key] = first_for_value
+                elif key != "for":
+                    # For other fields, just use the value as-is
+                    forwarded_info[key] = value
+    except Exception as e:
+        logger.warning(f"Failed to parse Forwarded header '{forwarded_header}': {e}")
+        return {}
+
+    return forwarded_info
+
+
+def get_base_url(request: Request) -> str:
+    """
+    Get the request's base URL, accounting for forwarded headers from load balancers/proxies.
+
+    This function handles both the standard Forwarded header (RFC 7239) and legacy
+    X-Forwarded-* headers to reconstruct the original client URL when the service
+    is deployed behind load balancers or reverse proxies.
+
+    Args:
+        request: The Starlette request object
+
+    Returns:
+        The reconstructed client base URL
+
+    Example:
+        >>> # With Forwarded header
+        >>> request.headers = {"Forwarded": "for=192.0.2.43; proto=https; host=api.example.com"}
+        >>> get_base_url(request)
+        "https://api.example.com/"
+
+        >>> # With X-Forwarded-* headers
+        >>> request.headers = {"X-Forwarded-Host": "api.example.com", "X-Forwarded-Proto": "https"}
+        >>> get_base_url(request)
+        "https://api.example.com/"
+
+    """
+    # Check for standard Forwarded header first (RFC 7239)
+    forwarded_header = request.headers.get("Forwarded")
+    if forwarded_header:
+        try:
+            forwarded_info = parse_forwarded_header(forwarded_header)
+            # Only use Forwarded header if we successfully parsed it and got useful info
+            if forwarded_info and (
+                "proto" in forwarded_info or "host" in forwarded_info
+            ):
+                scheme = forwarded_info.get("proto", request.url.scheme)
+                host = forwarded_info.get("host", request.url.netloc)
+                # Note: Forwarded header doesn't include path, so we use request.base_url.path
+                path = request.base_url.path
+                return f"{scheme}://{host}{path}"
+        except Exception as e:
+            logger.warning(f"Failed to parse Forwarded header: {e}")
+
+    # Fall back to legacy X-Forwarded-* headers
+    forwarded_host = request.headers.get("X-Forwarded-Host")
+    forwarded_proto = request.headers.get("X-Forwarded-Proto")
+    forwarded_path = request.headers.get("X-Forwarded-Path")
+
+    if forwarded_host:
+        # Use forwarded headers to reconstruct the original client URL
+        scheme = forwarded_proto or request.url.scheme
+        netloc = forwarded_host
+        # Use forwarded path if available, otherwise use request base URL path
+        path = forwarded_path or request.base_url.path
+    else:
+        # Fall back to the request's base URL if no forwarded headers
+        scheme = request.url.scheme
+        netloc = request.url.netloc
+        path = request.base_url.path
+
+    return f"{scheme}://{netloc}{path}"