From 8ee0a0d082737c2f8f1e7447695ac3660ca9eb02 Mon Sep 17 00:00:00 2001 From: Rohit P Date: Fri, 24 Jul 2026 17:47:22 -0700 Subject: [PATCH] security: strengthen session handling --- coworker/server/app.py | 172 ++++++++++++++++++++++++++++++------- coworker/server/manager.py | 16 ++-- coworker/server/run.py | 6 +- surfaces/gui/src/App.tsx | 6 ++ surfaces/gui/src/types.ts | 1 + tests/test_server.py | 121 +++++++++++++++++++++++++- 6 files changed, 277 insertions(+), 45 deletions(-) diff --git a/coworker/server/app.py b/coworker/server/app.py index e43033a2..67d77024 100644 --- a/coworker/server/app.py +++ b/coworker/server/app.py @@ -12,6 +12,7 @@ import json import os import re import uuid +from collections import deque from contextlib import asynccontextmanager from pathlib import Path from typing import Any, Optional @@ -39,25 +40,25 @@ def _origin_allowed(origin: str | None) -> bool: return origin is None or bool(_ALLOWED_ORIGIN_RE.match(origin)) -# Caps on a single inbound `user_message` frame. The loopback socket is unauthenticated -# (any local process can reach it), so an oversized frame is a cheap way to spike memory — -# these bound one message's text and its attachments before we build content / start a turn. +# Caps on inbound WebSocket traffic. The loopback socket is unauthenticated (any local +# process can reach it), so bound frames, messages, and per-connection request rate before +# building model content or starting a turn. +_WS_MAX_FRAME_BYTES = 16 * 1024 * 1024 +_WS_RATE_LIMIT_COUNT = 30 +_WS_RATE_LIMIT_WINDOW_SECONDS = 10.0 _MAX_MESSAGE_TEXT_CHARS = 200_000 -_MAX_ATTACHMENTS = 25 -_MAX_ATTACHMENTS_BYTES = 32 * 1024 * 1024 # ~32 MB total across a message's attachments +_MAX_ATTACHMENTS_BYTES = 15_000_000 # leaves JSON overhead below the 16 MiB frame cap -def _attachments_size(attachments: list) -> int: - """Approximate on-wire byte size of a message's attachments (data URLs dominate).""" - total = 0 - for a in attachments: - if isinstance(a, str): - total += len(a) - elif isinstance(a, dict): - for v in a.values(): - if isinstance(v, str): - total += len(v) - return total +def _json_value_size(value: Any) -> int: + """Conservative UTF-8 size of parsed JSON without allocating another giant string.""" + if isinstance(value, str): + return len(value.encode("utf-8")) + if isinstance(value, dict): + return sum(_json_value_size(k) + _json_value_size(v) for k, v in value.items()) + if isinstance(value, list): + return sum(_json_value_size(v) for v in value) + return 8 # numbers, booleans, null, separators # Brand colors for the connector badge riding the ✓ (UX-DECISIONS §30). The GUI owns the @@ -145,7 +146,13 @@ _CONNECT_FAILED_DETAIL = ( "Close this tab and try again from OpenWorker." ) -from ..attachments import build_user_content +from ..attachments import ( + MAX_ATTACHMENTS as _MAX_ATTACHMENTS, + MAX_IMAGE_CHARS, + MAX_PDF_CHARS, + MAX_TEXT_CHARS, + build_user_content, +) from ..engine import ApprovalOutcome from ..inbox import VIS_INBOX, VIS_INLINE, args_preview from ..permissions import Mode @@ -1618,9 +1625,8 @@ def create_app(manager: SessionManager) -> FastAPI: } async def run_turn(content, *, retry: bool = False) -> None: - manager.mark_running( - session_id - ) # busy → self-wakes steer instead of colliding + # The receive loop atomically claims this session before scheduling the task. + # Keeping the claim outside prevents two back-to-back frames from both starting. try: events = engine.retry() if retry else engine.run(content) async for event in events: @@ -1641,10 +1647,48 @@ def create_app(manager: SessionManager) -> FastAPI: # This socket is now a live view of the session; background turns (channel delivery, # self-wake, durable resume) broadcast here too, not just locally driven run_turns. manager.register_session_client(session_id, ws.send_json) + inbound_times: deque[float] = deque() + + async def reject_input(reason: str) -> None: + # Input validation failures are not provider failures and must not offer "Retry" + # or flush an in-progress assistant stream in the GUI. + await ws.send_json({"type": "input_rejected", "data": {"error": reason}}) + + async def claim_turn(*, retry: bool = False, content=None) -> None: + if not manager.try_mark_running(session_id): + await reject_input( + "This session is already running a turn. Wait for it to finish or stop it." + ) + return + asyncio.create_task(run_turn(content, retry=retry)) + try: while True: - message = await ws.receive_json() + try: + message = await ws.receive_json() + except (json.JSONDecodeError, UnicodeDecodeError): + await reject_input("Invalid WebSocket message: expected JSON.") + continue + + now = asyncio.get_running_loop().time() + while ( + inbound_times + and now - inbound_times[0] > _WS_RATE_LIMIT_WINDOW_SECONDS + ): + inbound_times.popleft() + if len(inbound_times) >= _WS_RATE_LIMIT_COUNT: + await reject_input("Too many WebSocket messages; reconnect and try again.") + await ws.close(code=1008) + return + inbound_times.append(now) + + if not isinstance(message, dict): + await reject_input("Invalid WebSocket message: expected an object.") + continue kind = message.get("type") + if not isinstance(kind, str): + await reject_input("Invalid WebSocket message: missing string type.") + continue if kind == "approval": _resolve_pending(message.get("decision", "deny")) elif kind == "directory_response": @@ -1674,22 +1718,33 @@ def create_app(manager: SessionManager) -> FastAPI: elif kind == "retry": # Re-run after a provider error (engine guards on the error-notice # tail, so a stray frame is a no-op that still ends with turn_done). - if not manager.is_running(session_id): - asyncio.create_task(run_turn(None, retry=True)) + await claim_turn(retry=True) elif kind == "set_mode": try: engine.permissions.mode = Mode(message.get("mode")) - except ValueError: + except (TypeError, ValueError): pass elif kind == "set_model": - await _apply_model(message.get("model")) + model = message.get("model") + if model is not None and not isinstance(model, str): + await reject_input("Invalid model: expected a string.") + else: + await _apply_model(model) elif kind == "user_message": - text = (message.get("text") or "").strip() - attachments = message.get("attachments") or [] + raw_text = message.get("text") + if raw_text is None: + raw_text = "" + if not isinstance(raw_text, str): + await reject_input("Invalid message text: expected a string.") + continue + text = raw_text.strip() + raw_attachments = message.get("attachments") + attachments = [] if raw_attachments is None else raw_attachments # Reject an oversized frame instead of buffering it into a turn. Send a # visible error so the surface can tell the user, and drop the message. if not isinstance(attachments, list): - attachments = [] + await reject_input("Invalid attachments: expected a list.") + continue reject = None if len(text) > _MAX_MESSAGE_TEXT_CHARS: reject = ( @@ -1701,18 +1756,69 @@ def create_app(manager: SessionManager) -> FastAPI: f"Too many attachments ({len(attachments)}; " f"limit {_MAX_ATTACHMENTS})." ) - elif _attachments_size(attachments) > _MAX_ATTACHMENTS_BYTES: - reject = "Attachments too large (limit 32 MB per message)." + elif any(not isinstance(a, dict) for a in attachments): + reject = "Invalid attachment: expected an object." + elif _json_value_size(attachments) > _MAX_ATTACHMENTS_BYTES: + reject = "Attachments too large (limit 15 MB per message)." + else: + for attachment in attachments: + attachment_kind = attachment.get("kind") + name = attachment.get("name") + mime = attachment.get("mime") + if attachment_kind not in {"image", "pdf", "text"}: + reject = "Invalid attachment kind." + elif name is not None and ( + not isinstance(name, str) or len(name) > 1024 + ): + reject = "Invalid attachment name." + elif mime is not None and ( + not isinstance(mime, str) or len(mime) > 255 + ): + reject = "Invalid attachment MIME type." + elif attachment_kind == "image": + data = attachment.get("data_url") + if ( + not isinstance(data, str) + or not data.startswith("data:image/") + or ";base64," not in data + or len(data) > MAX_IMAGE_CHARS + ): + reject = "Invalid or oversized image attachment." + elif attachment_kind == "pdf": + data = attachment.get("data_url") + if ( + not isinstance(data, str) + or not data.startswith( + "data:application/pdf;base64," + ) + or len(data) > MAX_PDF_CHARS + ): + reject = "Invalid or oversized PDF attachment." + else: + body = attachment.get("text") + if ( + not isinstance(body, str) + or len(body) > MAX_TEXT_CHARS + ): + reject = "Invalid or oversized text attachment." + if reject is not None: + break if reject is not None: - await ws.send_json({"type": "error", "data": {"error": reject}}) + await reject_input(reject) continue # The composer sends its visible model with every message — the FIRST # one binds the session (race-proof across reconnects; see api.ts # Session.userMessage), later ones may switch it (notice persisted). - await _apply_model(message.get("model")) + model = message.get("model") + if model is not None and not isinstance(model, str): + await reject_input("Invalid model: expected a string.") + continue + await _apply_model(model) if text or attachments: content = build_user_content(text, attachments) - asyncio.create_task(run_turn(content)) + await claim_turn(content=content) + else: + await reject_input(f"Unknown WebSocket message type: {kind}.") except WebSocketDisconnect: pass finally: diff --git a/coworker/server/manager.py b/coworker/server/manager.py index 2501e9b9..1ff2ecbf 100644 --- a/coworker/server/manager.py +++ b/coworker/server/manager.py @@ -2465,6 +2465,13 @@ class SessionManager: def mark_running(self, session_id: str) -> None: self._running_sessions.add(session_id) + def try_mark_running(self, session_id: str) -> bool: + """Atomically claim an idle session for one turn on the server event loop.""" + if session_id in self._running_sessions: + return False + self._running_sessions.add(session_id) + return True + def mark_idle(self, session_id: str) -> None: self._running_sessions.discard(session_id) # Every turn path (WS, background delivery, durable resume) marks idle when it @@ -2488,15 +2495,12 @@ class SessionManager: by self-wake and channel-subscription delivery. `source` is the display-only MessageSource sidecar for connector messages (framed `message` stays the model-facing text). """ - if self.is_running(session_id): - engine = self._engines.get(session_id) - if engine is not None: - engine.queue_steering(message, source) - return engine = self.get_engine(session_id) if engine is None: return - self.mark_running(session_id) + if not self.try_mark_running(session_id): + engine.queue_steering(message, source) + return try: async for event in engine.run(message, source=source): # Stream every event to any socket viewing this session, so a background turn diff --git a/coworker/server/run.py b/coworker/server/run.py index e5d2ff7e..7c10799d 100644 --- a/coworker/server/run.py +++ b/coworker/server/run.py @@ -10,7 +10,7 @@ from pathlib import Path from ..config import load_config from ..permissions import Mode from ..secrets import state_dir -from .app import create_app +from .app import _WS_MAX_FRAME_BYTES, create_app from .manager import SessionManager @@ -149,7 +149,9 @@ def main(argv=None) -> None: _exit_when_orphaned() app = build_app(args.cwd, args.model, args.mode) - uvicorn.run(app, host=args.host, port=args.port) + uvicorn.run( + app, host=args.host, port=args.port, ws_max_size=_WS_MAX_FRAME_BYTES + ) if __name__ == "__main__": diff --git a/surfaces/gui/src/App.tsx b/surfaces/gui/src/App.tsx index 18204164..a2a0a124 100644 --- a/surfaces/gui/src/App.tsx +++ b/surfaces/gui/src/App.tsx @@ -695,6 +695,12 @@ export function App() { { kind: "notice", tone: "warn", text: "Error: " + (d.error || "unknown"), retriable: true }, ]); break; + case "input_rejected": + setItems((p) => [ + ...p, + { kind: "notice", tone: "warn", text: d.error || "That message was rejected." }, + ]); + break; case "turn_done": setRunning(false); refreshSessions(); diff --git a/surfaces/gui/src/types.ts b/surfaces/gui/src/types.ts index 74692a05..e0104802 100644 --- a/surfaces/gui/src/types.ts +++ b/surfaces/gui/src/types.ts @@ -15,6 +15,7 @@ export type EventType = | "iteration_end" | "turn_end" | "error" + | "input_rejected" | "interrupted" | "model_changed" | "turn_done"; diff --git a/tests/test_server.py b/tests/test_server.py index 71f4ab7e..9a85caea 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -325,19 +325,22 @@ def test_ws_simple_turn(tmp_path): def test_ws_rejects_oversized_message(tmp_path): from coworker.server import app as app_mod + from coworker.attachments import MAX_ATTACHMENTS client = _client(tmp_path, [_text("should not run")]) with client.websocket_connect("/ws/session/big") as ws: assert ws.receive_json()["type"] == "ready" - # Oversized text → single error frame, no turn runs. + # Oversized text → single input-rejected frame, no turn runs. ws.send_json( {"type": "user_message", "text": "x" * (app_mod._MAX_MESSAGE_TEXT_CHARS + 1)} ) evt = ws.receive_json() - assert evt["type"] == "error" and "too long" in evt["data"]["error"].lower() + assert evt["type"] == "input_rejected" + assert "too long" in evt["data"]["error"].lower() - # Too many attachments → error frame. + # The ingress cap is the same cap the attachment builder enforces. + assert app_mod._MAX_ATTACHMENTS == MAX_ATTACHMENTS ws.send_json( { "type": "user_message", @@ -346,13 +349,123 @@ def test_ws_rejects_oversized_message(tmp_path): } ) evt = ws.receive_json() - assert evt["type"] == "error" and "attachment" in evt["data"]["error"].lower() + assert evt["type"] == "input_rejected" + assert "attachment" in evt["data"]["error"].lower() # A normal message still works afterwards (the socket wasn't torn down). ws.send_json({"type": "user_message", "text": "hello"}) assert "turn_done" in _drain(ws) +def test_ws_rejects_malformed_payloads_without_killing_socket(tmp_path): + client = _client(tmp_path, [_text("normal")]) + with client.websocket_connect("/ws/session/malformed") as ws: + assert ws.receive_json()["type"] == "ready" + + invalid = [ + [], + {"type": "user_message", "text": ["not", "text"]}, + {"type": "user_message", "text": "x", "attachments": {}}, + { + "type": "user_message", + "text": "x", + "attachments": [{"kind": "image", "data_url": "https://example.com/x"}], + }, + {"type": "set_model", "model": {"unexpected": True}}, + {"type": "unknown"}, + ] + for payload in invalid: + ws.send_json(payload) + evt = ws.receive_json() + assert evt["type"] == "input_rejected" + + ws.send_json({"type": "user_message", "text": "still works"}) + assert "turn_done" in _drain(ws) + + +def test_ws_allows_only_one_inflight_turn_per_session(tmp_path): + import threading + import time + + class SlowProvider(ProviderClient): + def __init__(self): + self._lock = threading.Lock() + self.active = 0 + self.max_active = 0 + + def complete(self, *, model, messages, tools=None, **settings): + with self._lock: + self.active += 1 + self.max_active = max(self.max_active, self.active) + try: + time.sleep(0.08) + return _text("done") + finally: + with self._lock: + self.active -= 1 + + def capabilities(self, model): + return ModelCapabilities() + + provider = SlowProvider() + manager = SessionManager(workspace=tmp_path, provider=provider) + client = TestClient(create_app(manager)) + with client.websocket_connect("/ws/session/serialized") as ws: + assert ws.receive_json()["type"] == "ready" + ws.send_json({"type": "user_message", "text": "first"}) + ws.send_json({"type": "user_message", "text": "second"}) + + types = [] + while "turn_done" not in types: + types.append(ws.receive_json()["type"]) + + assert "input_rejected" in types + assert provider.max_active == 1 + engine = manager._engines["serialized"] + user_messages = [m for m in engine.messages if m.get("role") == "user"] + assert [m["content"] for m in user_messages] == ["first"] + + +def test_ws_rate_limits_inbound_frames(tmp_path): + from coworker.server import app as app_mod + from starlette.websockets import WebSocketDisconnect + + client = _client(tmp_path, []) + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/ws/session/rate") as ws: + assert ws.receive_json()["type"] == "ready" + for _ in range(app_mod._WS_RATE_LIMIT_COUNT): + ws.send_json({"type": "unknown"}) + assert ws.receive_json()["type"] == "input_rejected" + ws.send_json({"type": "unknown"}) + assert ws.receive_json()["type"] == "input_rejected" + ws.receive_json() + + +def test_server_sets_explicit_websocket_frame_limit(tmp_path, monkeypatch): + import sys + from types import SimpleNamespace + + from coworker.server import run as server_run + + seen = {} + fake_app = object() + + monkeypatch.setattr(server_run, "_ensure_ca_bundle", lambda: None) + monkeypatch.setattr(server_run, "_exit_when_orphaned", lambda: None) + monkeypatch.setattr(server_run, "build_app", lambda *args: fake_app) + monkeypatch.setitem( + sys.modules, + "uvicorn", + SimpleNamespace(run=lambda app, **kwargs: seen.update(app=app, **kwargs)), + ) + + server_run.main(["--cwd", str(tmp_path), "--port", "8766"]) + + assert seen["app"] is fake_app + assert seen["ws_max_size"] == server_run._WS_MAX_FRAME_BYTES + + def test_ws_error_persists_notice_and_retry_reruns(tmp_path): class FlakyProvider(ProviderClient): def __init__(self):