fix: fetch existing record in-process for transaction validation

src/stac_auth_proxy/app.py

@@ -140,10 +140,7 @@ def configure_app(
 
     if settings.items_filter or settings.collections_filter:
         app.add_middleware(Cql2ValidateResponseBodyMiddleware)
-        app.add_middleware(
-            Cql2ValidateTransactionMiddleware,
-            upstream_url=str(settings.upstream_url),
-        )
+        app.add_middleware(Cql2ValidateTransactionMiddleware)
         app.add_middleware(Cql2ApplyFilterBodyMiddleware)
         app.add_middleware(Cql2ApplyFilterQueryStringMiddleware)
         app.add_middleware(Cql2RewriteLinksFilterMiddleware)

src/stac_auth_proxy/middleware/Cql2ValidateTransactionMiddleware.py

@@ -2,11 +2,10 @@
 
 import json
 import re
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from logging import getLogger
 from typing import Optional
 
-import httpx
 from cql2 import Expr
 from starlette.requests import Request
 from starlette.responses import JSONResponse
@@ -17,6 +16,10 @@
 logger = getLogger(__name__)
 
 
+class UpstreamError(Exception):
+    """Raised when the existing record could not be fetched."""
+
+
 def _deep_merge(base: dict, override: dict) -> dict:
     """Recursively merge override into base. Non-dict values from override replace base."""
     result = {**base}
@@ -38,19 +41,12 @@ class Cql2ValidateTransactionMiddleware:
     """Middleware to validate transaction requests against a CQL2 filter."""
 
     app: ASGIApp
-    upstream_url: str
     state_key: str = "cql2_filter"
 
-    _client: httpx.AsyncClient = field(init=False)
-
     # Transaction endpoint patterns
     items_pattern = r"^/collections/([^/]+)/(items|bulk_items)(?:/([^/]+))?$"
     collections_pattern = r"^/collections(?:/([^/]+))?$"
 
-    def __post_init__(self):
-        """Initialize the HTTP client."""
-        self._client = httpx.AsyncClient(base_url=self.upstream_url)
-
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Validate transaction requests against the CQL2 filter."""
         if scope["type"] != "http":
@@ -74,25 +70,21 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
                 return await self._handle_create(scope, receive, send, cql2_filter)
             if method in ("PUT", "PATCH"):
                 return await self._handle_update(
-                    scope, receive, send, cql2_filter, path, method
+                    scope, receive, send, cql2_filter, method
                 )
             if method == "DELETE":
-                return await self._handle_delete(
-                    scope, receive, send, cql2_filter, path
-                )
+                return await self._handle_delete(scope, receive, send, cql2_filter)
 
         # Match collections endpoints: /collections, /collections/{id}
         if re.match(self.collections_pattern, path):
             if method == "POST":
                 return await self._handle_create(scope, receive, send, cql2_filter)
             if method in ("PUT", "PATCH"):
                 return await self._handle_update(
-                    scope, receive, send, cql2_filter, path, method
+                    scope, receive, send, cql2_filter, method
                 )
             if method == "DELETE":
-                return await self._handle_delete(
-                    scope, receive, send, cql2_filter, path
-                )
+                return await self._handle_delete(scope, receive, send, cql2_filter)
 
         # Not a transaction endpoint, pass through
         return await self.app(scope, receive, send)
@@ -120,13 +112,54 @@ async def new_receive():
 
         return new_receive
 
-    async def _fetch_existing(self, path: str) -> Optional[dict]:
-        """Fetch the existing record from upstream."""
-        response = await self._client.get(path)
-        if response.status_code == 404:
+    async def _fetch_existing(self, scope: Scope) -> Optional[dict]:
+        """
+        Fetch the existing record by sending a GET for the same path to the
+        downstream app, in-process.
+
+        When deployed as a proxy, this reaches the upstream via the reverse proxy
+        handler; when deployed as middleware, it reaches the STAC API's routes
+        directly. Either way the request never re-enters the auth middleware, so no
+        credentials need to be forwarded.
+        """
+        sub_scope = {
+            **scope,
+            "method": "GET",
+            "query_string": b"",
+            # Drop the caller's headers (body sizing, conditionals, encodings) so the
+            # downstream always answers with a complete, plain JSON body.
+            "headers": [(k, v) for k, v in scope["headers"] if k == b"host"]
+            + [(b"accept", b"application/json")],
+            # Copy so downstream writes to request.state can't leak into the caller's
+            # request.
+            "state": dict(scope.get("state", {})),
+        }
+        status = None
+        body = b""
+
+        async def receive():
+            return {"type": "http.request", "body": b"", "more_body": False}
+
+        async def send(message):
+            nonlocal status, body
+            if message["type"] == "http.response.start":
+                status = message["status"]
+            else:
+                body += message.get("body", b"")
+
+        try:
+            await self.app(sub_scope, receive, send)
+        except Exception as e:
+            raise UpstreamError("Failed to fetch existing record") from e
+
+        if status == 404:
             return None
-        response.raise_for_status()
-        return response.json()
+        if status != 200:
+            raise UpstreamError(f"Unexpected status {status} fetching existing record")
+        try:
+            return json.loads(body)
+        except json.JSONDecodeError as e:
+            raise UpstreamError("Existing record is not valid JSON") from e
 
     async def _handle_create(
         self,
@@ -221,7 +254,6 @@ async def _handle_update(
         receive: Receive,
         send: Send,
         cql2_filter: Expr,
-        path: str,
         method: str,
     ) -> None:
         """Validate update requests."""
@@ -241,8 +273,8 @@ async def _handle_update(
 
         # Fetch existing record
         try:
-            existing = await self._fetch_existing(path)
-        except httpx.HTTPError:
+            existing = await self._fetch_existing(scope)
+        except UpstreamError:
             response = JSONResponse(
                 {
                     "code": "UpstreamError",
@@ -294,12 +326,11 @@ async def _handle_delete(
         receive: Receive,
         send: Send,
         cql2_filter: Expr,
-        path: str,
     ) -> None:
         """Validate delete requests."""
         try:
-            existing = await self._fetch_existing(path)
-        except httpx.HTTPError:
+            existing = await self._fetch_existing(scope)
+        except UpstreamError:
             response = JSONResponse(
                 {
                     "code": "UpstreamError",

tests/test_cql2_validate_transaction_middleware.py

@@ -7,10 +7,13 @@
 import pytest
 from cql2 import Expr
 from fastapi import FastAPI, Request
+from starlette.responses import JSONResponse, PlainTextResponse
 from starlette.testclient import TestClient
 
+from stac_auth_proxy.handlers import ReverseProxyHandler
 from stac_auth_proxy.middleware.Cql2ValidateTransactionMiddleware import (
     Cql2ValidateTransactionMiddleware,
+    UpstreamError,
     _deep_merge,
 )
 
@@ -28,12 +31,9 @@ def cql2_filter():
 def app_with_middleware():
     """Create a FastAPI app with the transaction middleware."""
 
-    def _create(upstream_url="http://upstream:8080"):
+    def _create():
         app = FastAPI()
-        app.add_middleware(
-            Cql2ValidateTransactionMiddleware,
-            upstream_url=upstream_url,
-        )
+        app.add_middleware(Cql2ValidateTransactionMiddleware)
 
         @app.post("/collections/{collection_id}/items")
         async def create_item(request: Request):
@@ -540,7 +540,7 @@ def test_upstream_unreachable(
             Cql2ValidateTransactionMiddleware,
             "_fetch_existing",
             new_callable=AsyncMock,
-            side_effect=httpx.ConnectError("Connection refused"),
+            side_effect=UpstreamError("Connection refused"),
         ):
             response = getattr(client, method)(path, **kwargs)
         assert response.status_code == 502
@@ -574,3 +574,169 @@ def test_record_not_found(
             response = getattr(client, method)(path, **kwargs)
         assert response.status_code == 404
         assert response.json()["code"] == "NotFoundError"
+
+
+CALLER_HEADERS = {
+    "Authorization": "Bearer caller-token",
+    "If-None-Match": "*",
+    "Accept-Encoding": "zstd",
+}
+LEAKED_HEADERS = {"authorization", "if-none-match", "accept-encoding"}
+
+
+class TestFetchExistingMiddlewareMode:
+    """Existing record is fetched in-process from the wrapped STAC API's routes."""
+
+    def _create(self, get_status=200, existing=None):
+        seen = []
+        app = FastAPI()
+        app.add_middleware(Cql2ValidateTransactionMiddleware)
+        _set_cql2_filter(app, Expr(ITEM_FILTER))
+
+        # Stand-in for EnforceAuthMiddleware, which sits outside the transaction
+        # middleware; the in-process GET must never need to pass through it again.
+        @app.middleware("http")
+        async def require_auth(request, call_next):
+            if "authorization" not in request.headers:
+                return JSONResponse({"code": "Unauthorized"}, status_code=401)
+            return await call_next(request)
+
+        @app.get("/collections/{collection_id}/items/{item_id}")
+        async def get_item(request: Request):
+            seen.append(request.headers)
+            request.state.set_by_get = True
+            if get_status == "raise":
+                raise RuntimeError("downstream blew up")
+            if get_status == "not-json":
+                return PlainTextResponse("not json")
+            if get_status != 200:
+                return JSONResponse({"code": "Error"}, status_code=get_status)
+            return existing or {"id": "item1", "collection": "allowed"}
+
+        @app.put("/collections/{collection_id}/items/{item_id}")
+        async def put_item(request: Request):
+            return {"state_leaked": hasattr(request.state, "set_by_get")}
+
+        @app.delete("/collections/{collection_id}/items/{item_id}")
+        async def delete_item():
+            return {"deleted": True}
+
+        return TestClient(app), seen
+
+    @pytest.mark.parametrize("method", ["put", "delete"])
+    def test_fetches_without_caller_headers(self, method):
+        """The GET reaches the route without auth, conditional, or encoding headers."""
+        client, seen = self._create()
+        kwargs = (
+            {"json": {"id": "item1", "collection": "allowed"}}
+            if method == "put"
+            else {}
+        )
+        response = getattr(client, method)(
+            "/collections/allowed/items/item1", headers=CALLER_HEADERS, **kwargs
+        )
+        assert response.status_code == 200
+        assert len(seen) == 1
+        assert not LEAKED_HEADERS & set(seen[0].keys())
+
+    def test_state_not_shared_with_caller(self):
+        """State written during the in-process GET doesn't leak into the real request."""
+        client, _ = self._create()
+        response = client.put(
+            "/collections/allowed/items/item1",
+            json={"id": "item1", "collection": "allowed"},
+            headers=CALLER_HEADERS,
+        )
+        assert response.json() == {"state_leaked": False}
+
+    @pytest.mark.parametrize(
+        "get_status,existing,expected_status,code",
+        [
+            pytest.param(
+                200,
+                {"id": "item1", "collection": "denied"},
+                404,
+                "NotFoundError",
+                id="existing-denied",
+            ),
+            pytest.param(404, None, 404, "NotFoundError", id="missing"),
+            pytest.param(500, None, 502, "UpstreamError", id="downstream-error"),
+            pytest.param("raise", None, 502, "UpstreamError", id="downstream-raises"),
+            pytest.param("not-json", None, 502, "UpstreamError", id="not-json"),
+        ],
+    )
+    def test_existing_record_outcomes(
+        self, get_status, existing, expected_status, code
+    ):
+        """Status of the in-process GET maps onto the transaction response."""
+        client, _ = self._create(get_status=get_status, existing=existing)
+        response = client.put(
+            "/collections/allowed/items/item1",
+            json={"id": "item1", "collection": "allowed"},
+            headers=CALLER_HEADERS,
+        )
+        assert response.status_code == expected_status
+        assert response.json()["code"] == code
+
+
+class TestFetchExistingProxyMode:
+    """Existing record is fetched from the upstream via the reverse proxy handler."""
+
+    def _create(self, get_status=200):
+        upstream_requests = []
+
+        def upstream(request: httpx.Request):
+            upstream_requests.append(request)
+            if request.method == "GET":
+                if get_status != 200:
+                    return httpx.Response(get_status)
+                return httpx.Response(
+                    200, json={"id": "item1", "collection": "allowed"}
+                )
+            return httpx.Response(200, json={"ok": True})
+
+        proxy = ReverseProxyHandler(
+            upstream="http://upstream",
+            client=httpx.AsyncClient(
+                transport=httpx.MockTransport(upstream), base_url="http://upstream"
+            ),
+        )
+        app = FastAPI()
+        app.add_middleware(Cql2ValidateTransactionMiddleware)
+        _set_cql2_filter(app, Expr(ITEM_FILTER))
+        app.add_api_route(
+            "/{path:path}",
+            proxy.proxy_request,
+            methods=["GET", "POST", "PUT", "PATCH", "DELETE"],
+        )
+        return TestClient(app), upstream_requests
+
+    def test_put_fetches_then_forwards(self):
+        """PUT issues a clean GET to the upstream, then forwards the PUT."""
+        client, upstream_requests = self._create()
+        response = client.put(
+            "/collections/allowed/items/item1",
+            json={"id": "item1", "collection": "allowed"},
+            headers=CALLER_HEADERS,
+        )
+        assert response.status_code == 200
+        get, put = upstream_requests
+        assert (get.method, get.url.path) == (
+            "GET",
+            "/collections/allowed/items/item1",
+        )
+        assert not LEAKED_HEADERS & set(get.headers.keys())
+        assert get.headers["host"] == "upstream"
+        assert put.method == "PUT"
+        assert put.headers["authorization"] == "Bearer caller-token"
+
+    def test_upstream_error(self):
+        """A failing upstream GET is reported as 502 and the PUT is not forwarded."""
+        client, upstream_requests = self._create(get_status=503)
+        response = client.put(
+            "/collections/allowed/items/item1",
+            json={"id": "item1", "collection": "allowed"},
+        )
+        assert response.status_code == 502
+        assert response.json()["code"] == "UpstreamError"
+        assert [r.method for r in upstream_requests] == ["GET"]