fix(craft): bound recorded-response memory, editor-id links, draft view

An undeclared-length recorded response streams through a capped observer
instead of buffering whole, and decoding is bounded (zlib with a capped
output) so neither a chunked origin nor a compression bomb can grow proxy
memory before the refine cap applies. Google editor APIs link through
their product id keys (documentId, spreadsheetId, presentationId, which
carry no mimeType), and Gmail draft receipts link the drafts view since a
draft is not in #all until sent.

Author: nmgarza5

backend/onyx/sandbox_proxy/addons/gate.py

@@ -11,6 +11,7 @@
 import operator
 import socket
 import threading
+import zlib
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Protocol
@@ -217,6 +218,43 @@ class _ApprovalGrant:
 RECEIPT_REFINE_MAX_BODY_BYTES = 1024 * 1024
 
 
+class _CappedBodyCapture:
+    """Streams a recorded flow's response through untouched while keeping a
+    copy for the extractors, up to the refine cap. Oversize drops the copy."""
+
+    def __init__(self) -> None:
+        self.body = bytearray()
+        self.overflowed = False
+
+    def __call__(self, chunk: bytes) -> bytes:
+        if chunk and not self.overflowed:
+            self.body.extend(chunk)
+            if len(self.body) > RECEIPT_REFINE_MAX_BODY_BYTES:
+                self.overflowed = True
+                self.body = bytearray()
+        return chunk
+
+
+def _bounded_decode(data: bytes, content_encoding: str) -> bytes | None:
+    """Undo Content-Encoding without trusting it: output is capped so a small
+    compressed body cannot balloon in memory. Unsupported or broken codings
+    refine nothing rather than risk an unbounded decode."""
+    coding = content_encoding.strip().lower()
+    if coding in ("", "identity"):
+        return data if len(data) <= RECEIPT_REFINE_MAX_BODY_BYTES else None
+    if coding not in ("gzip", "deflate"):
+        return None
+    wbits = 16 + zlib.MAX_WBITS if coding == "gzip" else zlib.MAX_WBITS
+    try:
+        decompressor = zlib.decompressobj(wbits=wbits)
+        decoded = decompressor.decompress(data, RECEIPT_REFINE_MAX_BODY_BYTES + 1)
+    except zlib.error:
+        return None
+    if len(decoded) > RECEIPT_REFINE_MAX_BODY_BYTES or decompressor.unconsumed_tail:
+        return None
+    return decoded
+
+
 class ParkedApprovals:
     """Approvals the proxy is currently parked on, grouped by tenant.
 
@@ -376,15 +414,18 @@ def responseheaders(self, flow: http.HTTPFlow) -> None:
         if flow.response is None:
             return
         # A recorded flow's response feeds the receipt extractors: buffer it
-        # unless declared oversize (chunked declares nothing and must buffer).
-        # The decoded-length cap in `response` bounds the read either way.
+        # when its declared size fits, otherwise stream through a capped
+        # observer so chunked or lying origins can never grow proxy memory.
         if RECEIPT_FLOW_KEY in flow.metadata:
             declared = flow.response.headers.get("content-length")
             try:
                 length = int(declared) if declared is not None else None
             except ValueError:
                 length = None
-            if length is None or length <= RECEIPT_REFINE_MAX_BODY_BYTES:
+            if length is not None and length <= RECEIPT_REFINE_MAX_BODY_BYTES:
+                return
+            if length is None:
+                flow.response.stream = _CappedBodyCapture()
                 return
         flow.response.stream = True
 
@@ -627,12 +668,21 @@ async def response(self, flow: http.HTTPFlow) -> None:
             return
         confirmed = flow.response is not None and flow.response.status_code < 400
         response_body: bytes | None = None
-        if flow.response is not None and not flow.response.stream:
-            # get_content undoes Content-Encoding. raw_content is still
-            # compressed on the wire and would parse as garbage.
-            decoded = flow.response.get_content(strict=False)
-            if decoded is not None and len(decoded) <= RECEIPT_REFINE_MAX_BODY_BYTES:
-                response_body = decoded
+        if flow.response is not None:
+            stream = flow.response.stream
+            wire: bytes | None = None
+            if isinstance(stream, _CappedBodyCapture):
+                wire = bytes(stream.body) if not stream.overflowed else None
+            elif not stream:
+                raw = flow.response.raw_content
+                if raw is not None and len(raw) <= RECEIPT_REFINE_MAX_BODY_BYTES:
+                    wire = raw
+            if wire is not None:
+                # The wire body is still Content-Encoding compressed and
+                # would parse as garbage undecoded.
+                response_body = _bounded_decode(
+                    wire, flow.response.headers.get("content-encoding", "")
+                )
         await self._finalize_recorded(
             recorded,
             ReceiptStatus.CONFIRMED if confirmed else ReceiptStatus.FAILED,

backend/onyx/sandbox_proxy/receipt_extractors.py

@@ -137,7 +137,7 @@ def response_facts(action_type: str, response_body: bytes | None) -> ResponseFac
     if action_type.startswith("gdrive."):
         return _gdrive_facts(body)
     if action_type.startswith("gmail."):
-        return _gmail_facts(body)
+        return _gmail_facts(action_type, body)
     if action_type.startswith("linear."):
         return _linear_facts(body)
     return ResponseFacts()
@@ -185,25 +185,38 @@ def _slack_facts(action_type: str, body: dict[str, Any]) -> ResponseFacts:
     return ResponseFacts(link=link, operation_key=operation_key)
 
 
+# The dedicated editor APIs name their id after the product and return no
+# mimeType, so the key alone picks the link.
+_GDRIVE_LINKS_BY_ID_KEY = {
+    "documentId": "https://docs.google.com/document/d/{id}",
+    "spreadsheetId": "https://docs.google.com/spreadsheets/d/{id}",
+    "presentationId": "https://docs.google.com/presentation/d/{id}",
+}
+
+
 def _gdrive_facts(body: dict[str, Any]) -> ResponseFacts:
-    file_id = _safe_id(
-        body.get("id") or body.get("documentId") or body.get("spreadsheetId")
-    )
+    for key, template in _GDRIVE_LINKS_BY_ID_KEY.items():
+        editor_id = _safe_id(body.get(key))
+        if editor_id is not None:
+            return ResponseFacts(link=template.format(id=editor_id))
+    file_id = _safe_id(body.get("id"))
     if file_id is None:
         return ResponseFacts()
     mime = body.get("mimeType")
     template = _GDRIVE_LINKS_BY_MIME.get(str(mime), _GDRIVE_GENERIC_LINK)
     return ResponseFacts(link=template.format(id=file_id))
 
 
-def _gmail_facts(body: dict[str, Any]) -> ResponseFacts:
+def _gmail_facts(action_type: str, body: dict[str, Any]) -> ResponseFacts:
     # Drafts nest the message, sends carry it at the top level.
     nested = body.get("message")
     message = nested if isinstance(nested, dict) else body
     message_id = _safe_id(message.get("id"))
-    if message_id is not None:
-        return ResponseFacts(link=f"https://mail.google.com/mail/u/0/#all/{message_id}")
-    return ResponseFacts()
+    if message_id is None:
+        return ResponseFacts()
+    # A draft is not in #all until sent, so link the drafts view instead.
+    view = "drafts" if action_type.startswith("gmail.drafts.") else "all"
+    return ResponseFacts(link=f"https://mail.google.com/mail/u/0/#{view}/{message_id}")
 
 
 def _linear_facts(body: dict[str, Any]) -> ResponseFacts:

backend/tests/unit/sandbox_proxy/test_gate_receipts.py

@@ -28,6 +28,7 @@
 from onyx.sandbox_proxy.addons.gate import (
     RECEIPT_FLOW_KEY,
     GateAddon,
+    _CappedBodyCapture,
     _IdentityResolver,
 )
 from onyx.sandbox_proxy.credential_injection import CredentialInjectionDispatcher
@@ -226,24 +227,31 @@ def test_oversized_decoded_body_refines_nothing(
     asyncio.run(addon.response(flow))
     assert calls[0]["response_body"] is None
 
-    # An undecodable body falls back to its raw bytes rather than raising;
-    # the extractors then simply fail to parse it.
+    # An undecodable body refines nothing rather than raising.
     flow = _recorded_flow(200)
     flow.response.headers["content-encoding"] = "gzip"
     flow.response.raw_content = b"not gzip"
     asyncio.run(addon.response(flow))
-    assert calls[1]["response_body"] == b"not gzip"
+    assert calls[1]["response_body"] is None
 
 
-def test_responseheaders_buffers_recorded_flows_without_declared_length() -> None:
+def test_responseheaders_caps_undeclared_recorded_flows(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     addon = _addon()
 
-    # Chunked responses declare no length, but a recorded flow still buffers
-    # so the extractors can read it. Anything else streams as usual.
-    recorded = _recorded_flow(200)
-    del recorded.response.headers["content-length"]
-    addon.responseheaders(recorded)
-    assert recorded.response.stream is False
+    # Chunked responses declare no length, so a recorded flow streams through
+    # a capped observer the response hook reads. Small declared bodies buffer,
+    # oversize ones and unrecorded flows stream unobserved.
+    chunked = _recorded_flow(200)
+    del chunked.response.headers["content-length"]
+    addon.responseheaders(chunked)
+    capture = chunked.response.stream
+    assert isinstance(capture, _CappedBodyCapture)
+
+    small = _recorded_flow(200)
+    addon.responseheaders(small)
+    assert small.response.stream is False
 
     declared_oversize = _recorded_flow(200)
     declared_oversize.response.headers["content-length"] = str(2 * 1024 * 1024)
@@ -254,6 +262,35 @@ def test_responseheaders_buffers_recorded_flows_without_declared_length() -> Non
     addon.responseheaders(unrecorded)
     assert unrecorded.response is not None and unrecorded.response.stream is True
 
+    # The observer passes chunks through untouched and the response hook
+    # reads the copy, so chunked provider responses still refine.
+    calls = _finalize_capture(monkeypatch)
+    payload = b'{"ok": true}'
+    assert capture(payload) == payload
+    assert capture(b"") == b""
+    asyncio.run(addon.response(chunked))
+    assert calls[0]["response_body"] == payload
+
+
+def test_capped_capture_drops_oversized_streams(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    addon = _addon()
+    calls = _finalize_capture(monkeypatch)
+    flow = _recorded_flow(200)
+    del flow.response.headers["content-length"]
+    addon.responseheaders(flow)
+    capture = flow.response.stream
+    assert isinstance(capture, _CappedBodyCapture)
+
+    chunk = b"x" * (512 * 1024)
+    for _ in range(5):
+        assert capture(chunk) == chunk
+    assert capture.overflowed and not capture.body
+
+    asyncio.run(addon.response(flow))
+    assert calls[0]["response_body"] is None
+
 
 def test_always_policy_writes_still_record(
     monkeypatch: pytest.MonkeyPatch,

backend/tests/unit/sandbox_proxy/test_receipt_extractors.py

@@ -7,6 +7,7 @@
 """
 
 import json
+from typing import Any
 from urllib.parse import urlencode
 
 import pytest
@@ -21,7 +22,7 @@
 )
 
 
-def _body(payload: dict) -> bytes:
+def _body(payload: dict[str, Any]) -> bytes:
     return json.dumps(payload).encode()
 
 
@@ -129,10 +130,24 @@ def test_gdrive_links_follow_the_mime_type(mime: str, prefix: str) -> None:
 def test_gmail_send_and_draft_shapes_both_link() -> None:
     sent = response_facts("gmail.messages.send", _body({"id": "m1", "threadId": "t1"}))
     assert sent.link == "https://mail.google.com/mail/u/0/#all/m1"
+    # A draft is not in #all until sent, so its link opens the drafts view.
     draft = response_facts(
         "gmail.drafts.create", _body({"id": "d1", "message": {"id": "m2"}})
     )
-    assert draft.link == "https://mail.google.com/mail/u/0/#all/m2"
+    assert draft.link == "https://mail.google.com/mail/u/0/#drafts/m2"
+
+
+@pytest.mark.parametrize(
+    "key,prefix",
+    [
+        ("documentId", "https://docs.google.com/document/d/"),
+        ("spreadsheetId", "https://docs.google.com/spreadsheets/d/"),
+        ("presentationId", "https://docs.google.com/presentation/d/"),
+    ],
+)
+def test_gdrive_editor_api_ids_link_without_a_mime_type(key: str, prefix: str) -> None:
+    facts = response_facts("gdrive.slides.create", _body({key: "p123"}))
+    assert facts.link == f"{prefix}p123"
 
 
 def test_linear_errors_fail_and_only_shaped_urls_pass() -> None: