fink: enhance ProcessLinksMiddleware with base URL handling and link transformation

- Added `get_base_url` utility to reconstruct the client's base URL from forwarded headers.
- Updated `ProcessLinksMiddleware` to utilize the new utility for transforming links in responses.
- Improved link transformation logic to handle various scenarios, including different hostnames and ports.
- Refactored tests for `ProcessLinksMiddleware` to cover new functionality and edge cases.

Author: alukach

src/stac_auth_proxy/middleware/ProcessLinksMiddleware.py

@@ -11,6 +11,7 @@
 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,6 +41,11 @@ 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."""
+        # 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):
             href = link.get("href")
             if not href:
@@ -48,12 +54,25 @@ def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, An
             try:
                 parsed_link = urlparse(href)
 
-                # Ignore links that are not for this proxy
-                if parsed_link.netloc != request.headers.get("host"):
+                if parsed_link.netloc not in [
+                    parsed_req_url.netloc,
+                    parsed_upstream_url.netloc,
+                ]:
+                    logger.warning(
+                        "Ignoring link %s because it is not for an endpoint behind this proxy (%s or %s)",
+                        href,
+                        parsed_req_url.netloc,
+                        parsed_upstream_url.netloc,
+                    )
                     continue
 
-                # Remove the upstream_url path from the link if it exists
-                parsed_upstream_url = urlparse(self.upstream_url)
+                if parsed_link.netloc == parsed_upstream_url.netloc:
+                    # Replace the upstream host with the client's host
+                    parsed_link = parsed_link._replace(
+                        netloc=parsed_req_url.netloc
+                    )._replace(scheme=parsed_req_url.scheme)
+
+                # Rewrite the link path
                 if parsed_upstream_url.path != "/" and parsed_link.path.startswith(
                     parsed_upstream_url.path
                 ):
@@ -68,6 +87,7 @@ def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, An
                     )
 
                 link["href"] = urlunparse(parsed_link)
+
             except Exception as e:
                 logger.error(
                     "Failed to parse link href %r, (ignoring): %s", href, str(e)

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