diff --git a/toolgate/api/dispatch.py b/toolgate/api/dispatch.py index a073abc..a7e87f3 100644 --- a/toolgate/api/dispatch.py +++ b/toolgate/api/dispatch.py @@ -1,6 +1,7 @@ """Dispatch of one approved tool call to its backend (HTTP, MemoryGate, Ollama, Gemini, research, system executors).""" import hashlib +import json import os import re from typing import NoReturn @@ -114,15 +115,22 @@ def _execute_http_json(execution: dict, args: dict) -> dict: # Rendered before the client opens, so a missing body argument is refused before any send. body = _render_template(execution["body"], args) if execution.get("body") is not None else None try: - with public_client(timeout=timeout, headers=headers) as client: - response = client.request(str(execution.get("method", "GET")).upper(), url, json=body) - if 300 <= response.status_code < 400: - deny("DESTINATION_DENIED", "The upstream attempted an unapproved redirect", 502) - if response.status_code >= 400: - deny("UPSTREAM_ERROR", f"The upstream returned HTTP {response.status_code}", 502) - if len(response.content) > max_bytes: - deny("RESPONSE_TOO_LARGE", "The upstream response exceeded this tool's byte limit", 502) - value = response.json() + method = str(execution.get("method", "GET")).upper() + with ( + public_client(timeout=timeout, headers=headers) as client, + client.stream(method, url, json=body) as response, + ): + if 300 <= response.status_code < 400: + deny("DESTINATION_DENIED", "The upstream attempted an unapproved redirect", 502) + if response.status_code >= 400: + deny("UPSTREAM_ERROR", f"The upstream returned HTTP {response.status_code}", 502) + # Read up to the cap and stop: a huge or endless body never sits in memory whole. + content = bytearray() + for chunk in response.iter_bytes(): + content.extend(chunk) + if len(content) > max_bytes: + deny("RESPONSE_TOO_LARGE", "The upstream response exceeded this tool's byte limit", 502) + value = json.loads(bytes(content)) except HTTPException: raise except (httpx.HTTPError, ValueError) as exc: diff --git a/toolgate/tests/test_http_response_cap.py b/toolgate/tests/test_http_response_cap.py new file mode 100644 index 0000000..edf606f --- /dev/null +++ b/toolgate/tests/test_http_response_cap.py @@ -0,0 +1,46 @@ +"""An HTTP tool stops reading at its byte cap instead of downloading the whole body first +(hunt toolgate-1 F3).""" + +import httpx +import pytest + +from toolgate.api import dispatch, invocation +from toolgate.core import control_plane + +CHUNK = 64 * 1024 + + +@pytest.fixture +def tool(tmp_path, monkeypatch): + monkeypatch.setattr(control_plane, "DB_PATH", tmp_path / "gate.db") + monkeypatch.setattr(dispatch, "public_url", lambda url: True) + return {"id": "http", "name": "HTTP", "version": 1, "authorization": "auto", "status": "active", + "inputs": [], "outputs": [], "execution": { + "type": "http_json", "billing": {"mode": "free"}, "url": "https://api.example.com/x", + "allowed_hosts": ["api.example.com"], "max_response_bytes": 4 * CHUNK}} + + +def upstream(monkeypatch, body): + monkeypatch.setattr(dispatch, "public_client", lambda timeout, headers: httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, content=body)), + headers=headers)) + + +def test_a_huge_body_is_cut_off_just_past_the_cap(tool, monkeypatch): + pulled = [] + + def endless(): + for _ in range(400): # 25 MiB if it were all read + pulled.append(CHUNK) + yield b"x" * CHUNK + + upstream(monkeypatch, endless()) + result = invocation.invoke_tool(tool, {}, "Actor", actor_id="agent", action_id="huge") + assert result["code"] == "OUTCOME_UNKNOWN" # sent, so the upstream may have acted + assert len(pulled) <= 6 + + +def test_a_body_within_the_cap_is_parsed(tool, monkeypatch): + upstream(monkeypatch, iter([b'{"items": ', b'[1, 2]}'])) + result = invocation.invoke_tool(tool, {}, "Actor", actor_id="agent", action_id="small") + assert (result["code"], result["result"]["result"]) == ("OK", {"items": [1, 2]}) diff --git a/toolgate/tests/test_tool_secret_refs.py b/toolgate/tests/test_tool_secret_refs.py index 690d47f..514c3ae 100644 --- a/toolgate/tests/test_tool_secret_refs.py +++ b/toolgate/tests/test_tool_secret_refs.py @@ -49,7 +49,7 @@ def __enter__(self): def __exit__(self, *_exc): return False - def request(self, method, url, json=None): + def stream(self, method, url, json=None): raise dispatch.httpx.ConnectError("stop here") monkeypatch.setattr(dispatch, "public_url", lambda _url: True) diff --git a/toolgate/tests/test_tool_settings.py b/toolgate/tests/test_tool_settings.py index ad71bca..8c44452 100644 --- a/toolgate/tests/test_tool_settings.py +++ b/toolgate/tests/test_tool_settings.py @@ -5,7 +5,7 @@ import sqlite3 import time -from contextlib import closing, contextmanager +from contextlib import closing, contextmanager, nullcontext from unittest.mock import Mock import pytest @@ -257,9 +257,8 @@ def test_web_api_memory_and_ai_calls_wait_no_longer_than_the_limit(store, monkey @contextmanager def client(timeout, headers): waits.append(("http_json", timeout)) - answer = Mock(status_code=200, content=b"{}") - answer.json.return_value = {} - yield Mock(request=Mock(return_value=answer)) + answer = Mock(status_code=200, iter_bytes=Mock(return_value=iter([b"{}"]))) + yield Mock(stream=Mock(return_value=nullcontext(answer))) def post(*_args, timeout, **_kwargs): waits.append(("post", timeout))