diff --git a/coderag/surfaces/webui.py b/coderag/surfaces/webui.py index f0d6c2d..237075f 100644 --- a/coderag/surfaces/webui.py +++ b/coderag/surfaces/webui.py @@ -220,19 +220,16 @@ def create_ui_app(cr: "CodeRAG") -> "FastAPI": demo_sessions: Dict[str, Dict[str, Any]] = {} demo_lock = threading.Lock() - def demo_session(request: Request) -> "tuple[str, Optional[str]]": - """Return ``(sid, fresh_cookie)`` for the demo quota. + def mint_demo_cookie(request: Request) -> Optional[str]: + """A fresh session token to Set-Cookie for first-time visitors, else None. - ``fresh_cookie`` is a newly-minted token to Set-Cookie for first-time visitors, - or None for browsers that already carry one. The client's cookie value is used - only as the in-memory quota key and is never written back into a Set-Cookie - header — so no user-supplied input is ever reflected into a cookie. + Only checks whether the cookie is *present* — it never reads the client-supplied + value — so nothing user-controlled can flow into the Set-Cookie header. When a + cookie is written, its value is always a freshly minted secret. """ - existing = request.cookies.get(_DEMO_COOKIE) - if existing: - return existing, None - minted = secrets.token_urlsafe(16) - return minted, minted + if _DEMO_COOKIE in request.cookies: + return None + return secrets.token_urlsafe(16) def demo_remaining(sid: str) -> int: with demo_lock: @@ -359,8 +356,7 @@ def home( ctx["answer_qs"] = urlencode({"q": q.strip(), "k": k}) resp = templates.TemplateResponse(request, "index.html", ctx) if demo: - _, fresh_cookie = demo_session(request) - apply_demo_cookie(resp, fresh_cookie) + apply_demo_cookie(resp, mint_demo_cookie(request)) return resp @app.get("/file", response_class=HTMLResponse) @@ -415,7 +411,10 @@ def answer( fresh_cookie: Optional[str] = None if demo: - sid, fresh_cookie = demo_session(request) + fresh_cookie = mint_demo_cookie(request) + # New visitor → charge under the freshly minted id; returning visitor → + # under their existing cookie (read only as the quota key, never re-set). + sid = fresh_cookie or request.cookies.get(_DEMO_COOKIE) or "" ok, msg = demo_gate(sid, q) if not ok: blocked = Response(msg, media_type="text/plain; charset=utf-8") diff --git a/tests/test_webui.py b/tests/test_webui.py index fa897c1..6309bf6 100644 --- a/tests/test_webui.py +++ b/tests/test_webui.py @@ -166,3 +166,16 @@ def test_demo_answer_cooldown_blocks_rapid_followups(tmp_path): cr, client = _demo_client(tmp_path, max_answers=5, cooldown=30) client.get("/answer", params={"q": "a"}) # charges, starts the cooldown assert "wait" in client.get("/answer", params={"q": "b"}).text.lower() + + +def test_demo_cookie_is_minted_not_reflected(tmp_path): + # First visit (no cookie) → the server mints and sets one. + cr, client = _demo_client(tmp_path) + client.cookies.clear() + r = client.get("/") + assert "coderag_demo=" in r.headers.get("set-cookie", "") + # A request that already carries a (crafted) cookie must NOT have that value + # echoed back via Set-Cookie — the server never reflects client input. + client.cookies.clear() + r2 = client.get("/", cookies={"coderag_demo": "attacker-controlled"}) + assert "set-cookie" not in {k.lower() for k in r2.headers}