1 changed files with 282 additions and 0 deletions
@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import socket
+import threading
+import time
+from dataclasses import dataclass
+from typing import Any
+
+
+PROTOCOL_VERSION = 2
+MAX_LINE_BYTES = 64 * 1024
+DEFAULT_TIMEOUT_SECONDS = 3.0
+
+
+@dataclass
+class SessionInfo:
+    player_id: int
+    server_protocol_version: int
+    world_state_host_player_id: int | None
+
+
+class ProtocolV2Client:
+    """Small headless Commonwealth Online client for protocol and load testing."""
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        *,
+        name: str = "synthetic-client",
+        timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
+    ) -> None:
+        self.host = host
+        self.port = int(port)
+        self.name = name
+        self.timeout_seconds = float(timeout_seconds)
+        self.socket: socket.socket | None = None
+        self.session: SessionInfo | None = None
+        self._recv_buffer = b""
+        self._send_lock = threading.Lock()
+
+    def connect(self) -> SessionInfo:
+        if self.socket is not None:
+            raise RuntimeError("Client is already connected.")
+
+        connection = socket.create_connection((self.host, self.port), timeout=self.timeout_seconds)
+        connection.settimeout(self.timeout_seconds)
+        self.socket = connection
+
+        try:
+            welcome = self.recv_packet()
+            if welcome.get("type") != "welcome":
+                raise RuntimeError(f"Expected welcome, received {welcome.get('type')!r}")
+            server_protocol = _require_positive_int(welcome.get("protocolVersion"), "welcome.protocolVersion")
+            if server_protocol < PROTOCOL_VERSION:
+                raise RuntimeError(
+                    f"Server protocol {server_protocol} does not support client protocol {PROTOCOL_VERSION}."
+                )
+
+            self.send_packet(
+                {
+                    "type": "hello",
+                    "protocolVersion": PROTOCOL_VERSION,
+                    "clientVersion": "co-server-synthetic/1",
+                    "capabilities": ["synthetic-client-v1", "interest-v1"],
+                    "clientName": self.name,
+                }
+            )
+
+            while True:
+                packet = self.recv_packet()
+                packet_type = packet.get("type")
+                if packet_type == "sessionEnded":
+                    raise RuntimeError(
+                        f"Session rejected: {packet.get('code', 'unknown')} {packet.get('reason', '')}".strip()
+                    )
+                if packet_type != "sessionReady":
+                    continue
+
+                player_id = _require_positive_int(packet.get("playerId"), "sessionReady.playerId")
+                negotiated = _require_positive_int(packet.get("protocolVersion"), "sessionReady.protocolVersion")
+                server_version = _require_positive_int(
+                    packet.get("serverProtocolVersion"), "sessionReady.serverProtocolVersion"
+                )
+                host_id = packet.get("worldStateHostPlayerId")
+                if host_id is not None:
+                    host_id = _require_positive_int(host_id, "sessionReady.worldStateHostPlayerId")
+                self.session = SessionInfo(player_id, server_version, host_id)
+                if negotiated != PROTOCOL_VERSION:
+                    raise RuntimeError(f"Unexpected negotiated protocol version {negotiated}.")
+                return self.session
+        except Exception:
+            self.close()
+            raise
+
+    def close(self) -> None:
+        connection = self.socket
+        self.socket = None
+        self.session = None
+        self._recv_buffer = b""
+        if connection is not None:
+            try:
+                connection.shutdown(socket.SHUT_RDWR)
+            except OSError:
+                pass
+            try:
+                connection.close()
+            except OSError:
+                pass
+
+    def __enter__(self) -> ProtocolV2Client:
+        self.connect()
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close()
+
+    def send_packet(self, packet: dict[str, Any]) -> None:
+        connection = self.socket
+        if connection is None:
+            raise RuntimeError("Client is not connected.")
+        encoded = json.dumps(packet, separators=(",", ":"), allow_nan=False).encode("utf-8") + b"\n"
+        if len(encoded) > MAX_LINE_BYTES + 1:
+            raise ValueError("Packet exceeds protocol line limit.")
+        with self._send_lock:
+            connection.sendall(encoded)
+
+    def recv_packet(self) -> dict[str, Any]:
+        connection = self.socket
+        if connection is None:
+            raise RuntimeError("Client is not connected.")
+
+        while b"\n" not in self._recv_buffer:
+            chunk = connection.recv(4096)
+            if not chunk:
+                raise ConnectionError("Server closed the connection.")
+            self._recv_buffer += chunk
+            if len(self._recv_buffer) > MAX_LINE_BYTES and b"\n" not in self._recv_buffer:
+                raise ValueError("Server response exceeded protocol line limit.")
+
+        line, self._recv_buffer = self._recv_buffer.split(b"\n", 1)
+        if len(line) > MAX_LINE_BYTES:
+            raise ValueError("Server response exceeded protocol line limit.")
+        packet = json.loads(line.decode("utf-8"), parse_constant=_reject_non_finite)
+        if not isinstance(packet, dict):
+            raise ValueError("Server response must be a JSON object.")
+        return packet
+
+    def send_transform(
+        self,
+        *,
+        x: float,
+        y: float,
+        z: float,
+        angle_z: float,
+        cell_id: int,
+        worldspace_id: int = 0,
+        movement_speed: float = 0.0,
+        animation_direction: float = 0.0,
+        is_moving: bool = False,
+    ) -> None:
+        if self.session is None:
+            raise RuntimeError("Protocol session is not ready.")
+        if not all(math.isfinite(value) for value in (x, y, z, angle_z, movement_speed, animation_direction)):
+            raise ValueError("Transform values must be finite.")
+        packet: dict[str, Any] = {
+            "type": "transform",
+            "playerId": self.session.player_id,
+            "x": float(x),
+            "y": float(y),
+            "z": float(z),
+            "angleZ": float(angle_z),
+            "cellId": f"{int(cell_id) & 0xFFFFFFFF:08X}",
+            "isMoving": bool(is_moving),
+            "movementSpeed": float(movement_speed),
+            "animationDirection": float(animation_direction),
+        }
+        if worldspace_id:
+            packet["worldspaceId"] = f"{int(worldspace_id) & 0xFFFFFFFF:08X}"
+        self.send_packet(packet)
+
+
+def run_load_scenario(
+    host: str,
+    port: int,
+    *,
+    client_count: int,
+    cell_count: int,
+    duration_seconds: float,
+    send_rate_hz: float,
+) -> None:
+    if client_count < 1:
+        raise ValueError("client_count must be at least 1")
+    if cell_count < 1:
+        raise ValueError("cell_count must be at least 1")
+    if duration_seconds <= 0:
+        raise ValueError("duration_seconds must be positive")
+    if send_rate_hz <= 0:
+        raise ValueError("send_rate_hz must be positive")
+
+    clients: list[ProtocolV2Client] = []
+    try:
+        for index in range(client_count):
+            client = ProtocolV2Client(host, port, name=f"load-{index + 1}")
+            session = client.connect()
+            clients.append(client)
+            print(f"connected {client.name}: player={session.player_id} protocol={session.server_protocol_version}")
+
+        interval = 1.0 / send_rate_hz
+        deadline = time.monotonic() + duration_seconds
+        tick = 0
+        while time.monotonic() < deadline:
+            started = time.monotonic()
+            for index, client in enumerate(clients):
+                group = index % cell_count
+                cell_id = 0x01000000 + group + 1
+                x = float((group * 20000) + (index * 64) + tick)
+                y = float(index * 32)
+                client.send_transform(
+                    x=x,
+                    y=y,
+                    z=0.0,
+                    angle_z=0.0,
+                    cell_id=cell_id,
+                    movement_speed=send_rate_hz,
+                    animation_direction=0.0,
+                    is_moving=True,
+                )
+            tick += 1
+            remaining = interval - (time.monotonic() - started)
+            if remaining > 0:
+                time.sleep(remaining)
+
+        print(
+            f"completed: clients={client_count} cells={cell_count} duration={duration_seconds:.1f}s "
+            f"rate={send_rate_hz:.1f}Hz"
+        )
+    finally:
+        for client in clients:
+            client.close()
+
+
+def _require_positive_int(value: Any, field: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(f"{field} must be a positive integer.")
+    return value
+
+
+def _reject_non_finite(value: str) -> None:
+    raise ValueError(f"Non-finite JSON number is not allowed: {value}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Commonwealth Online Protocol V2 synthetic load client")
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=7777)
+    parser.add_argument("--clients", type=int, default=8)
+    parser.add_argument("--cells", type=int, default=2)
+    parser.add_argument("--duration", type=float, default=10.0)
+    parser.add_argument("--rate", type=float, default=10.0)
+    args = parser.parse_args()
+
+    try:
+        run_load_scenario(
+            args.host,
+            args.port,
+            client_count=args.clients,
+            cell_count=args.cells,
+            duration_seconds=args.duration,
+            send_rate_hz=args.rate,
+        )
+    except (OSError, RuntimeError, ValueError, ConnectionError) as error:
+        print(f"load test failed: {error}")
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())