WebSocket 是会话,不是页面

"""Probe WebSocket session semantics with stdlib only: handshake + heartbeat."""
import base64
import hashlib
import socket
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer

GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
# RFC 6455 section 1.2 example key and its expected accept value.
KEY = "dGhlIHNhbXBsZSBub25jZQ=="
EXPECTED_ACCEPT = "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="


def compute_accept(key: str) -> str:
    return base64.b64encode(hashlib.sha1((key + GUID).encode()).digest()).decode()


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def do_GET(self):
        key = self.headers.get("Sec-WebSocket-Key")
        upgrade = (self.headers.get("Upgrade") or "").lower()
        connection = (self.headers.get("Connection") or "").lower()
        ok = (
            self.path == "/chat"
            and key is not None
            and upgrade == "websocket"
            and "upgrade" in connection
        )
        if not ok:
            body = b"expected websocket upgrade"
            self.send_response(400)
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        accept = compute_accept(key.strip())
        self.send_response(101, "Switching Protocols")
        self.send_header("Upgrade", "websocket")
        self.send_header("Connection", "Upgrade")
        self.send_header("Sec-WebSocket-Accept", accept)
        self.send_header("Content-Length", "0")
        self.end_headers()

    def log_message(self, *args):
        pass


def read_http_response(sock: socket.socket) -> str:
    data = b""
    sock.settimeout(5)
    while b"\r\n\r\n" not in data:
        chunk = sock.recv(4096)
        if not chunk:
            break
        data += chunk
    return data.decode("latin-1")


def main() -> None:
    # 1. The accept derivation itself must match the RFC 6455 test vector.
    assert compute_accept(KEY) == EXPECTED_ACCEPT, "accept mismatch"
    print(f"accept vector ok: {EXPECTED_ACCEPT}")

    # 2. A real loopback handshake: raw request, assert 101 + accept echo.
    server = HTTPServer(("127.0.0.1", 0), Handler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    try:
        port = server.server_address[1]
        with socket.create_connection(("127.0.0.1", port), timeout=5) as sock:
            sock.sendall(
                (
                    "GET /chat HTTP/1.1\r\n"
                    "Host: 127.0.0.1\r\n"
                    "Upgrade: websocket\r\n"
                    "Connection: Upgrade\r\n"
                    f"Sec-WebSocket-Key: {KEY}\r\n"
                    "Sec-WebSocket-Version: 13\r\n"
                    "\r\n"
                ).encode()
            )
            head = read_http_response(sock)
        assert "101" in head.splitlines()[0], head
        assert f"Sec-WebSocket-Accept: {EXPECTED_ACCEPT}" in head, head
        print("handshake ok: 101 + Sec-WebSocket-Accept matches key")

        # 3. A wrong path is not a session: it must stay plain HTTP.
        with socket.create_connection(("127.0.0.1", port), timeout=5) as sock:
            sock.sendall(b"GET /page HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n")
            head = read_http_response(sock)
        assert "400" in head.splitlines()[0], head
        print("non-upgrade ok: plain GET gets 400, no session created")

        # 4. Heartbeat semantics over a socketpair: a pong in time keeps the
        #    session, silence past the deadline marks a ghost connection.
        alive, peer = socket.socketpair()
        with alive, peer:
            alive.settimeout(2)
            peer.settimeout(2)
            alive.sendall(b"ping-1")
            assert peer.recv(1024) == b"ping-1"
            peer.sendall(b"pong-1")
            assert alive.recv(1024) == b"pong-1"
            print("heartbeat ok: ping/pong round-trip keeps session alive")
            peer.settimeout(0.2)
            try:
                peer.recv(1024)
                raise AssertionError("expected silence, got data")
            except socket.timeout:
                print("ghost detection ok: missed pong past deadline = stale session")
    finally:
        server.shutdown()
        server.server_close()
    print("PASS: handshake, rejection, heartbeat, and ghost detection hold")


if __name__ == "__main__":
    main()