feat(BA-2257): Forward client Date header through Webserver proxy

changes/5737.fix.md

@@ -0,0 +1 @@
+Forward the client's `Date` header through `Webserver` proxy by accepting per-call header overrides on `Request.fetch()`, fixing signature-based authentication that binds the signature to the `Date` header.

src/ai/backend/client/request.py

@@ -23,6 +23,7 @@
 import appdirs
 import attrs
 from aiohttp.client import _RequestContextManager, _WSRequestContextManager
+from dateutil.parser import parse as parse_datetime
 from dateutil.tz import tzutc
 from multidict import CIMultiDict
 from yarl import URL
@@ -277,6 +278,20 @@ def _pack_content(self) -> RequestContent | aiohttp.FormData:
             return data
         return self._content
 
+    def _apply_header_overrides(self, overrides: Mapping[str, str]) -> None:
+        """
+        Apply caller-supplied header overrides on top of the auto-set headers.
+
+        If ``Date`` is overridden, parse it back into ``self.date`` so that
+        request signing uses the same value the upstream server will see.
+        Required for use cases like proxying a pre-signed request where the
+        original ``Date`` must be preserved end-to-end.
+        """
+        for key, value in overrides.items():
+            self.headers[key] = value
+            if key.lower() == "date":
+                self.date = parse_datetime(value)
+
     def _build_url(self) -> URL:
         base_url = self.config.endpoint.path.rstrip("/")
         query_path = self.path.lstrip("/") if self.path is not None and len(self.path) > 0 else ""
@@ -291,7 +306,12 @@ def _build_url(self) -> URL:
 
     # TODO: attach rate-limit information
 
-    def fetch(self, **kwargs: Any) -> FetchContextManager:
+    def fetch(
+        self,
+        *,
+        headers: Mapping[str, str] | None = None,
+        **kwargs: Any,
+    ) -> FetchContextManager:
         """
         Sends the request to the server and reads the response.
 
@@ -307,6 +327,11 @@ def fetch(self, **kwargs: Any) -> FetchContextManager:
             rqst = Request('GET', ...)
             async with rqst.fetch() as resp:
               print(await resp.text())
+
+        :param headers: Header overrides applied after the auto-populated
+            headers. Useful for proxy use cases that need to forward the
+            client's original ``Date`` header (and matching signature)
+            instead of having it refreshed.
         """
         if self.method not in self._allowed_methods:
             raise ValueError(f"Disallowed HTTP method: {self.method}")
@@ -316,6 +341,8 @@ def fetch(self, **kwargs: Any) -> FetchContextManager:
         self.headers["Date"] = self.date.isoformat()
         if self.content_type is not None and "Content-Type" not in self.headers:
             self.headers["Content-Type"] = self.content_type
+        if headers:
+            self._apply_header_overrides(headers)
         force_anonymous = kwargs.pop("anonymous", False)
 
         def _rqst_ctx_builder() -> _RequestContextManager:
@@ -345,7 +372,11 @@ def _rqst_ctx_builder() -> _RequestContextManager:
         return FetchContextManager(self.session, _rqst_ctx_builder, self._session_mode, **kwargs)
 
     def connect_websocket(
-        self, protocols: Iterable[str] = tuple(), **kwargs: Any
+        self,
+        protocols: Iterable[str] = tuple(),
+        *,
+        headers: Mapping[str, str] | None = None,
+        **kwargs: Any,
     ) -> WebSocketContextManager:
         """
         Creates a WebSocket connection.
@@ -354,6 +385,9 @@ def connect_websocket(
 
           This method only works with
           :class:`~ai.backend.client.session.AsyncSession`.
+
+        :param headers: Header overrides applied after the auto-populated
+            headers. See :meth:`fetch` for details.
         """
         if not isinstance(self.session, AsyncSession):
             raise RuntimeError("Cannot use websockets with sessions in the synchronous mode")
@@ -365,6 +399,8 @@ def connect_websocket(
         self.headers["Date"] = self.date.isoformat()
         # websocket is always a "binary" stream.
         self.content_type = "application/octet-stream"
+        if headers:
+            self._apply_header_overrides(headers)
 
         def _ws_ctx_builder() -> _WSRequestContextManager:
             full_url = self._build_url()
@@ -385,14 +421,22 @@ def _ws_ctx_builder() -> _WSRequestContextManager:
 
         return WebSocketContextManager(self.session, _ws_ctx_builder, **kwargs)
 
-    def connect_events(self, **kwargs: Any) -> SSEContextManager:
+    def connect_events(
+        self,
+        *,
+        headers: Mapping[str, str] | None = None,
+        **kwargs: Any,
+    ) -> SSEContextManager:
         """
         Creates a Server-Sent Events connection.
 
         .. warning::
 
           This method only works with
           :class:`~ai.backend.client.session.AsyncSession`.
+
+        :param headers: Header overrides applied after the auto-populated
+            headers. See :meth:`fetch` for details.
         """
         if not isinstance(self.session, AsyncSession):
             raise RuntimeError("Cannot use event streams with sessions in the synchronous mode")
@@ -403,6 +447,8 @@ def connect_events(self, **kwargs: Any) -> SSEContextManager:
             raise RuntimeError("Failed to set request date")
         self.headers["Date"] = self.date.isoformat()
         self.content_type = "application/octet-stream"
+        if headers:
+            self._apply_header_overrides(headers)
 
         def _rqst_ctx_builder() -> _RequestContextManager:
             timeout_config = aiohttp.ClientTimeout(

src/ai/backend/web/proxy.py

@@ -255,7 +255,13 @@ async def web_handler(
                     continue
                 if (value := frontend_rqst.headers.get(key)) is not None:
                     backend_rqst.headers[key] = value
-            async with backend_rqst.fetch() as backend_resp:
+            # Preserve the client's original Date header so that signature-based
+            # auth schemes that bind the signature to Date keep working through
+            # the proxy. fetch() otherwise refreshes Date unconditionally.
+            fetch_header_overrides: dict[str, str] = {}
+            if (client_date := frontend_rqst.headers.get("Date")) is not None:
+                fetch_header_overrides["Date"] = client_date
+            async with backend_rqst.fetch(headers=fetch_header_overrides) as backend_resp:
                 frontend_resp_hdrs = {
                     key: value
                     for key, value in backend_resp.headers.items()
@@ -427,8 +433,14 @@ async def web_handler_with_jwt(
                 if (value := frontend_rqst.headers.get(key)) is not None:
                     backend_rqst.headers[key] = value
 
+            # Preserve the client's original Date header so that signature-based
+            # auth schemes that bind the signature to Date keep working through
+            # the proxy. fetch() otherwise refreshes Date unconditionally.
+            fetch_header_overrides: dict[str, str] = {}
+            if (client_date := frontend_rqst.headers.get("Date")) is not None:
+                fetch_header_overrides["Date"] = client_date
             # Fetch from backend and stream response
-            async with backend_rqst.fetch() as backend_resp:
+            async with backend_rqst.fetch(headers=fetch_header_overrides) as backend_resp:
                 frontend_resp_hdrs = {
                     key: value
                     for key, value in backend_resp.headers.items()
@@ -530,7 +542,13 @@ async def web_plugin_handler(
             for key in HTTP_HEADERS_TO_FORWARD:
                 if (value := frontend_rqst.headers.get(key)) is not None:
                     backend_rqst.headers[key] = value
-            async with backend_rqst.fetch() as backend_resp:
+            # Preserve the client's original Date header so that signature-based
+            # auth schemes that bind the signature to Date keep working through
+            # the proxy. fetch() otherwise refreshes Date unconditionally.
+            fetch_header_overrides: dict[str, str] = {}
+            if (client_date := frontend_rqst.headers.get("Date")) is not None:
+                fetch_header_overrides["Date"] = client_date
+            async with backend_rqst.fetch(headers=fetch_header_overrides) as backend_resp:
                 frontend_resp_hdrs = {
                     key: value
                     for key, value in backend_resp.headers.items()

tests/unit/client/test_request.py

@@ -267,3 +267,49 @@ async def test_response_async(defconfig: APIConfig, dummy_endpoint: str) -> None
             async with rqst.fetch() as resp:
                 assert await resp.text() == '{"test": 5678}'
                 assert await resp.json() == {"test": 5678}
+
+
+async def test_fetch_preserves_date_header_override(dummy_endpoint: str) -> None:
+    fixed_date = "Tue, 02 Sep 2025 08:00:00 GMT"
+    with aioresponses() as m:
+        m.post(dummy_endpoint + "function", status=HTTPStatus.OK, body=b"")
+        async with AsyncSession():
+            rqst = Request("POST", "function")
+            async with rqst.fetch(headers={"Date": fixed_date}):
+                pass
+
+        sent_kwargs = next(iter(m.requests.values()))[0].kwargs
+        assert sent_kwargs["headers"]["Date"] == fixed_date
+
+    # The internal datetime must also reflect the override so that signing
+    # done with `self.date` matches the Date header sent on the wire.
+    assert rqst.date is not None
+    assert rqst.date.year == 2025
+    assert rqst.date.month == 9
+    assert rqst.date.day == 2
+
+
+async def test_fetch_default_date_header_when_no_override(dummy_endpoint: str) -> None:
+    with aioresponses() as m:
+        m.post(dummy_endpoint + "function", status=HTTPStatus.OK, body=b"")
+        async with AsyncSession():
+            rqst = Request("POST", "function")
+            async with rqst.fetch():
+                pass
+
+        sent_kwargs = next(iter(m.requests.values()))[0].kwargs
+        # Without override, fetch() auto-populates Date with self.date.isoformat().
+        assert rqst.date is not None
+        assert sent_kwargs["headers"]["Date"] == rqst.date.isoformat()
+
+
+async def test_fetch_passes_through_arbitrary_header_overrides(dummy_endpoint: str) -> None:
+    with aioresponses() as m:
+        m.post(dummy_endpoint + "function", status=HTTPStatus.OK, body=b"")
+        async with AsyncSession():
+            rqst = Request("POST", "function")
+            async with rqst.fetch(headers={"X-Custom-Header": "custom-value"}):
+                pass
+
+        sent_kwargs = next(iter(m.requests.values()))[0].kwargs
+        assert sent_kwargs["headers"]["X-Custom-Header"] == "custom-value"