Reject encoded discovery responses and bound payload copies

Author: FrancescAlted

src/blosc2/caterva2_url.py

@@ -144,13 +144,15 @@ def discover_service(value, *, required=False, auth_token=None):
     auth_token = discovery_auth_token(value, required=required, auth_token=auth_token)
     url = _server_url(value.rstrip("/"), "api/roots")
     client = _sync_client()
+    headers = dict(_auth_headers(auth_token) or {})
+    headers["Accept-Encoding"] = "identity"
     deadline = time.monotonic() + 3
     for _ in range(4):
         remaining = deadline - time.monotonic()
         if remaining <= 0:
             raise TimeoutError("Caterva2 service discovery exceeded 3 seconds")
         with client.stream(
-            "GET", url, headers=_auth_headers(auth_token), timeout=remaining, follow_redirects=False
+            "GET", url, headers=headers, timeout=remaining, follow_redirects=False
         ) as response:
             if response.status_code in {301, 302, 303, 307, 308}:
                 from urllib.parse import urljoin
@@ -176,13 +178,16 @@ def discover_service(value, *, required=False, auth_token=None):
                     response.raise_for_status()
                 return None
             response.raise_for_status()
+            # Reject before iter_bytes can decompress an untrusted response.
+            if response.headers.get("content-encoding", "identity").strip().lower() != "identity":
+                raise ValueError("Encoded Caterva2 roots responses are unsupported")
             payload = bytearray()
             for chunk in response.iter_bytes():
                 if time.monotonic() > deadline:
                     raise TimeoutError("Caterva2 service discovery exceeded 3 seconds")
-                payload.extend(chunk)
-                if len(payload) > 1 << 20:
+                if len(payload) + len(chunk) > 1 << 20:
                     raise ValueError("Caterva2 roots response exceeds 1 MiB")
+                payload.extend(chunk)
             import json
 
             try:

tests/test_caterva2_access.py

@@ -271,6 +271,70 @@ def redirect(request):
     assert seen == ["https://host/demo/api/roots", "https://host/demo/api/roots/"]
 
 
+@pytest.mark.parametrize("encoding", ["gzip", "br", "deflate"])
+def test_discovery_rejects_encoding_before_reading(encoding, monkeypatch):
+    import httpx
+
+    from blosc2.caterva2_url import discover_service
+
+    class UnreadableStream(httpx.SyncByteStream):
+        def __iter__(self):
+            pytest.fail("Encoded response was read or decompressed")
+            yield b""  # Make this a stream iterator without reading any payload.
+
+    def respond(request):
+        assert request.headers["accept-encoding"] == "identity"
+        return httpx.Response(200, headers={"content-encoding": encoding}, stream=UnreadableStream())
+
+    with httpx.Client(transport=httpx.MockTransport(respond)) as client:
+        monkeypatch.setattr(blosc2.c2array, "_sync_client", lambda: client)
+        with pytest.raises(ValueError, match="Encoded"):
+            discover_service("https://host")
+
+
+@pytest.mark.parametrize("fragmented", [False, True])
+def test_discovery_checks_size_before_copying(fragmented, monkeypatch):
+    import httpx
+
+    import blosc2.caterva2_url as discovery
+
+    class BoundedPayload(bytearray):
+        def extend(self, data):
+            assert len(self) + len(data) <= 1 << 20, "Oversized fragment copied before checking"
+            super().extend(data)
+
+    class Fragments(httpx.SyncByteStream):
+        def __iter__(self):
+            if fragmented:
+                yield b" " * (1 << 19)
+                yield b" " * ((1 << 19) + 1)
+            else:
+                yield b" " * ((1 << 20) + 1)
+
+    monkeypatch.setattr(discovery, "bytearray", BoundedPayload, raising=False)
+    with httpx.Client(
+        transport=httpx.MockTransport(lambda _: httpx.Response(200, stream=Fragments()))
+    ) as client:
+        monkeypatch.setattr(blosc2.c2array, "_sync_client", lambda: client)
+        with pytest.raises(ValueError, match="exceeds 1 MiB"):
+            discovery.discover_service("https://host")
+
+
+def test_discovery_accepts_identity_at_exact_byte_limit(monkeypatch):
+    import httpx
+
+    from blosc2.caterva2_url import discover_service
+
+    payload = b"{}" + b" " * ((1 << 20) - 2)
+    with httpx.Client(
+        transport=httpx.MockTransport(
+            lambda _: httpx.Response(200, headers={"content-encoding": "identity"}, content=payload)
+        )
+    ) as client:
+        monkeypatch.setattr(blosc2.c2array, "_sync_client", lambda: client)
+        assert discover_service("https://host") == {}
+
+
 def test_repository_freezes_auth_context(caterva2_source):  # noqa: F811
     base, _, _, stats = caterva2_source
     stats["roots"]["@broken"] = {"name": "@broken"}