WebSockets

import asyncio
import json
import secrets

from tornado import web
from tornado.httpclient import HTTPClientError, HTTPRequest
from tornado.httpserver import HTTPServer
from tornado.testing import bind_unused_port
from tornado.websocket import (
    WebSocketClosedError,
    WebSocketHandler,
    websocket_connect,
)


async def serve(app: web.Application) -> tuple[HTTPServer, int]:
    """Start `app` on a free loopback port and return (server, port)."""
    sock, port = bind_unused_port()
    server = HTTPServer(app)
    server.add_socket(sock)
    return server, port

class EchoSocket(WebSocketHandler):
    def open(self):
        # The handshake is complete; the connection is usable.
        self.write_message({"type": "hello", "path": self.request.path})

    def on_message(self, message):
        # `message` is str for a text frame and bytes for a binary frame.
        if isinstance(message, bytes):
            self.write_message(message, binary=True)
        else:
            self.write_message("You said: " + message)

    def on_close(self):
        # close_code / close_reason are set when the peer sent them.
        print("closed:", self.close_code, self.close_reason)


async def echo_demo():
    server, port = await serve(web.Application([(r"/echo", EchoSocket)]))
    conn = await websocket_connect(f"ws://127.0.0.1:{port}/echo")
    print(await conn.read_message())          # the JSON hello
    await conn.write_message("hello")
    print(await conn.read_message())          # You said: hello
    await conn.write_message(b"\x00\x01", binary=True)
    print(await conn.read_message())          # b'\x00\x01'
    conn.close(1000, "done")
    print(await conn.read_message())          # None: the connection is closed
    await asyncio.sleep(0.05)                 # let the server run on_close
    server.stop()


asyncio.run(echo_demo())

class ReviewSocket(WebSocketHandler):
    # Origins of the web front ends allowed to open this socket.
    allowed_origins = {"https://bookshelf.example", "https://www.bookshelf.example"}

    def check_origin(self, origin: str) -> bool:
        # Exact match of scheme + host + port, never a substring or endswith test.
        return origin in self.allowed_origins

    def open(self):
        self.write_message("welcome")


async def origin_demo():
    server, port = await serve(web.Application([(r"/live", ReviewSocket)]))
    url = f"ws://127.0.0.1:{port}/live"

    # A browser page on an unknown site: rejected during the handshake.
    try:
        await websocket_connect(
            HTTPRequest(url, headers={"Origin": "https://evil.example"})
        )
    except HTTPClientError as exc:
        print("rejected:", exc.code)

    # The real front end: accepted.
    conn = await websocket_connect(
        HTTPRequest(url, headers={"Origin": "https://bookshelf.example"})
    )
    print(await conn.read_message())

    # No Origin header (a script, not a browser): the check is skipped.
    other = await websocket_connect(url)
    print(await other.read_message())
    conn.close()
    other.close()
    server.stop()


asyncio.run(origin_demo())

class PongCounter(WebSocketHandler):
    pongs = 0

    def on_pong(self, data: bytes) -> None:
        type(self).pongs += 1

    def on_message(self, message):
        self.write_message(f"received {len(message)} characters")


async def settings_demo():
    app = web.Application(
        [(r"/ws", PongCounter)],
        websocket_ping_interval=0.05,     # seconds; use 20-30 in production
        websocket_ping_timeout=0.05,
        websocket_max_message_size=1024,  # bytes; the default is 10 MiB
    )
    server, port = await serve(app)
    conn = await websocket_connect(f"ws://127.0.0.1:{port}/ws")
    await asyncio.sleep(0.3)                  # idle: only pings and pongs flow
    print("pongs received:", PongCounter.pongs >= 2)

    await conn.write_message("x" * 4000)      # larger than the limit
    print(await conn.read_message())          # None: the server closed it
    print("close code:", conn.close_code)     # 1009
    server.stop()


asyncio.run(settings_demo())

class VersionedSocket(WebSocketHandler):
    SUPPORTED = ("bookshelf.v2", "bookshelf.v1")

    def select_subprotocol(self, subprotocols):
        for name in self.SUPPORTED:         # server preference order
            if name in subprotocols:
                return name
        return None

    def get_compression_options(self):
        return {"compression_level": 6, "mem_level": 5}

    def open(self):
        self.write_message({"protocol": self.selected_subprotocol})


async def negotiation_demo():
    server, port = await serve(web.Application([(r"/v", VersionedSocket)]))
    conn = await websocket_connect(
        f"ws://127.0.0.1:{port}/v",
        subprotocols=["bookshelf.v1", "bookshelf.v2"],
        compression_options={},            # same meaning as on the server
    )
    print("selected:", conn.selected_subprotocol)
    print(await conn.read_message())
    conn.close()
    server.stop()


asyncio.run(negotiation_demo())

class Rooms:
    """Per-process registry: room name -> open WebSocket handlers."""

    def __init__(self):
        self._rooms: dict[str, set["ReviewRoomSocket"]] = {}

    def join(self, room: str, member: "ReviewRoomSocket") -> None:
        self._rooms.setdefault(room, set()).add(member)

    def leave(self, room: str, member: "ReviewRoomSocket") -> None:
        members = self._rooms.get(room)
        if members is not None:
            members.discard(member)
            if not members:
                del self._rooms[room]       # do not leak empty rooms

    def size(self, room: str) -> int:
        return len(self._rooms.get(room, ()))

    async def broadcast(self, room: str, payload: dict, skip=None) -> int:
        text = json.dumps(payload)
        targets = [m for m in self._rooms.get(room, ()) if m is not skip]
        results = await asyncio.gather(
            *(m.send(text) for m in targets), return_exceptions=True
        )
        return sum(1 for r in results if r is None)


class ReviewRoomSocket(WebSocketHandler):
    def initialize(self, rooms: Rooms):
        self.rooms = rooms                   # passed from the route's kwargs
        self.room = None

    def open(self, book_id: str):
        self.room = f"book-{book_id}"
        self.rooms.join(self.room, self)

    async def on_message(self, message):
        try:
            data = json.loads(message)
            text = str(data["text"])[:500]   # bound what one client can send
        except (ValueError, KeyError, TypeError):
            return self.write_message({"type": "error"})
        await self.rooms.broadcast(
            self.room, {"type": "review", "text": text}, skip=self
        )

    def on_close(self):
        if self.room is not None:
            self.rooms.leave(self.room, self)

    async def send(self, text: str) -> None:
        try:
            await self.write_message(text)
        except WebSocketClosedError:
            self.rooms.leave(self.room, self)


async def rooms_demo():
    rooms = Rooms()
    app = web.Application(
        [(r"/ws/books/(\d+)/live", ReviewRoomSocket, dict(rooms=rooms))]
    )
    server, port = await serve(app)
    url = f"ws://127.0.0.1:{port}/ws/books"
    a = await websocket_connect(f"{url}/7/live")
    b = await websocket_connect(f"{url}/7/live")
    c = await websocket_connect(f"{url}/9/live")       # another book
    await asyncio.sleep(0.05)
    print("members of book-7:", rooms.size("book-7"))

    await a.write_message(json.dumps({"text": "Great read"}))
    print("B got:", await b.read_message())
    try:
        await asyncio.wait_for(c.read_message(), 0.2)
    except asyncio.TimeoutError:
        print("C (other room) got nothing")

    b.close()
    await asyncio.sleep(0.1)
    print("members after B left:", rooms.size("book-7"))
    for conn in (a, c):
        conn.close()
    server.stop()


asyncio.run(rooms_demo())

class LoginHandler(web.RequestHandler):
    def post(self):
        name = self.get_body_argument("user")
        # Credential check omitted: see the OAuth and authentication pages for real flows.
        self.set_signed_cookie(
            "user", name, httponly=True, samesite="Strict"
        )  # add secure=True behind HTTPS
        self.write({"ok": True})


class PrivateSocket(WebSocketHandler):
    def get_current_user(self):
        raw = self.get_signed_cookie("user")      # bytes or None; signature verified
        return raw.decode() if raw else None

    def check_origin(self, origin: str) -> bool:
        return origin in {"https://bookshelf.example"}

    def prepare(self):
        if self.current_user is None:
            raise web.HTTPError(401, reason="Login required")

    def open(self):
        self.write_message({"type": "welcome", "user": self.current_user})


def make_private_app() -> web.Application:
    return web.Application(
        [(r"/login", LoginHandler), (r"/private", PrivateSocket)],
        cookie_secret=secrets.token_hex(32),       # in production: read it from the environment
    )


async def auth_demo():
    from tornado.httpclient import AsyncHTTPClient

    server, port = await serve(make_private_app())
    origin = {"Origin": "https://bookshelf.example"}
    url = f"ws://127.0.0.1:{port}/private"

    # Anonymous handshake: rejected before the upgrade.
    try:
        await websocket_connect(HTTPRequest(url, headers=origin))
    except HTTPClientError as exc:
        print("anonymous:", exc.code)

    # Log in over HTTP, then reuse the cookie on the handshake.
    client = AsyncHTTPClient()
    response = await client.fetch(
        f"http://127.0.0.1:{port}/login", method="POST", body="user=alice"
    )
    cookie = response.headers["Set-Cookie"].split(";")[0]
    conn = await websocket_connect(
        HTTPRequest(url, headers={**origin, "Cookie": cookie})
    )
    print(await conn.read_message())
    conn.close()
    server.stop()


asyncio.run(auth_demo())

async def client_demo():
    server, port = await serve(web.Application([(r"/echo", EchoSocket)]))
    url = f"ws://127.0.0.1:{port}/echo"

    # Pull style: a read loop that ends when the connection closes.
    conn = await websocket_connect(url, connect_timeout=5)
    received = []
    await conn.write_message("one")
    await conn.write_message("two")
    while len(received) < 3:                      # hello + two replies
        message = await conn.read_message()
        if message is None:
            break
        received.append(message)
    print(received)
    conn.close()

    # Push style: a callback receives every message, then None on close.
    seen = []
    done = asyncio.Event()

    def on_message(message):
        if message is None:
            done.set()
        else:
            seen.append(message)

    pushed = await websocket_connect(url, on_message_callback=on_message)
    await pushed.write_message("three")
    await asyncio.sleep(0.1)
    pushed.close()
    await asyncio.wait_for(done.wait(), 2)
    print(seen)
    server.stop()


asyncio.run(client_demo())

class FakeBroker:
    """In-memory stand-in for Redis pub/sub, shared by the 'processes' below."""

    def __init__(self):
        self._queues: list[asyncio.Queue] = []

    async def publish(self, channel: str, data: str) -> None:
        for queue in list(self._queues):
            queue.put_nowait((channel, data))

    async def messages(self):
        queue: asyncio.Queue = asyncio.Queue()
        self._queues.append(queue)
        try:
            while True:
                yield await queue.get()
        finally:
            self._queues.remove(queue)


class Hub(Rooms):
    """A process-local Rooms registry fed by a broker subscription."""

    def __init__(self, broker):
        super().__init__()
        self.broker = broker
        self._task = None

    async def publish(self, room: str, payload: dict) -> None:
        await self.broker.publish(f"room:{room}", json.dumps(payload))

    def start(self) -> None:
        self._task = asyncio.create_task(self._listen())   # keep the reference

    async def _listen(self) -> None:
        async for channel, data in self.broker.messages():
            room = channel.removeprefix("room:")
            await self.broadcast(room, json.loads(data))

    def stop(self) -> None:
        if self._task:
            self._task.cancel()


class ClusterSocket(ReviewRoomSocket):
    def initialize(self, rooms: Hub):
        super().initialize(rooms)

    async def on_message(self, message):
        try:
            text = str(json.loads(message)["text"])[:500]
        except (ValueError, KeyError, TypeError):
            return self.write_message({"type": "error"})
        await self.rooms.publish(self.room, {"type": "review", "text": text})


async def cluster_demo():
    broker = FakeBroker()
    hubs, servers, ports = [], [], []
    for _ in range(2):                                # two "worker processes"
        hub = Hub(broker)
        hub.start()
        server, port = await serve(
            web.Application([(r"/ws/books/(\d+)/live", ClusterSocket, dict(rooms=hub))])
        )
        hubs.append(hub)
        servers.append(server)
        ports.append(port)
    await asyncio.sleep(0.05)

    a = await websocket_connect(f"ws://127.0.0.1:{ports[0]}/ws/books/7/live")
    b = await websocket_connect(f"ws://127.0.0.1:{ports[1]}/ws/books/7/live")
    await asyncio.sleep(0.05)
    await a.write_message(json.dumps({"text": "Seen on another process"}))
    print("B (other process) got:", await asyncio.wait_for(b.read_message(), 2))
    print("A (sender) got:", await asyncio.wait_for(a.read_message(), 2))

    for conn in (a, b):
        conn.close()
    for hub in hubs:
        hub.stop()
    for server in servers:
        server.stop()


asyncio.run(cluster_demo())

# norun: needs a running Redis server and the redis package
import os

import redis.asyncio as redis


class RedisBroker:
    def __init__(self, url: str | None = None):
        self._redis = redis.from_url(url or os.environ["REDIS_URL"], decode_responses=True)

    async def publish(self, channel: str, data: str) -> None:
        await self._redis.publish(channel, data)

    async def messages(self):
        pubsub = self._redis.pubsub()
        await pubsub.psubscribe("room:*")
        try:
            async for item in pubsub.listen():
                if item["type"] == "pmessage":
                    yield item["channel"], item["data"]
        finally:
            await pubsub.aclose()

location /ws/ {
    proxy_pass http://tornado_backend;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_set_header Host $host;
    proxy_read_timeout 120s;
}

import unittest

from tornado.testing import AsyncHTTPTestCase, gen_test


class LiveTest(AsyncHTTPTestCase):
    def get_app(self):
        self.rooms = Rooms()
        return web.Application(
            [
                (r"/echo", EchoSocket),
                (r"/live", ReviewSocket),
                (r"/ws/books/(\d+)/live", ReviewRoomSocket, dict(rooms=self.rooms)),
            ]
        )

    def ws_url(self, path: str) -> str:
        return f"ws://127.0.0.1:{self.get_http_port()}{path}"

    @gen_test
    async def test_echo(self):
        conn = await websocket_connect(self.ws_url("/echo"))
        await conn.read_message()                 # the hello message
        await conn.write_message("ping")
        self.assertEqual(await conn.read_message(), "You said: ping")
        conn.close()

    @gen_test
    async def test_foreign_origin_is_rejected(self):
        request = HTTPRequest(self.ws_url("/live"), headers={"Origin": "https://evil.example"})
        with self.assertRaises(HTTPClientError) as ctx:
            await websocket_connect(request)
        self.assertEqual(ctx.exception.code, 403)

    @gen_test
    async def test_room_fan_out(self):
        a = await websocket_connect(self.ws_url("/ws/books/1/live"))
        b = await websocket_connect(self.ws_url("/ws/books/1/live"))
        await a.write_message(json.dumps({"text": "hi"}))
        reply = json.loads(await b.read_message())
        self.assertEqual(reply, {"type": "review", "text": "hi"})
        a.close()
        b.close()


suite = unittest.defaultTestLoader.loadTestsFromTestCase(LiveTest)
result = unittest.TextTestRunner(verbosity=0).run(suite)
assert result.wasSuccessful()

class SlowJobSocket(WebSocketHandler):
    _tasks: set[asyncio.Task] = set()      # the loop keeps only weak references to tasks

    async def on_message(self, message):
        if message == "slow":
            task = asyncio.create_task(self._slow_job())
            self._tasks.add(task)
            task.add_done_callback(self._tasks.discard)
            await self.write_message("started")
        else:
            await self.write_message("fast reply")

    async def _slow_job(self):
        await asyncio.sleep(0.2)
        try:
            await self.write_message("slow done")
        except WebSocketClosedError:
            pass                            # the client left meanwhile


async def pitfall_demo():
    server, port = await serve(web.Application([(r"/slow", SlowJobSocket)]))
    conn = await websocket_connect(f"ws://127.0.0.1:{port}/slow")
    await conn.write_message("slow")
    await conn.write_message("hello")
    for _ in range(3):
        print(await conn.read_message())   # started, fast reply, slow done
    conn.close()
    server.stop()


asyncio.run(pitfall_demo())