From 4ee776af1ca46177cd733430870cfa5618e512e3 Mon Sep 17 00:00:00 2001 From: Adam Cohen Hillel Date: Sun, 5 Jul 2026 18:07:48 -0700 Subject: [PATCH 1/2] Stream model responses through the broker instead of buffering Detect streaming requests via Accept: text/event-stream or "stream": true and relay the upstream body chunk-by-chunk with per-chunk flushes, closing the connection to delimit the response. Auth, rate-limit, and privacy checks still run before proxying; the smoke test now covers an SSE case against the fake provider and asserts chunks arrive incrementally. Closes #81 Co-Authored-By: Claude Opus 4.7 --- scripts/smoke-test-model-broker.sh | 74 ++++++++++++++++++- .../model-broker/openphone_model_broker.py | 60 +++++++++++++++ 2 files changed, 132 insertions(+), 2 deletions(-) diff --git a/scripts/smoke-test-model-broker.sh b/scripts/smoke-test-model-broker.sh index e8931ca..176ea84 100755 --- a/scripts/smoke-test-model-broker.sh +++ b/scripts/smoke-test-model-broker.sh @@ -65,6 +65,7 @@ import http.server import json import pathlib import sys +import time state_path = pathlib.Path(sys.argv[1]) @@ -72,7 +73,23 @@ state_path = pathlib.Path(sys.argv[1]) class Handler(http.server.BaseHTTPRequestHandler): def do_POST(self): length = int(self.headers.get("Content-Length", "0")) - self.rfile.read(length) + body = self.rfile.read(length) + try: + payload = json.loads(body.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + payload = {} + if isinstance(payload, dict) and payload.get("stream") is True: + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Connection", "close") + self.end_headers() + for index in range(3): + self.wfile.write(f"data: chunk-{index}\n\n".encode("utf-8")) + self.wfile.flush() + time.sleep(0.4) + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + return attempts = int(state_path.read_text(encoding="utf-8")) if state_path.exists() else 0 attempts += 1 state_path.write_text(str(attempts), encoding="utf-8") @@ -318,6 +335,59 @@ expect_status 200 \ -H 'Content-Type: application/json' \ --data '{"model":"gpt-4.1-mini","metadata":{"openphone_task":"true","risk_flags":""},"input":[{"role":"user","content":[{"type":"input_text","text":"hello"}]}]}' +# Streaming: the broker must relay SSE chunks incrementally instead of +# buffering the full upstream response before replying. +python3 - <<'PY' "$port" "$static_token" +import http.client +import json +import sys +import time + +port, token = sys.argv[1:] +body = json.dumps({ + "model": "gpt-4.1-mini", + "stream": True, + "metadata": {"openphone_task": "true", "risk_flags": ""}, + "input": [{"role": "user", "content": [{"type": "input_text", "text": "hi"}]}], +}) +connection = http.client.HTTPConnection("127.0.0.1", int(port), timeout=30) +connection.request( + "POST", + "/v1/responses", + body=body, + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "Accept": "text/event-stream", + }, +) +response = connection.getresponse() +if response.status != 200: + raise SystemExit(f"streaming request failed: HTTP {response.status}") +content_type = response.headers.get("Content-Type", "") +if not content_type.startswith("text/event-stream"): + raise SystemExit(f"streaming response has wrong content-type: {content_type}") + +chunks = [] +arrival_times = [] +while True: + chunk = response.read1(65536) + if not chunk: + break + chunks.append(chunk) + arrival_times.append(time.monotonic()) +payload = b"".join(chunks).decode("utf-8") +for expected in ("data: chunk-0", "data: chunk-1", "data: chunk-2", "data: [DONE]"): + if expected not in payload: + raise SystemExit(f"streaming response missing {expected!r}: {payload!r}") +if len(arrival_times) < 2: + raise SystemExit("streaming response arrived as a single buffered chunk") +spread = arrival_times[-1] - arrival_times[0] +if spread < 0.2: + raise SystemExit(f"streaming chunks were not delivered incrementally (spread={spread:.3f}s)") +print(f"streaming smoke case passed ({len(chunks)} chunks over {spread:.2f}s)") +PY + expect_status 429 \ -X POST "http://127.0.0.1:$port/v1/responses" \ -H "Authorization: Bearer $static_token" \ @@ -374,7 +444,7 @@ for event in events: if "body" in event or "request" in event or "response" in event: raise SystemExit(f"audit event contains body-like key: {event}") proxied = [event for event in events if event.get("outcome") == "proxied"] -if not proxied or proxied[-1].get("provider_attempts") != 2: +if not any(event.get("provider_attempts") == 2 for event in proxied): raise SystemExit("broker did not record retried provider forwarding") PY diff --git a/services/model-broker/openphone_model_broker.py b/services/model-broker/openphone_model_broker.py index cbfd268..d6aa964 100644 --- a/services/model-broker/openphone_model_broker.py +++ b/services/model-broker/openphone_model_broker.py @@ -211,6 +211,7 @@ def do_POST(self) -> None: # noqa: N802 body_size, started_at, model, + stream=self._wants_stream(payload), ) return @@ -326,6 +327,7 @@ def _proxy_to_openai( body_size: int, started_at: float, model: str | None, + stream: bool = False, ) -> None: headers = { "Authorization": f"Bearer {self.server.config.api_key}", @@ -338,6 +340,17 @@ def _proxy_to_openai( attempts += 1 try: with urllib.request.urlopen(request, timeout=120) as response: + if stream: + self._stream_provider_response( + response, + token_subject, + body_size, + started_at, + endpoint, + model, + attempts, + ) + return response_body = response.read() self._audit_request( "proxied", @@ -404,6 +417,53 @@ def _proxy_to_openai( ) return + def _wants_stream(self, payload: dict[str, object]) -> bool: + accept = self.headers.get("Accept", "") + if "text/event-stream" in accept.lower(): + return True + return payload.get("stream") is True + + def _stream_provider_response( + self, + response: object, + token_subject: str, + body_size: int, + started_at: float, + endpoint: str, + model: str | None, + attempts: int, + ) -> None: + """Relay the upstream body chunk-by-chunk instead of buffering it. + + The response is delimited by connection close (HTTP/1.0 semantics), + which every SSE-capable client handles; no Content-Length is sent. + """ + self._audit_request( + "proxied", + token_subject, + response.status, + body_size, + started_at, + endpoint, + model, + provider_attempts=attempts, + ) + self.send_response(response.status) + self.send_header( + "Content-Type", + response.headers.get("Content-Type", "text/event-stream"), + ) + self.send_header("Cache-Control", "no-store") + self.send_header("Connection", "close") + self.end_headers() + self.close_connection = True + while True: + chunk = response.read1(65536) + if not chunk: + break + self.wfile.write(chunk) + self.wfile.flush() + def _should_retry_provider(self, status: int, attempts: int) -> bool: if attempts > self.server.config.provider_max_retries: return False From e9b61b0bbca2a149fdee89d8a6c882e2795658e9 Mon Sep 17 00:00:00 2001 From: Adam Cohen Hillel Date: Sun, 5 Jul 2026 18:08:17 -0700 Subject: [PATCH 2/2] Add pytest coverage for broker security primitives Cover signed-token expiry and tampering, wrong HMAC key versions, device attestation clock-skew boundaries, rate limiter concurrency and stale-key eviction, admin/static token compare_digest paths, and adversarial count_input_images payloads. Wire the suite into the repo-checks CI job. Closes #87 Co-Authored-By: Claude Opus 4.7 --- .github/workflows/ci.yml | 5 + services/model-broker/tests/conftest.py | 8 + services/model-broker/tests/test_broker.py | 371 +++++++++++++++++++++ 3 files changed, 384 insertions(+) create mode 100644 services/model-broker/tests/conftest.py create mode 100644 services/model-broker/tests/test_broker.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bdd7242..3d77eb7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -45,5 +45,10 @@ jobs: - name: Run OpenPhone checks run: ./scripts/check.sh + - name: Run model broker unit tests + run: | + pip install pytest + python3 -m pytest services/model-broker/tests -q + - name: Check whitespace run: git diff --check diff --git a/services/model-broker/tests/conftest.py b/services/model-broker/tests/conftest.py new file mode 100644 index 0000000..bb5caa2 --- /dev/null +++ b/services/model-broker/tests/conftest.py @@ -0,0 +1,8 @@ +"""Make the stdlib-only broker module importable from the tests directory.""" + +import pathlib +import sys + +BROKER_DIR = pathlib.Path(__file__).resolve().parent.parent +if str(BROKER_DIR) not in sys.path: + sys.path.insert(0, str(BROKER_DIR)) diff --git a/services/model-broker/tests/test_broker.py b/services/model-broker/tests/test_broker.py new file mode 100644 index 0000000..2265831 --- /dev/null +++ b/services/model-broker/tests/test_broker.py @@ -0,0 +1,371 @@ +"""Unit tests for the OpenPhone model broker security primitives. + +The broker is deliberately stdlib-only; these tests import the module +directly and exercise the token, attestation, rate-limiting, and privacy +logic without starting an HTTP server. +""" + +import hashlib +import hmac +import threading +import time +import types + +import openphone_model_broker as broker + + +SECRET = b"unit-test-token-secret" +OTHER_SECRET = b"unit-test-token-secret-v2" + + +# --------------------------------------------------------------------------- +# Signed session tokens +# --------------------------------------------------------------------------- + + +class TestSignedTokens: + def test_valid_token_round_trip(self): + token = broker.mint_signed_token(SECRET, "device-1", ttl_seconds=60) + assert broker.validate_signed_token(token, SECRET) == "signed:device-1" + + def test_expired_token_rejected(self): + expires_at = int(time.time()) - 10 + encoded_subject = broker._b64url_encode(b"device-1") + nonce = "unit-nonce" + signature = broker._token_signature(SECRET, expires_at, encoded_subject, nonce) + token = f"op1.{expires_at}.{encoded_subject}.{nonce}.{signature}" + assert broker.validate_signed_token(token, SECRET) is None + + def test_token_valid_until_expiry_boundary(self): + # A token expiring in the future must validate right up to expiry. + token = broker.mint_signed_token(SECRET, "device-1", ttl_seconds=2) + assert broker.validate_signed_token(token, SECRET) is not None + + def test_bit_flipped_subject_rejected(self): + token = broker.mint_signed_token(SECRET, "device-1", ttl_seconds=60) + prefix, expires_at, encoded_subject, nonce, signature = token.split(".") + flipped = ("B" if encoded_subject[0] != "B" else "C") + encoded_subject[1:] + tampered = ".".join([prefix, expires_at, flipped, nonce, signature]) + assert broker.validate_signed_token(tampered, SECRET) is None + + def test_bit_flipped_expiry_rejected(self): + token = broker.mint_signed_token(SECRET, "device-1", ttl_seconds=60) + prefix, expires_at, encoded_subject, nonce, signature = token.split(".") + extended = str(int(expires_at) + 86400) + tampered = ".".join([prefix, extended, encoded_subject, nonce, signature]) + assert broker.validate_signed_token(tampered, SECRET) is None + + def test_bit_flipped_signature_rejected(self): + token = broker.mint_signed_token(SECRET, "device-1", ttl_seconds=60) + prefix, expires_at, encoded_subject, nonce, signature = token.split(".") + flipped = ("A" if signature[-1] != "A" else "B") + signature[1:] + tampered = ".".join([prefix, expires_at, encoded_subject, nonce, flipped]) + assert broker.validate_signed_token(tampered, SECRET) is None + + def test_wrong_hmac_key_rejected(self): + # A token minted under a rotated-out key version must not validate. + token = broker.mint_signed_token(SECRET, "device-1", ttl_seconds=60) + assert broker.validate_signed_token(token, OTHER_SECRET) is None + + def test_malformed_tokens_rejected(self): + assert broker.validate_signed_token("", SECRET) is None + assert broker.validate_signed_token("op1.only.three.parts", SECRET) is None + assert broker.validate_signed_token("op2.1.a.b.c", SECRET) is None + assert broker.validate_signed_token("op1.not-a-number.a.b.c", SECRET) is None + + def test_mint_requires_positive_ttl(self): + try: + broker.mint_signed_token(SECRET, "device-1", ttl_seconds=0) + except ValueError: + pass + else: + raise AssertionError("mint_signed_token accepted a non-positive TTL") + + +# --------------------------------------------------------------------------- +# Device attestation clock skew +# --------------------------------------------------------------------------- + + +ATTESTATION_SECRET = b"unit-test-attestation-secret" +ATTESTATION_SUBJECT = "unit-device" + + +def _attestation_server(max_skew_seconds: int) -> types.SimpleNamespace: + """Build a stand-in for BrokerServer exposing only .config.""" + config = types.SimpleNamespace( + device_attestation_secrets={ATTESTATION_SUBJECT: ATTESTATION_SECRET}, + device_attestation_max_skew_seconds=max_skew_seconds, + ) + return types.SimpleNamespace(config=config) + + +def _attestation_payload(timestamp: int, nonce: str = "unit-nonce") -> dict: + body = f"{ATTESTATION_SUBJECT}.{timestamp}.{nonce}".encode("utf-8") + signature = hmac.new(ATTESTATION_SECRET, body, hashlib.sha256).hexdigest() + return { + "attestation_timestamp": timestamp, + "attestation_nonce": nonce, + "attestation_signature": signature, + } + + +class TestDeviceAttestation: + def _verify(self, server, subject, payload): + return broker.BrokerServer.verify_device_attestation(server, subject, payload) + + def test_timestamp_inside_skew_window_accepted(self): + server = _attestation_server(max_skew_seconds=300) + payload = _attestation_payload(int(time.time()) - 250) + assert self._verify(server, ATTESTATION_SUBJECT, payload) is None + + def test_future_timestamp_inside_skew_window_accepted(self): + server = _attestation_server(max_skew_seconds=300) + payload = _attestation_payload(int(time.time()) + 250) + assert self._verify(server, ATTESTATION_SUBJECT, payload) is None + + def test_timestamp_outside_skew_window_rejected(self): + server = _attestation_server(max_skew_seconds=300) + payload = _attestation_payload(int(time.time()) - 301) + assert self._verify(server, ATTESTATION_SUBJECT, payload) == "attestation_expired" + + def test_future_timestamp_outside_skew_window_rejected(self): + server = _attestation_server(max_skew_seconds=300) + payload = _attestation_payload(int(time.time()) + 301) + assert self._verify(server, ATTESTATION_SUBJECT, payload) == "attestation_expired" + + def test_tampered_signature_rejected(self): + server = _attestation_server(max_skew_seconds=300) + payload = _attestation_payload(int(time.time())) + signature = payload["attestation_signature"] + payload["attestation_signature"] = ("0" if signature[0] != "0" else "1") + signature[1:] + assert self._verify(server, ATTESTATION_SUBJECT, payload) == "attestation_invalid" + + def test_missing_fields_require_attestation(self): + server = _attestation_server(max_skew_seconds=300) + assert self._verify(server, ATTESTATION_SUBJECT, {}) == "attestation_required" + payload = _attestation_payload(int(time.time())) + del payload["attestation_signature"] + assert self._verify(server, ATTESTATION_SUBJECT, payload) == "attestation_required" + + def test_subject_without_secret_skips_attestation(self): + server = _attestation_server(max_skew_seconds=300) + assert self._verify(server, "unknown-device", {}) is None + + +# --------------------------------------------------------------------------- +# Rate limiter concurrency and eviction +# --------------------------------------------------------------------------- + + +class TestRateLimiterConcurrency: + def test_exactly_max_events_allowed_under_contention(self): + max_events = 50 + thread_count = 16 + attempts_per_thread = 25 # 400 total attempts against a budget of 50 + limiter = broker.RateLimiter(max_events, window_seconds=60) + allowed = [] + allowed_lock = threading.Lock() + start_barrier = threading.Barrier(thread_count) + + def hammer(): + start_barrier.wait() + local_allowed = 0 + for _ in range(attempts_per_thread): + if limiter.allow("shared-key"): + local_allowed += 1 + with allowed_lock: + allowed.append(local_allowed) + + threads = [threading.Thread(target=hammer) for _ in range(thread_count)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert sum(allowed) == max_events + + def test_keys_are_isolated(self): + limiter = broker.RateLimiter(1, window_seconds=60) + assert limiter.allow("a") is True + assert limiter.allow("a") is False + assert limiter.allow("b") is True + + def test_stale_key_eviction(self): + limiter = broker.RateLimiter(5, window_seconds=1) + now = time.monotonic() + # A key whose newest event fell out of the window is stale. + limiter._events["stale"].append(now - 10) + # A key with a recent event must survive eviction. + limiter._events["fresh"].append(now) + # Force the next allow() call to run the eviction sweep. + limiter._last_eviction = now - limiter._EVICTION_INTERVAL_SECONDS - 1 + assert limiter.allow("active") is True + assert "stale" not in limiter._events + assert "fresh" in limiter._events + assert "active" in limiter._events + + def test_events_refill_after_window(self): + limiter = broker.RateLimiter(1, window_seconds=60) + assert limiter.allow("key") is True + assert limiter.allow("key") is False + # Age the stored event past the window; the next call must succeed. + limiter._events["key"][0] = time.monotonic() - 61 + assert limiter.allow("key") is True + + +class TestByteRateLimiter: + def test_byte_budget_enforced(self): + limiter = broker.ByteRateLimiter(100, window_seconds=60) + assert limiter.allow("key", 60) is True + assert limiter.allow("key", 60) is False + assert limiter.allow("key", 40) is True + + def test_zero_budget_disables_limiting(self): + limiter = broker.ByteRateLimiter(0, window_seconds=60) + assert limiter.allow("key", 10**9) is True + + def test_concurrent_byte_accounting(self): + max_bytes = 1000 + chunk = 10 + thread_count = 8 + attempts_per_thread = 50 # 4000 bytes attempted against a 1000 budget + limiter = broker.ByteRateLimiter(max_bytes, window_seconds=60) + allowed = [] + allowed_lock = threading.Lock() + start_barrier = threading.Barrier(thread_count) + + def hammer(): + start_barrier.wait() + local_allowed = 0 + for _ in range(attempts_per_thread): + if limiter.allow("shared-key", chunk): + local_allowed += 1 + with allowed_lock: + allowed.append(local_allowed) + + threads = [threading.Thread(target=hammer) for _ in range(thread_count)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert sum(allowed) * chunk == max_bytes + + +# --------------------------------------------------------------------------- +# Admin and static token validation +# --------------------------------------------------------------------------- + + +def _token_server( + session_tokens=frozenset(), + admin_tokens=frozenset(), + token_secret=None, +) -> types.SimpleNamespace: + config = types.SimpleNamespace( + session_tokens=frozenset(session_tokens), + admin_tokens=frozenset(admin_tokens), + token_secret=token_secret, + ) + return types.SimpleNamespace(config=config) + + +class TestAdminTokenValidation: + def test_matching_admin_token_accepted(self): + server = _token_server(admin_tokens={"admin-token-1", "admin-token-2"}) + subject = broker.BrokerServer.validate_admin_token(server, "admin-token-2") + expected_digest = hashlib.sha256(b"admin-token-2").hexdigest()[:16] + assert subject == f"admin:{expected_digest}" + + def test_wrong_admin_token_rejected(self): + server = _token_server(admin_tokens={"admin-token-1"}) + assert broker.BrokerServer.validate_admin_token(server, "admin-token-x") is None + + def test_none_token_rejected(self): + server = _token_server(admin_tokens={"admin-token-1"}) + assert broker.BrokerServer.validate_admin_token(server, None) is None + + def test_empty_admin_set_rejects_everything(self): + server = _token_server() + assert broker.BrokerServer.validate_admin_token(server, "anything") is None + + +class TestSessionTokenValidation: + def test_static_token_accepted(self): + server = _token_server(session_tokens={"static-token"}) + subject = broker.BrokerServer.validate_token(server, "static-token") + expected_digest = hashlib.sha256(b"static-token").hexdigest()[:16] + assert subject == f"static:{expected_digest}" + + def test_unknown_token_without_secret_rejected(self): + server = _token_server(session_tokens={"static-token"}) + assert broker.BrokerServer.validate_token(server, "other") is None + + def test_signed_token_fallback(self): + server = _token_server(session_tokens={"static-token"}, token_secret=SECRET) + token = broker.mint_signed_token(SECRET, "device-9", ttl_seconds=60) + assert broker.BrokerServer.validate_token(server, token) == "signed:device-9" + + def test_none_token_rejected(self): + server = _token_server(session_tokens={"static-token"}) + assert broker.BrokerServer.validate_token(server, None) is None + + +# --------------------------------------------------------------------------- +# Privacy: image counting on adversarial payloads +# --------------------------------------------------------------------------- + + +class TestCountInputImages: + def test_simple_input_image(self): + payload = {"input": [{"type": "input_image", "image_url": "data:image/png;base64,AA=="}]} + assert broker.count_input_images(payload) == 1 + + def test_data_uri_without_input_image_type_counts(self): + # Smuggling a data:image URL under a different type must still count. + payload = {"input": [{"type": "input_text", "image_url": "data:image/jpeg;base64,AA=="}]} + assert broker.count_input_images(payload) == 1 + + def test_input_image_with_data_uri_counts_once(self): + payload = { + "input": [{"type": "input_image", "image_url": "data:image/jpeg;base64,AA=="}] + } + assert broker.count_input_images(payload) == 1 + + def test_deeply_nested_images_are_found(self): + payload = { + "metadata": { + "hidden": [ + {"deeper": {"type": "input_image"}}, + [{"type": "input_image"}], + ] + }, + "input": [ + { + "role": "user", + "content": [ + {"type": "input_image", "image_url": "data:image/png;base64,AA=="} + ], + } + ], + } + assert broker.count_input_images(payload) == 3 + + def test_non_image_data_uri_not_counted(self): + payload = {"input": [{"type": "input_text", "image_url": "data:text/plain;base64,AA=="}]} + assert broker.count_input_images(payload) == 0 + + def test_non_string_image_url_not_counted(self): + payload = {"input": [{"image_url": {"url": "data:image/png;base64,AA=="}}]} + # The nested dict has no type/image_url string, so nothing counts. + assert broker.count_input_images(payload) == 0 + + def test_scalar_and_empty_payloads(self): + assert broker.count_input_images(None) == 0 + assert broker.count_input_images("data:image/png;base64,AA==") == 0 + assert broker.count_input_images(123) == 0 + assert broker.count_input_images({}) == 0 + assert broker.count_input_images([]) == 0 + + def test_many_images_all_counted(self): + images = [{"type": "input_image"} for _ in range(25)] + assert broker.count_input_images({"input": images}) == 25