python-socketio

FROM python:3.13-slim
WORKDIR /srv
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY store.py server.py announce.py ./
EXPOSE 8000
CMD ["uvicorn", "server:asgi", "--host", "0.0.0.0", "--port", "8000"]

services:
  redis:
    image: redis:8.8.3-alpine
    ports:
      - "127.0.0.1:6330:6379"
  chatterbox-1:
    build: .
    environment: &env
      REDIS_URL: redis://redis:6379/0
      ALLOWED_ORIGINS: https://chat.example.test
      PRESENCE_GRACE_SECONDS: "1"   # the protocol says 15; shortened so the demo finishes quickly
    ports:
      - "127.0.0.1:8130:8000"
    depends_on: [redis]
  chatterbox-2:
    build: .
    environment: *env
    ports:
      - "127.0.0.1:8131:8000"
    depends_on: [redis]
  proxy:
    image: nginx:1.28-alpine
    volumes:
      - ./nginx.conf:/etc/nginx/conf.d/default.conf:ro
    ports:
      - "127.0.0.1:8132:80"
    depends_on: [chatterbox-1, chatterbox-2]

"""Redis-backed state for the Chatterbox server: tickets, rooms, seq, history, idempotency, presence."""
import json
import os
import re
import secrets
import time

import redis.asyncio as aioredis

REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
r = aioredis.from_url(REDIS_URL, decode_responses=True)

CROCKFORD = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
ROOM_RE = re.compile(r"^[a-z0-9-]{1,64}$")
USER_RE = re.compile(r"^[a-z0-9]{1,32}$")
HISTORY_PAGE_BYTES = 262144
EVENT = "chatterbox"  # every envelope travels as the single argument of this Socket.IO event


def ulid() -> str:
    """26-character ULID: 48-bit millisecond timestamp followed by 80 random bits."""
    n = (int(time.time() * 1000) << 80) | secrets.randbits(80)
    return "".join(CROCKFORD[(n >> (5 * i)) & 31] for i in range(25, -1, -1))


def server_id() -> str:
    return "s_" + ulid()


def now_ms() -> int:
    return int(time.time() * 1000)


def dumps(value) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def utf16_len(text: str) -> int:
    return len(text.encode("utf-16-le")) // 2  # same unit as JavaScript String.length


class ChatError(Exception):
    CODES = {"bad_request": False, "unsupported_version": False, "unauthorized": False,
             "forbidden": False, "not_found": False, "conflict": False, "too_large": False,
             "rate_limited": True, "internal": True, "unavailable": True}

    def __init__(self, code: str, message: str):
        super().__init__(message)
        self.code, self.message, self.retryable = code, message, self.CODES[code]


def envelope(type_, room, payload, id_=None):
    return {"v": 1, "type": type_, "id": id_ or server_id(), "room": room, "ts": now_ms(), "payload": payload}


def display_name(user_id: str) -> str:
    return user_id.capitalize()


# --- tickets (one-time, 30 s) ---------------------------------------------------------------
async def issue_ticket(user_id: str) -> str:
    if not USER_RE.match(user_id):
        raise ChatError("bad_request", "user must match [a-z0-9]{1,32}")
    ticket = "tk_" + secrets.token_urlsafe(24)
    await r.set(f"ticket:{ticket}", user_id, ex=30)
    return ticket


async def redeem_ticket(ticket: str):
    return await r.getdel(f"ticket:{ticket}")  # atomic read-and-delete: a replay finds nothing


# --- idempotency ----------------------------------------------------------------------------
async def cached_ack(user_id: str, env_id: str):
    raw = await r.get(f"idem:{user_id}:{env_id}")
    return json.loads(raw) if raw else None


async def remember_ack(user_id: str, env_id: str, ack_payload: dict):
    await r.set(f"idem:{user_id}:{env_id}", dumps(ack_payload), ex=86400)


# --- rooms and membership -------------------------------------------------------------------
def check_room_access(room: str, user_id: str):
    if not ROOM_RE.match(room):
        raise ChatError("bad_request", "room must match [a-z0-9-]{1,64}")
    if room.startswith("dm-"):
        parts = room.split("-")
        if len(parts) != 3 or parts[1] > parts[2] or user_id not in parts[1:]:
            raise ChatError("forbidden", "direct-message room")


async def add_member(room: str, user_id: str) -> bool:
    added = await r.sadd(f"room:{room}:members", user_id)
    await r.sadd(f"user:{user_id}:rooms", room)
    return bool(added)


async def remove_member(room: str, user_id: str) -> bool:
    removed = await r.srem(f"room:{room}:members", user_id)
    await r.srem(f"user:{user_id}:rooms", room)
    return bool(removed)


async def is_member(room: str, user_id: str) -> bool:
    return bool(await r.sismember(f"room:{room}:members", user_id))


async def user_rooms(user_id: str) -> list[str]:
    return sorted(await r.smembers(f"user:{user_id}:rooms"))


async def members(room: str) -> list[dict]:
    out = []
    for uid in sorted(await r.smembers(f"room:{room}:members")):
        out.append({"userId": uid, "name": display_name(uid),
                    "status": await r.hget("presence", uid) or "offline"})
    return out


async def last_seq(room: str) -> int:
    return int(await r.get(f"room:{room}:seq") or 0)


# --- messages -------------------------------------------------------------------------------
async def append_message(room, user_id, msg_id, text, reply_to=None, attachments=None):
    """Assign the next seq, store the message and return the broadcast envelope."""
    if reply_to is not None and not await r.hexists(f"room:{room}:ids", reply_to):
        raise ChatError("not_found", "replyTo is not a message of this room")
    seq = await r.incr(f"room:{room}:seq")
    if not await r.hsetnx(f"room:{room}:ids", msg_id, seq):
        raise ChatError("conflict", "message id already used")
    env = envelope("message", room, {
        "from": {"userId": user_id, "name": display_name(user_id)}, "text": text,
        "replyTo": reply_to, "attachments": attachments or [], "seq": seq, "edited": False}, id_=msg_id)
    member = dumps(env)
    await r.zadd(f"room:{room}:msgs", {member: seq})
    await r.zadd(f"room:{room}:log", {member: seq})
    return env


async def history(room, before, limit):
    upper = "+inf" if before is None else f"({before}"
    raws = await r.zrevrangebyscore(f"room:{room}:msgs", upper, "-inf", start=0, num=limit + 1)
    page, size = [], 0
    for raw in raws[:limit]:
        size += len(raw.encode())
        if page and size > HISTORY_PAGE_BYTES:  # keep a page under 256 KiB, return fewer messages
            break
        page.append(json.loads(raw))
    has_more = len(page) < len(raws)
    page.reverse()
    return {"messages": page, "hasMore": has_more, "next": page[0]["payload"]["seq"] if has_more else None}


async def log_since(room, since):
    return [json.loads(x) for x in await r.zrangebyscore(f"room:{room}:log", f"({since}", "+inf")]


# --- sessions and presence --------------------------------------------------------------------
async def save_session(session_id, user_id, ttl=None):
    await r.set(f"sess:{session_id}", user_id, ex=ttl or 120)


async def session_owner(session_id):
    return await r.get(f"sess:{session_id}")


async def set_status(user_id, status) -> bool:
    """Returns True when the status changed (and so must be broadcast)."""
    previous = await r.hget("presence", user_id)
    await r.hset("presence", user_id, status)
    return previous != status

"""Chatterbox protocol v1 over Socket.IO: one event name, one envelope per event."""
import logging
import os
import time
from dataclasses import dataclass, field

import socketio
from redis.exceptions import RedisError
from fastapi import FastAPI, HTTPException

import store
from store import EVENT, ChatError, envelope

log = logging.getLogger("chatterbox")
CLOSE_EVENT = "chatterbox.close"  # Socket.IO has no close codes, so the code travels in this event
HEARTBEAT_MS = 20000
MAX_FRAME = 16384
GRACE = float(os.environ.get("PRESENCE_GRACE_SECONDS", "15"))
ORIGINS = os.environ.get("ALLOWED_ORIGINS", "https://chat.example.test").split(",")

mgr = socketio.AsyncRedisManager(store.REDIS_URL)
sio = socketio.AsyncServer(
    async_mode="asgi",
    client_manager=mgr,
    cors_allowed_origins=ORIGINS,
    transports=["websocket"],           # required with several workers or instances
    ping_interval=HEARTBEAT_MS / 1000,
    ping_timeout=HEARTBEAT_MS / 1000,   # silent for 2 x heartbeatMs -> disconnected
    max_http_buffer_size=MAX_FRAME,
    async_handlers=False,               # one client's envelopes are handled strictly in arrival order
)


@dataclass
class Conn:
    user_id: str | None = None
    session_id: str | None = None
    ready: bool = False
    tokens: float = 40.0
    stamp: float = field(default_factory=time.monotonic)
    limited: list = field(default_factory=list)


CONNS: dict[str, Conn] = {}


# --- sending ------------------------------------------------------------------------------------
async def send(sid, env):
    await sio.emit(EVENT, env, to=sid, ignore_queue=True)  # sid lives on this process: skip Redis


async def send_error(sid, ref, err: ChatError, room=None):
    await send(sid, envelope("error", room, {"ref": ref, "code": err.code, "message": err.message,
                                             "retryable": err.retryable}))


async def close(sid, code, reason):
    await sio.emit(CLOSE_EVENT, {"code": code, "reason": reason}, to=sid, ignore_queue=True)
    await sio.disconnect(sid)


async def to_room(room, env, skip_sid=None):
    await sio.emit(EVENT, env, room=room, skip_sid=skip_sid)  # goes through Redis to every instance


async def presence_env(room, user_id, status):
    return envelope("presence", room, {"userId": user_id, "name": store.display_name(user_id), "status": status})


async def broadcast_presence(user_id, status, skip_sid=None):
    for room in await store.user_rooms(user_id):
        await to_room(room, await presence_env(room, user_id, status), skip_sid)


# --- connection lifecycle -----------------------------------------------------------------------
@sio.event
async def connect(sid, environ, auth):
    CONNS[sid] = Conn()
    sio.start_background_task(hello_deadline, sid)


async def hello_deadline(sid):
    await sio.sleep(5)
    conn = CONNS.get(sid)
    if conn and not conn.ready:
        await send_error(sid, None, ChatError("unauthorized", "hello expected within 5 s"))
        await close(sid, 4408, "hello timeout")


@sio.event
async def disconnect(sid, reason):
    conn = CONNS.pop(sid, None)
    if not conn or not conn.ready:
        return
    await store.save_session(conn.session_id, conn.user_id)  # keeps the session for the 120 s resume window
    if await store.r.decr(f"conns:{conn.user_id}") <= 0:
        sio.start_background_task(go_offline_later, conn.user_id)


async def go_offline_later(user_id):
    await sio.sleep(GRACE)
    if int(await store.r.get(f"conns:{user_id}") or 0) <= 0 and await store.set_status(user_id, "offline"):
        await broadcast_presence(user_id, "offline")


# --- the one event ------------------------------------------------------------------------------
@sio.on(EVENT)
async def on_envelope(sid, data):
    conn = CONNS.get(sid)
    if conn is None:
        return
    ref = data.get("id") if isinstance(data, dict) else None
    room = data.get("room") if isinstance(data, dict) else None
    try:
        if not take_token(conn):
            conn.limited = [t for t in conn.limited if time.monotonic() - t < 10] + [time.monotonic()]
            if len(conn.limited) >= 10:
                return await close(sid, 4429, "rate limited")
            raise ChatError("rate_limited", "slow down")
        validate(data)
        if data["v"] != 1:
            await send_error(sid, ref, ChatError("unsupported_version", "v must be 1"))
            return await close(sid, 4426, "unsupported version")
        kind = data["type"]
        if kind == "hello":
            return await do_hello(sid, conn, data)
        if not conn.ready:
            await send_error(sid, ref, ChatError("unauthorized", "hello first"))
            return await close(sid, 4401, "hello required")
        if kind in FIRE_AND_FORGET:
            return await FIRE_AND_FORGET[kind](sid, conn, data)
        handler = ACKED.get(kind)
        if handler is None:
            raise ChatError("bad_request", f"unknown type {kind!r}")
        cached = await store.cached_ack(conn.user_id, data["id"])
        if cached:
            cached["dup"] = True
            return await send(sid, envelope("ack", room, cached))
        seq, result = await handler(sid, conn, data)
        ack = {"ref": data["id"], "ok": True, "seq": seq, "dup": False, "result": result}
        await store.remember_ack(conn.user_id, data["id"], ack)
        await send(sid, envelope("ack", room, ack))
        if kind == "join" and isinstance(data["payload"].get("since"), int):
            for missed in await store.log_since(data["room"], data["payload"]["since"]):
                await send(sid, missed)  # replay follows the ack
    except ChatError as err:
        await send_error(sid, ref, err, room if isinstance(room, str) else None)
    except RedisError:
        log.warning("redis unavailable")
        await send_error(sid, ref, ChatError("unavailable", "state store unavailable"), room if isinstance(room, str) else None)
    except Exception:
        log.exception("handler failed")
        await send_error(sid, ref, ChatError("internal", "internal error"))


def take_token(conn: Conn) -> bool:
    now = time.monotonic()
    conn.tokens = min(40.0, conn.tokens + (now - conn.stamp) * 20)
    conn.stamp = now
    if conn.tokens >= 1:
        conn.tokens -= 1
        return True
    return False


def validate(data):
    if not isinstance(data, dict) or not isinstance(data.get("v"), int):
        raise ChatError("bad_request", "envelope must be an object with an integer v")
    if not isinstance(data.get("type"), str) or not isinstance(data.get("payload"), dict):
        raise ChatError("bad_request", "type and payload are required")
    if not isinstance(data.get("id"), str) or not 26 <= len(data["id"]) <= 36:
        raise ChatError("bad_request", "id must be 26 to 36 characters")
    if data.get("room") is not None and not isinstance(data["room"], str):
        raise ChatError("bad_request", "room must be a string or null")
    if len(store.dumps(data).encode()) > MAX_FRAME:
        raise ChatError("too_large", "frame exceeds 16384 bytes")


# --- handlers -----------------------------------------------------------------------------------
async def do_hello(sid, conn, env):
    if conn.ready:
        raise ChatError("conflict", "already authenticated")
    payload = env["payload"]
    user_id = await store.redeem_ticket(str(payload.get("token", "")))
    if not user_id:
        await send_error(sid, env["id"], ChatError("unauthorized", "invalid or expired ticket"))
        return await close(sid, 4401, "unauthorized")
    resume = payload.get("resume") or {}
    resumed = bool(resume.get("sessionId")) and await store.session_owner(resume["sessionId"]) == user_id
    conn.user_id, conn.ready = user_id, True
    conn.session_id = resume["sessionId"] if resumed else store.ulid()
    await store.save_session(conn.session_id, user_id)
    rooms = await store.user_rooms(user_id)
    for room in rooms:
        await sio.enter_room(sid, room)
    if await store.r.incr(f"conns:{user_id}") == 1 and await store.set_status(user_id, "online"):
        await broadcast_presence(user_id, "online", skip_sid=sid)
    await send(sid, envelope("hello", None, {
        "ref": env["id"], "userId": user_id, "name": store.display_name(user_id), "sessionId": conn.session_id,
        "heartbeatMs": HEARTBEAT_MS, "maxFrameBytes": MAX_FRAME, "resumed": resumed, "rooms": rooms}))
    if resumed:
        for room, seen in (resume.get("rooms") or {}).items():
            if room in rooms and isinstance(seen, int):
                for missed in await store.log_since(room, seen):
                    await send(sid, missed)


def require_room(env) -> str:
    room = env.get("room")
    if not room:
        raise ChatError("bad_request", "room is required")
    return room


async def need_member(conn, room):
    if not await store.is_member(room, conn.user_id):
        raise ChatError("forbidden", "join the room first")


async def do_join(sid, conn, env):
    room = require_room(env)
    store.check_room_access(room, conn.user_id)
    since = env["payload"].get("since")
    if since is not None and not isinstance(since, int):
        raise ChatError("bad_request", "since must be an integer or null")
    newly = await store.add_member(room, conn.user_id)
    await sio.enter_room(sid, room)
    if newly:
        status = await store.r.hget("presence", conn.user_id) or "online"
        await to_room(room, await presence_env(room, conn.user_id, status), skip_sid=sid)
    return None, {"members": await store.members(room), "lastSeq": await store.last_seq(room)}


async def do_leave(sid, conn, env):
    room = require_room(env)
    if not await store.remove_member(room, conn.user_id):
        raise ChatError("not_found", "not a member")
    await sio.leave_room(sid, room)
    return None, {}


async def do_message(sid, conn, env):
    room, p = require_room(env), env["payload"]
    await need_member(conn, room)
    text = p.get("text")
    if not isinstance(text, str) or not 1 <= store.utf16_len(text) <= 4000:
        raise ChatError("bad_request", "text must be 1 to 4000 characters")
    attachments = p.get("attachments") or []
    if not isinstance(attachments, list) or len(attachments) > 5 or not all(
            isinstance(a, dict) and str(a.get("url", "")).startswith("https://") for a in attachments):
        raise ChatError("bad_request", "at most 5 attachments with https urls")
    msg = await store.append_message(room, conn.user_id, env["id"], text, p.get("replyTo"), attachments)
    await to_room(room, msg)
    return msg["payload"]["seq"], {}


async def do_presence(sid, conn, env):
    status = env["payload"].get("status")
    if status not in ("online", "away"):
        raise ChatError("bad_request", "status must be online or away")
    await store.save_session(conn.session_id, conn.user_id)  # the heartbeat keeps the session alive
    if await store.set_status(conn.user_id, status):
        await broadcast_presence(conn.user_id, status, skip_sid=sid)
    return None, {}


async def do_history(sid, conn, env):
    room, p = require_room(env), env["payload"]
    await need_member(conn, room)
    limit = p.get("limit", 50)
    if not isinstance(limit, int) or not 1 <= limit <= 100:
        raise ChatError("bad_request", "limit must be 1 to 100")
    return None, await store.history(room, p.get("before"), limit)


async def do_typing(sid, conn, env):
    room = env.get("room")
    if room and await store.is_member(room, conn.user_id) and env["payload"].get("state") in ("start", "stop"):
        await to_room(room, envelope("typing", room, {"userId": conn.user_id, "state": env["payload"]["state"]}),
                      skip_sid=sid)


ACKED = {"join": do_join, "leave": do_leave, "message": do_message, "presence": do_presence, "history": do_history}
FIRE_AND_FORGET = {"typing": do_typing}


# --- HTTP part (FastAPI) and the combined ASGI app ----------------------------------------------
api = FastAPI(title="Chatterbox")


@api.post("/ticket")
async def ticket(user: str):
    try:
        return {"ticket": await store.issue_ticket(user)}
    except ChatError as err:
        raise HTTPException(400, err.message)
    except RedisError:
        raise HTTPException(503, "state store unavailable")


@api.get("/healthz")
async def healthz():
    return {"ok": True}


asgi = socketio.ASGIApp(sio, other_asgi_app=api)

"""Post a system message into a room from a process that holds no connections."""
import asyncio
import sys

import socketio

import store


async def main(room: str, text: str):
    # write_only=True: publish to Redis only, never subscribe (and so never receive callbacks)
    external = socketio.AsyncServer(client_manager=socketio.AsyncRedisManager(store.REDIS_URL, write_only=True))
    msg = await store.append_message(room, "system", store.ulid(), text)
    await external.emit(store.EVENT, msg, room=room)
    print("announced seq", msg["payload"]["seq"])


asyncio.run(main(sys.argv[1], sys.argv[2]))

upstream chatterbox {
    server chatterbox-1:8000;
    server chatterbox-2:8000;
}

map $http_upgrade $connection_upgrade {
    default upgrade;
    ""      close;
}

server {
    listen 80;

    location /socket.io/ {
        proxy_pass http://chatterbox;
        proxy_http_version 1.1;
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection $connection_upgrade;
        proxy_set_header Host $host;
        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto $scheme;
        proxy_read_timeout 75s;   # must exceed the 20 s Engine.IO ping interval
    }

    location / {
        proxy_pass http://chatterbox;
        proxy_set_header Host $host;
    }
}