undercooked/tests/integration/test_ws_abuse.py

"""The abuse suite (phase 06 task 06): the wire's guardrails, driven by
real sockets against a real `uvicorn.Server`.

The contract under test is PLAN §7 / `docs/PROTOCOL.md` "Inbound
guards": a malformed frame is an `err` + a counted shape violation (5
inside 10 s → close `4404`); a frame over `max_message_bytes` is
`too_big`, a frame over the token bucket `rate_limited` (3 consecutive
→ `4404`); the per-state whitelist is enforced before the semantics;
the host-only verbs are host-only no matter how fast they are
repeated; the emote throttle (2 per 1.0 s per connection) is a drop,
not an error. A flood of any of these must leave the **other**
connections — and the server itself — exactly as they were: the guard
is per-connection, the room and the process are not the blast radius.

Every test here bounds its own window: a live room's `snap` stream
never goes quiet while the round is `PLAYING`, so a drain waits for a
*budget* of frames or a *deadline*, never for silence that a healthy
server does not produce (the pre-task-06 drain design hung on a fixed
server and only ever passed on the pre-fix one that aborted).

Measured on this suite's machine (the `docs/PROTOCOL.md` budget table
carries the same numbers): 200 concurrent handshakes wall 0.18 s,
p95 164 ms, RSS delta 23.4 MB, `send_queue_max` 2.
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import time
from typing import Any

import httpx
import pytest
import websockets
from websockets.exceptions import ConnectionClosed

from undercooked.server.ws import (
    ABUSE_WINDOW_S,
    MAX_SEND_QUEUE,
    SHAPE_VIOLATIONS_TO_CLOSE,
)

pytestmark = pytest.mark.integration

#: The 4 KB boundary (PLAN §7's `max_message_bytes` default) in bytes.
MAX_MESSAGE_BYTES = 4096
#: The frame the cap refuses (64 KB — sixteen times the cap).
OVERSIZED_FRAME_BYTES = 64 * 1024
#: The `send_queue_max` bound for the burst test: the cap bounds
#: droppable traffic; non-droppable frames (the `pong`s a burst buys)
#: may push a queue past `MAX_SEND_QUEUE` until the pump drains it —
#: 4x the cap is the bound a 500-frame burst must stay inside
#: (measured peak: 2).
BURST_QUEUE_BOUND = 4 * MAX_SEND_QUEUE
#: The 200-connection memory bound (RSS delta, MB): the measured
#: footprint is 23.4 MB — the bound is 5x that, still 15x under a
#: per-connection "unbounded" regression would be visible at.
HANDSHAKE_MEMORY_BOUND_MB = 128.0
#: The p95 handshake bound (the task's: under 1 s). Measured 164 ms.
HANDSHAKE_P95_BOUND_S = 1.0


# ------------------------------------------------------------- the room


def _create_room(url: str, token: str) -> str:
    """One room, its creator at seat 0 as host."""
    response = httpx.post(f"{url}/api/rooms", json={"token": token}, timeout=10.0)
    assert response.status_code == 201, response.text
    return response.json()["room_code"]


def _join(url: str, code: str, token: str) -> None:
    """Seat `token` in `code`."""
    response = httpx.post(
        f"{url}/api/rooms/{code}/join", json={"token": token}, timeout=10.0
    )
    assert 200 <= response.status_code < 300, response.text


async def _connect(
    url: str, code: str, token: str, name: str | None = None
) -> websockets.ClientConnection:
    """A connected socket: the upgrade, the `hello`, the `welcome`."""
    ws = await websockets.connect(
        url.replace("http://", "ws://", 1) + f"/ws/{code}", open_timeout=10.0
    )
    hello: dict[str, Any] = {"type": "hello", "p": 1, "token": token}
    if name is not None:
        hello["name"] = name
    await ws.send(json.dumps(hello))
    welcome = json.loads(await ws.recv())
    assert welcome["type"] == "welcome", welcome
    return ws


async def _read_until_close(ws: Any) -> list[dict[str, Any]]:
    """Everything `ws` has left to say, until the close frame."""
    frames: list[dict[str, Any]] = []
    while True:
        try:
            frames.append(json.loads(await asyncio.wait_for(ws.recv(), 5.0)))
        except (
            TimeoutError,
            ConnectionClosed,
        ):
            break
    return frames


async def _alive(ws: Any, tag: str) -> None:
    """A keep-alive round trip: `ping` → `pong` (the other connection
    is still being served)."""
    await ws.send(json.dumps({"type": "ping", "t": 1}))
    deadline = asyncio.get_running_loop().time() + 5.0
    while True:
        frame = json.loads(
            await asyncio.wait_for(
                ws.recv(), max(0.1, deadline - asyncio.get_running_loop().time())
            )
        )
        if frame["type"] == "pong":
            return
        assert (
            asyncio.get_running_loop().time() < deadline
        ), f"{tag} lost its pong within 5 s: {frame}"


# ------------------------------------------------- 1: the shape storm


async def test_two_thousand_malformed_frames_cannot_take_the_server_down(
    ws_server: Any,
) -> None:
    """2000 malformed frames on one socket: the socket closes at the
    violation limit (`4404`), the *other* connection keeps being served,
    and the server is healthy afterwards."""
    with ws_server() as (url, app):
        code = _create_room(url, "tok-host")
        _join(url, code, "tok-attacker")
        host = await _connect(url, code, "tok-host", "Host")
        attacker = await _connect(url, code, "tok-attacker", "Attacker")
        room = app.state.supervisor.get(code)
        assert room is not None

        # The storm: 2000 shape violations (valid JSON, no `type` —
        # `parse_client_message` answers `bad_field` for each). The
        # close falls on the 5th inside the 10 s window; the rest of
        # the flood meets a closed socket (suppressed — the guard did
        # its job).
        for _ in range(2000):
            with contextlib.suppress(Exception):
                await attacker.send(json.dumps({"nope": 1}))

        # The attacker's socket is gone, by the abuse close.
        await _read_until_close(attacker)
        closed = attacker.close_code
        assert closed is not None, "the attacker's socket was not closed"
        assert closed == 4404, f"expected the 4404 abuse close, got {closed}"
        assert (
            SHAPE_VIOLATIONS_TO_CLOSE == 5
        )  # the bound the close fell on (the doc's "disconnect after 5")
        assert ABUSE_WINDOW_S == 10.0

        # The other connection is unaffected: a round trip on it, and
        # the room exactly as it was (LOBBY, two seats — the attacker's
        # seat is kept for the grace, not ejected by the flood).
        await _alive(host, "the host")
        assert room.state.value == "lobby"
        assert len(room.seats) == 2
        probe = await asyncio.to_thread(
            httpx.get, f"{url}/healthz", timeout=5.0
        )
        assert probe.status_code == 200


# ------------------------------------------------- 2: the rate burst


async def test_a_500_message_burst_hits_the_rate_limit_and_the_queue_bound(
    ws_server: Any,
) -> None:
    """500 pings on one socket: the bucket (40/40) empties, the
    overage is `rate_limited` (3 consecutive → `4404`), the send
    queue's high-water mark stays within its bound, and the server
    survives."""
    with ws_server() as (url, app):
        code = _create_room(url, "tok-host")
        host = await _connect(url, code, "tok-host", "Host")
        room = app.state.supervisor.get(code)
        assert room is not None

        # The burst: 500 pings, as fast as the socket lets us (the
        # first 40 take the bucket's burst; the overage is refused).
        for index in range(500):
            with contextlib.suppress(Exception):
                await host.send(json.dumps({"type": "ping", "t": index}))

        frames = await _read_until_close(host)
        codes = [
            frame.get("code")
            for frame in frames
            if frame["type"] == "err"
        ]
        assert "rate_limited" in codes, (
            "no rate_limited err frame arrived in the burst"
        )
        closed = host.close_code
        assert closed is not None, "the burst socket was not closed"
        assert closed == 4404, f"expected the 4404 abuse close, got {closed}"

        # The queue's high-water mark: the `pong`s the burst buys are
        # non-droppable, so a queue may lawfully run past the cap —
        # the bound is what must hold (measured peak: 2).
        stats = app.state.ws_registry.stats()
        assert stats["send_queue_max"] <= BURST_QUEUE_BOUND, stats

        # And the server is healthy.
        probe = await asyncio.to_thread(
            httpx.get, f"{url}/healthz", timeout=5.0
        )
        assert probe.status_code == 200
        assert room.state.value == "lobby"


# ------------------------------------------------- 3: the size cap


async def test_the_size_cap_refuses_64_kb_and_admits_the_4kb_boundary(
    ws_server: Any,
) -> None:
    """A 64 KB frame is `too_big` (refused before parsing, a rate
    violation); a frame of exactly `max_message_bytes` (4096 B — the
    cap's boundary) is admitted: the `ready` in it lands (the lobby
    re-broadcast) without a single `too_big`."""
    with ws_server() as (url, app):
        code = _create_room(url, "tok-host")
        host = await _connect(url, code, "tok-host", "Host")
        room = app.state.supervisor.get(code)
        assert room is not None
        seat = room.seat_for_token("tok-host")
        assert seat is not None and seat.ready is False

        # The 64 KB *text* frame (the wire's contract): any bytes do
        # (the cap is checked before the parse) — a `ready` padded
        # past the cap with JSON whitespace.
        oversized = b'{"type": "ready"}' + b" " * (
            OVERSIZED_FRAME_BYTES - len(b'{"type": "ready"}')
        )
        assert len(oversized) == OVERSIZED_FRAME_BYTES
        await host.send(oversized.decode("utf-8"))

        # The 64 KB *binary* frame: the same bytes, but the client
        # library encodes `bytes` as a binary frame — the regression
        # the `KeyError`'s fix covers (the cap answers, not the
        # stack's `receive_text()`).
        await host.send(oversized)

        # The 4 KB boundary: a `ready` padded with JSON whitespace to
        # exactly `max_message_bytes` (not over — the cap is `>`).
        payload = b'{"type": "ready"}'
        boundary = payload + b" " * (MAX_MESSAGE_BYTES - len(payload))
        assert len(boundary) == MAX_MESSAGE_BYTES
        await host.send(boundary.decode("utf-8"))

        # Read what the server said: two `too_big`s (the two 64 KB
        # frames), the `ready`'s lobby re-broadcast, nothing else.
        saw_too_big = 0
        saw_lobby = False
        deadline = asyncio.get_running_loop().time() + 5.0
        while asyncio.get_running_loop().time() < deadline:
            try:
                frame = json.loads(await asyncio.wait_for(host.recv(), 0.5))
            except TimeoutError:
                break
            if frame["type"] == "err":
                if frame["code"] == "too_big":
                    saw_too_big += 1
                else:
                    pytest.fail(f"an unexpected err frame: {frame}")
            elif frame["type"] == "lobby":
                saw_lobby = True

        assert saw_too_big == 2, (
            f"each 64 KB frame (text and binary) must be refused "
            f"exactly once (too_big); saw {saw_too_big}"
        )
        assert saw_lobby, "the 4 KB boundary frame did not land (no ready-up)"
        assert seat.ready is True, "the boundary frame's ready-up did not apply"


# ------------------------------------------------- 4: the lobby flood


async def test_a_lobby_flood_cannot_bypass_the_host_check_or_the_emote_throttle(
    ws_server: Any,
) -> None:
    """A non-host floods the host-only verbs (`start` x 50,
    `set_level` x 50) and the emote (x 20): the `not_host` refusals
    are per-frame errors, not shape violations (the connection is not
    closed by them), no transition happens, the level is untouched,
    and the emote throttle (2 per 1.0 s) drops the overage silently
    (no `err`, fewer `ev`s than the flood). The flood is paced to
    the token bucket (40/40 — waves of ≤ 40 frames a second apart):
    the `not_host` refusals must be tested *through* the rate guard,
    not by tripping it first."""
    with ws_server() as (url, app):
        code = _create_room(url, "tok-host")
        _join(url, code, "tok-b")
        host = await _connect(url, code, "tok-host", "Host")
        b = await _connect(url, code, "tok-b", "B")
        room = app.state.supervisor.get(code)
        assert room is not None
        original_level = room.level

        # The flood, paced to the token bucket (40/40): waves of ≤ 40
        # frames, a second apart (the refill outpaces the wave — the
        # `not_host` refusals are what the flood is meant to test).
        async def _wave(frames: list[dict[str, Any]]) -> None:
            for frame in frames:
                with contextlib.suppress(Exception):
                    await b.send(json.dumps(frame))
            await asyncio.sleep(1.1)

        await _wave([{"type": "start"} for _ in range(40)])
        await _wave(
            [{"type": "start"} for _ in range(10)]
            + [{"type": "set_level", "level": "not-a-level"} for _ in range(30)]
        )
        await _wave(
            [{"type": "set_level", "level": "not-a-level"} for _ in range(20)]
            + [{"type": "emote", "id": "wave"} for _ in range(10)]
        )
        await _wave(
            [{"type": "emote", "id": "wave"} for _ in range(10)]
        )

        # B's view: `not_host` for every start/set_level (the semantic
        # refusals — 100 of them), no close by the flood.
        not_host = 0
        emote_events_b = 0
        deadline = asyncio.get_running_loop().time() + 5.0
        while asyncio.get_running_loop().time() < deadline:
            try:
                frame = json.loads(await asyncio.wait_for(b.recv(), 0.5))
            except TimeoutError:
                break
            if frame["type"] == "err" and frame["code"] == "not_host":
                not_host += 1
            elif frame["type"] == "ev":
                for event in frame["events"]:
                    if event.get("k") == "emote":
                        emote_events_b += 1
        assert not_host >= 100, (
            f"every start/set_level needs its not_host refusal "
            f"(got {not_host})"
        )
        assert (
            b.close_code is None
        ), "the not_host refusals are semantic errors, not a 4404"

        # The room is exactly as it was: no transition, no level.
        assert room.state.value == "lobby", (
            "the start flood must not have left the lobby"
        )
        assert room.level is original_level, (
            "the set_level flood must not have changed the level"
        )

        # The emote overage is a silent drop: fewer events than the
        # flood (2 per 1.0 s window — the whole burst is one window on
        # a healthy loop; a slow one may roll into a second).
        assert emote_events_b < 20, (
            f"the emote throttle must drop the overage (got "
            f"{emote_events_b} events from a 20-emote burst)"
        )

        # And the host is unaffected.
        await _alive(host, "the host")


# ------------------------------------------------- 5: 200 handshakes


async def test_200_concurrent_handshakes_stay_within_the_memory_bound(
    ws_server: Any,
) -> None:
    """200 concurrent handshakes (50 rooms x 4 seats, every seat
    taken by a real socket): all 200 `welcome`s land, the p95
    handshake is under 1 s, the socket census sees them, and the
    process's RSS growth stays within the bound (measured: 23.4 MB
    → the bound is 5x that)."""
    with ws_server() as (url, app):
        code_of: dict[str, str] = {}
        for room_index in range(50):
            creator = f"tok-{room_index}-0"
            code = _create_room(url, creator)
            for seat_index in (1, 2, 3):
                _join(url, code, f"tok-{room_index}-{seat_index}")
            for seat_index in (0, 1, 2, 3):
                code_of[f"tok-{room_index}-{seat_index}"] = code
        assert len(code_of) == 200

        baseline_mb = app.state.ws_registry.memory_mb()

        async def handshake(token: str) -> tuple[str, float]:
            started = time.monotonic()
            ws = await websockets.connect(
                url.replace("http://", "ws://", 1) + f"/ws/{code_of[token]}",
                open_timeout=10.0,
            )
            await ws.send(
                json.dumps({"type": "hello", "p": 1, "token": token})
            )
            welcome = json.loads(await ws.recv())
            assert welcome["type"] == "welcome", welcome
            return token, time.monotonic() - started

        results = await asyncio.gather(
            *(handshake(token) for token in code_of)
        )
        assert len(results) == 200

        durations = sorted(duration for _, duration in results)
        p95 = durations[int(0.95 * len(durations)) - 1]
        assert p95 < HANDSHAKE_P95_BOUND_S, (
            f"the p95 handshake is {p95 * 1000:.0f} ms "
            f"(bound {HANDSHAKE_P95_BOUND_S * 1000:.0f} ms)"
        )

        # The census sees all 200 sockets (one per token — no
        # eviction: every token is distinct).
        stats = app.state.ws_registry.stats()
        assert stats["sockets"] == 200, stats

        # And the process's growth is bounded (the 200 connections'
        # footprint: sockets, queues, the room/seat bookkeeping).
        delta_mb = app.state.ws_registry.memory_mb() - baseline_mb
        assert delta_mb <= HANDSHAKE_MEMORY_BOUND_MB, (
            f"200 handshakes grew the process by {delta_mb:.1f} MB "
            f"(bound {HANDSHAKE_MEMORY_BOUND_MB:.0f} MB)"
        )

        # The server is healthy under the load.
        probe = await asyncio.to_thread(
            httpx.get, f"{url}/healthz", timeout=5.0
        )
        assert probe.status_code == 200