fleetmesh-node 0.2.0__py3-none-any.whl

node/transport/ws.py

@@ -0,0 +1,502 @@
+"""Enough of RFC 6455 for a node: a client, a server, text frames, ping and close.
+
+Standard library only, blocking sockets with timeouts. scripts/nostr.py has a
+smaller version of the client half for publishing one event; this one is for
+a process that holds connections open: subscriptions on a relay, and the
+relay face this node itself presents. Fragmented messages are reassembled;
+binary frames are accepted and returned as bytes; a ping is answered on the
+spot; a close frame is answered and raised as ConnectionClosed.
+"""
+from __future__ import annotations
+
+import base64
+import enum
+import errno
+import hashlib
+import os
+import socket
+import ssl
+import struct
+import threading
+import time
+from typing import Any, Callable, Mapping, Optional
+from urllib.parse import urlparse
+
+GUID: bytes = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
+class Opcode(enum.IntEnum):
+    """RFC 6455 frame opcodes. An IntEnum so the names read in a match
+    statement as value patterns and still compare equal to the raw ints the
+    frame decoder yields and that callers pass to send()."""
+    CONT = 0x0
+    TEXT = 0x1
+    BINARY = 0x2
+    CLOSE = 0x8
+    PING = 0x9
+    PONG = 0xA
+
+
+# Module-level aliases keep the existing spelling used across the tree.
+TEXT: int = Opcode.TEXT
+BINARY: int = Opcode.BINARY
+CLOSE: int = Opcode.CLOSE
+PING: int = Opcode.PING
+PONG: int = Opcode.PONG
+CONT: int = Opcode.CONT
+MAX_FRAME: int = 16 * 1024 * 1024
+# The whole HTTP upgrade, either direction. Headers are a few hundred bytes.
+HANDSHAKE_MAX: int = 64 * 1024
+
+
+class ConnectionClosed(ConnectionError):
+    pass
+
+
+class PlainHTTP(Exception):
+    """A request on the websocket port that did not ask to upgrade. The relay
+    face answers NIP-11 to these, and Blossom handles blob endpoints."""
+
+    def __init__(
+        self,
+        path: str,
+        headers: dict[bytes, bytes],
+        method: str = 'GET',
+        initial_body: bytes = b'',
+        sock: Optional[socket.socket] = None,
+    ) -> None:
+        super().__init__(path)
+        self.path: str = path
+        self.headers: dict[bytes, bytes] = headers
+        self.method: str = method
+        self.initial_body: bytes = initial_body
+        self.sock: Optional[socket.socket] = sock
+
+
+def http_response(
+    sock: socket.socket,
+    status: str,
+    body: bytes = b'',
+    content_type: str = 'application/json',
+    headers: Optional[Mapping[str, str]] = None,
+) -> None:
+    cl = str(len(body))
+    extra_hdrs: list[tuple[str, str]] = []
+    if headers:
+        for k, v in headers.items():
+            if k.lower() == 'content-length':
+                cl = str(v)
+            elif k.lower() == 'content-type':
+                content_type = str(v)
+            else:
+                extra_hdrs.append((k, str(v)))
+
+    hdrs = [
+        'HTTP/1.1 %s' % status,
+        'Content-Type: %s' % content_type,
+        'Content-Length: %s' % cl,
+        'Access-Control-Allow-Origin: *',
+        'Access-Control-Allow-Methods: GET, HEAD, PUT, POST, DELETE, OPTIONS',
+        'Access-Control-Allow-Headers: Authorization, Content-Type, X-SHA-256, X-CID',
+        'Connection: close',
+    ]
+    for k, v in extra_hdrs:
+        hdrs.append('%s: %s' % (k, v))
+    head = '\r\n'.join(hdrs) + '\r\n\r\n'
+    sock.sendall(head.encode('latin1') + body)
+
+
+def read_body(conn: socket.socket, req: PlainHTTP, max_bytes: int = 64 * 1024 * 1024) -> bytes:
+    """Read the full HTTP body from connection based on Content-Length."""
+    cl = req.headers.get(b'content-length')
+    if not cl:
+        return req.initial_body
+    try:
+        length = int(cl.decode('ascii'))
+    except (ValueError, TypeError):
+        return req.initial_body
+
+    if length > max_bytes:
+        raise ValueError('request body length %d exceeds limit %d' % (length, max_bytes))
+
+    buf = bytearray(req.initial_body)
+    while len(buf) < length:
+        chunk = conn.recv(min(65536, length - len(buf)))
+        if not chunk:
+            break
+        buf.extend(chunk)
+    return bytes(buf)
+
+
+def _read_exact(sock: socket.socket | _Prefixed, n: int) -> bytes:
+    buf = b''
+    while len(buf) < n:
+        chunk = sock.recv(n - len(buf))
+        if not chunk:
+            raise ConnectionClosed('peer closed the connection')
+        buf += chunk
+    return buf
+
+
+def _mask(payload: bytes, key: bytes) -> bytes:
+    """XOR every byte of payload with the 4-byte key, repeating (RFC 6455 §5.3).
+
+    Done as one wide integer XOR rather than a byte-at-a-time generator: the
+    result is identical and the work moves into C. Measured 6x faster on a
+    50-byte frame and 20-30x on 1KiB-1MiB, which is where a relay spends
+    its time when it is busy.
+    """
+    n = len(payload)
+    if not n:
+        return payload
+    full = key * (n // 4) + key[:n % 4]
+    return (int.from_bytes(payload, 'big') ^ int.from_bytes(full, 'big')).to_bytes(n, 'big')
+
+
+def encode_frame(opcode: int, payload: bytes, mask: bool) -> bytes:
+    head = bytes([0x80 | opcode])
+    n = len(payload)
+    flag = 0x80 if mask else 0
+    if n < 126:
+        head += bytes([flag | n])
+    elif n < (1 << 16):
+        head += bytes([flag | 126]) + struct.pack('!H', n)
+    else:
+        head += bytes([flag | 127]) + struct.pack('!Q', n)
+    if mask:
+        key = os.urandom(4)
+        return head + key + _mask(payload, key)
+    return head + payload
+
+
+def decode_frame(sock: socket.socket | _Prefixed) -> tuple[bool, int, bytes]:
+    """(fin, opcode, payload) for one frame, unmasked."""
+    b1, b2 = _read_exact(sock, 2)
+    fin, opcode = bool(b1 & 0x80), b1 & 0x0F
+    masked, n = bool(b2 & 0x80), b2 & 0x7F
+    if n == 126:
+        n = struct.unpack('!H', _read_exact(sock, 2))[0]
+    elif n == 127:
+        n = struct.unpack('!Q', _read_exact(sock, 8))[0]
+    if n > MAX_FRAME:
+        raise ConnectionClosed('frame of %d bytes refused' % n)
+    key = _read_exact(sock, 4) if masked else None
+    payload = _read_exact(sock, n)
+    if key:
+        payload = _mask(payload, key)
+    return fin, opcode, payload
+
+
+class WebSocket:
+    """One connection after the handshake. Clients mask, servers do not."""
+
+    def __init__(self, sock: socket.socket | _Prefixed, client_side: bool) -> None:
+        self.sock: socket.socket | _Prefixed = sock
+        self.mask: bool = client_side
+        self._lock: threading.Lock = threading.Lock()
+        self.closed: bool = False
+        self.last_rx: float = time.time()      # any frame at all, control frames included
+
+    def ping(self, payload: bytes = b'') -> None:
+        self.send(PING, payload)
+
+    def idle(self) -> float:
+        """Seconds since the peer last sent any frame."""
+        return time.time() - self.last_rx
+
+    def send(self, opcode: int, payload: bytes) -> None:
+        with self._lock:
+            self.sock.sendall(encode_frame(opcode, payload, self.mask))
+
+    def send_text(self, text: str) -> None:
+        self.send(TEXT, text.encode('utf-8'))
+
+    def recv(self) -> tuple[int, bytes]:
+        """The next complete message: (opcode, payload). Handles control frames.
+
+        MAX_FRAME bounds one frame in decode_frame; a message is any number of
+        continuation frames, and the sum was not bounded, so a peer could send
+        a megabyte at a time for as long as it liked and the list here grew
+        with it. The same cap applies to the whole message.
+        """
+        parts: list[bytes] = []
+        first: Optional[int] = None
+        total = 0
+        while True:
+            fin, opcode, payload = decode_frame(self.sock)
+            self.last_rx = time.time()
+            # Control frames may arrive between the fragments of a message
+            # (RFC 6455 §5.4), so they are answered here and never touch the
+            # buffer; anything that is not CONT begins a new message.
+            match opcode:
+                case Opcode.PING:
+                    self.send(PONG, payload)
+                    continue
+                case Opcode.PONG:
+                    continue
+                case Opcode.CLOSE:
+                    self.close(echo=payload[:2])
+                    raise ConnectionClosed('close frame received')
+                case Opcode.CONT:
+                    parts.append(payload)
+                    total += len(payload)
+                case _:
+                    first = opcode
+                    parts = [payload]
+                    total = len(payload)
+            if total > MAX_FRAME:
+                self.close(echo=struct.pack('!H', 1009))       # 1009: message too big
+                raise ConnectionClosed('message larger than %d bytes' % MAX_FRAME)
+            if fin:
+                assert first is not None
+                return first, b''.join(parts)
+
+    def recv_text(self) -> str:
+        opcode, payload = self.recv()
+        return payload.decode('utf-8') if opcode == TEXT else payload.decode('utf-8', errors='replace')
+
+    def close(self, echo: bytes = b'') -> None:
+        if self.closed:
+            return
+        self.closed = True
+        try:
+            self.send(CLOSE, echo or struct.pack('!H', 1000))
+        except OSError:
+            pass
+        try:
+            self.sock.close()
+        except OSError:
+            pass
+
+
+def connect(url: str, timeout: float = 15.0) -> WebSocket:
+    u = urlparse(url)
+    host = u.hostname or '127.0.0.1'
+    port = u.port or (443 if u.scheme == 'wss' else 80)
+    path = u.path or '/'
+    raw = socket.create_connection((host, port), timeout=timeout)
+    if u.scheme == 'wss':
+        try:
+            sock: socket.socket | ssl.SSLSocket = ssl.create_default_context().wrap_socket(raw, server_hostname=host)
+        except Exception:
+            raw.close()
+            raise
+    else:
+        sock = raw
+
+    try:
+        key = base64.b64encode(os.urandom(16))
+        sock.sendall(b'GET %s HTTP/1.1\r\nHost: %s\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n'
+                     b'Sec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\n\r\n'
+                     % (path.encode(), host.encode(), key))
+        resp = b''
+        while b'\r\n\r\n' not in resp:
+            chunk = sock.recv(4096)
+            if not chunk:
+                raise ConnectionClosed('no handshake response')
+            resp += chunk
+            if len(resp) > HANDSHAKE_MAX:
+                # The server side has had this cap since the start; a relay
+                # that answers the upgrade with an endless header is the
+                # same shape of peer as a client that sends one.
+                raise ConnectionClosed('handshake too large')
+        head, _, rest = resp.partition(b'\r\n\r\n')
+        lines = head.split(b'\r\n')
+        if b' 101 ' not in lines[0]:
+            raise ConnectionError('handshake refused: %s' % lines[0].decode(errors='replace'))
+        want = base64.b64encode(hashlib.sha1(key + GUID, usedforsecurity=False).digest())
+        accept = [l.split(b':', 1)[1].strip() for l in lines[1:] if l.lower().startswith(b'sec-websocket-accept:')]
+        if not accept or accept[0] != want:
+            raise ConnectionError('bad Sec-WebSocket-Accept')
+        final_sock: socket.socket | _Prefixed = _Prefixed(sock, rest) if rest else sock
+        return WebSocket(final_sock, client_side=True)
+    except Exception:
+        try:
+            sock.close()
+        except Exception:
+            pass
+        raise
+
+
+class _Prefixed:
+    """A socket with bytes already read; the first recv drains them."""
+
+    def __init__(self, sock: socket.socket | ssl.SSLSocket, buf: bytes) -> None:
+        self._sock: socket.socket | ssl.SSLSocket = sock
+        self._buf: bytes = buf
+
+    def recv(self, n: int) -> bytes:
+        if self._buf:
+            out, self._buf = self._buf[:n], self._buf[n:]
+            return out
+        return self._sock.recv(n)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._sock, name)
+
+
+def accept(sock: socket.socket) -> tuple[WebSocket, str]:
+    """Server side of the handshake on an accepted socket. Returns (WebSocket, path)."""
+    req = b''
+    while b'\r\n\r\n' not in req:
+        chunk = sock.recv(4096)
+        if not chunk:
+            raise ConnectionClosed('no handshake request')
+        req += chunk
+        if len(req) > HANDSHAKE_MAX:
+            raise ConnectionClosed('handshake too large')
+    head, _, rest = req.partition(b'\r\n\r\n')
+    lines = head.split(b'\r\n')
+    parts = lines[0].split(b' ')
+    method = parts[0].decode('latin1', errors='replace') if len(parts) > 0 else 'GET'
+    path = parts[1].decode('latin1', errors='replace') if len(parts) > 1 else '/'
+    headers: dict[bytes, bytes] = {}
+    for l in lines[1:]:
+        if b':' in l:
+            k, v = l.split(b':', 1)
+            headers[k.strip().lower()] = v.strip()
+    key = headers.get(b'sec-websocket-key')
+    if not key or b'websocket' not in headers.get(b'upgrade', b'').lower():
+        raise PlainHTTP(path, headers, method=method, initial_body=rest, sock=sock)
+    acc = base64.b64encode(hashlib.sha1(key + GUID, usedforsecurity=False).digest())
+    sock.sendall(b'HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n'
+                 b'Sec-WebSocket-Accept: %s\r\n\r\n' % acc)
+    final_sock: socket.socket | _Prefixed = _Prefixed(sock, rest) if rest else sock
+    return WebSocket(final_sock, client_side=False), path
+
+
+class Server:
+    """Accepts connections and hands each to handler(ws, path) on its own thread."""
+
+    # How long a bind waits for a port it has just been using to come free.
+    # A node's relay face is on a recorded port so that clients keep its
+    # address across a restart, and `systemctl restart` closes that port and
+    # rebinds it within the same second. On a small machine the kernel can
+    # still hold the closed listener - observed on a 1-vCPU droplet, where
+    # `ss` showed LISTEN on the port with no owning process after close()
+    # had returned - and the node then died on startup with EADDRINUSE
+    # instead of coming back. SO_REUSEADDR does not cover this; waiting does.
+    BIND_WAIT: float = 10.0
+
+    def __init__(
+        self,
+        handler: Callable[[WebSocket, str], Any],
+        host: str = '127.0.0.1',
+        port: int = 0,
+        http_handler: Optional[Callable[..., Any]] = None,
+        bind_wait: Optional[float] = None,
+    ) -> None:
+        self.handler: Callable[[WebSocket, str], Any] = handler
+        # Callable[..., Any] because two call shapes are honoured: the newer
+        # (conn, PlainHTTP) and the older (conn, path, headers), tried in turn
+        # in _serve. A fixed signature would misdescribe one of them.
+        self.http_handler: Optional[Callable[..., Any]] = http_handler
+        self.sock: socket.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        self._bind(host, port, self.BIND_WAIT if bind_wait is None else bind_wait)
+        self.sock.listen(16)
+        sock_name = self.sock.getsockname()
+        self.host: str = sock_name[0]
+        self.port: int = sock_name[1]
+        self._stop: threading.Event = threading.Event()
+        self._thread: threading.Thread = threading.Thread(target=self._loop, daemon=True)
+        self.connections: list[WebSocket] = []
+        self._conn_lock: threading.Lock = threading.Lock()
+        self.errors: list[str] = []
+
+    def _bind(self, host: str, port: int, wait: float) -> None:
+        """Bind, waiting out a port that is closing. Never waits on port 0.
+
+        Port 0 asks the kernel for anything free, so EADDRINUSE there is a
+        real fault and retrying it would only hide one. A named port is the
+        one an operator chose, and the only honest answer to "it is still
+        going" is to wait a moment and say so if it never comes.
+        """
+        deadline = time.time() + (wait if port else 0)
+        while True:
+            try:
+                self.sock.bind((host, port))
+                return
+            except OSError as e:
+                if e.errno != errno.EADDRINUSE or time.time() >= deadline:
+                    if port and wait:
+                        raise OSError(e.errno, '%s:%d is still in use after %gs - another '
+                                               'process holds it, or the last one has not let go'
+                                      % (host, port, wait))
+                    raise
+                time.sleep(0.2)
+
+    @property
+    def url(self) -> str:
+        return 'ws://%s:%d' % (self.host, self.port)
+
+    def start(self) -> Server:
+        self._thread.start()
+        return self
+
+    def _loop(self) -> None:
+        try:
+            self.sock.settimeout(0.2)
+        except OSError:
+            # Stopped before this thread was scheduled, which a test that
+            # starts and stops a face in the same breath does routinely. The
+            # accept below already treats a closed socket as the end; this line
+            # ran ahead of that guard and died with a traceback on stderr
+            # instead - noise in a log an operator reads to find real faults.
+            return
+        while not self._stop.is_set():
+            try:
+                conn, _ = self.sock.accept()
+            except socket.timeout:
+                continue
+            except OSError:
+                break
+            threading.Thread(target=self._serve, args=(conn,), daemon=True).start()
+
+    def _serve(self, conn: socket.socket) -> None:
+        try:
+            conn.settimeout(60)
+            ws, path = accept(conn)
+        except PlainHTTP as req:
+            try:
+                if self.http_handler:
+                    try:
+                        self.http_handler(conn, req)
+                    except TypeError:
+                        self.http_handler(conn, req.path, req.headers)
+                else:
+                    http_response(conn, '404 Not Found', b'{}')
+            finally:
+                try:
+                    conn.close()
+                except OSError:
+                    pass
+            return
+        except (ConnectionError, OSError):
+            conn.close()
+            return
+        with self._conn_lock:
+            self.connections.append(ws)
+        try:
+            self.handler(ws, path)
+        except (ConnectionError, OSError, ValueError):
+            pass
+        except Exception as e:          # a handler bug closes one connection, not the server
+            with self._conn_lock:
+                self.errors.append('%s: %s' % (type(e).__name__, e))
+                del self.errors[:-200]
+        finally:
+            with self._conn_lock:
+                if ws in self.connections:
+                    self.connections.remove(ws)
+            ws.close()
+
+    def stop(self) -> None:
+        self._stop.set()
+        try:
+            self.sock.close()
+        except OSError:
+            pass
+        with self._conn_lock:
+            conns = list(self.connections)
+            self.connections.clear()
+        for ws in conns:
+            ws.close()