diff --git a/src/bitworld/native_websocket.nim b/src/bitworld/native_websocket.nim index 16d06cf..71df672 100644 --- a/src/bitworld/native_websocket.nim +++ b/src/bitworld/native_websocket.nim @@ -19,6 +19,7 @@ type ipPlayer, ipCleanup PongState = enum psNone, psWaiting, psQueued, psSent + ReceiveObserver* = proc(data: string) {.closure.} NativeWebSocket* = ref object handle: PCurl socket: cint @@ -31,7 +32,9 @@ type pongState: Atomic[PongState] pongData: string maximum: int - closed: bool + closed, peerClosed: bool + received: Option[ReceiveObserver] + closePayload: string WebSocketConnection* = object kind*: WebSocketKind socket*: NativeWebSocket @@ -115,6 +118,10 @@ proc receiveBytes(ws: NativeWebSocket, deadline: MonoTime, release(ws.ioLock) if code == 0: if count == 0: return WebSocketResult(kind: wsClosed) + if ws.received.isSome: + var bytes = newString(int(count)) + copyMem(bytes[0].addr, buffer.addr, int(count)) + ws.received.get()(bytes) let offset = ws.incoming.len ws.incoming.setLen(offset + int(count)) copyMem(ws.incoming[offset].addr, buffer.addr, int(count)) @@ -262,6 +269,8 @@ proc receiveMessage(ws: NativeWebSocket, deadline: MonoTime, if code < 1000 or code >= 5000 or code in [1004, 1005, 1006, 1015] or (code >= 1016 and code < 3000) or validateUtf8(payload[2 .. ^1]) != -1: return WebSocketResult(kind: wsFailure, error: "Invalid WebSocket close reason") + ws.closePayload = payload + ws.peerClosed = true return WebSocketResult(kind: wsClosed) of 9: ws.pongData = payload @@ -307,8 +316,30 @@ proc closeNativeWebSocket*(ws: NativeWebSocket) = deinitLock(ws.ioLock) deinitLock(ws.sendLock) +proc finishNativeWebSocket*(ws: NativeWebSocket, deadline: MonoTime): WebSocketResult = + ## Call after joining all IO owners. Complete the close handshake within the + ## caller's existing cleanup deadline, then release the socket on every path. + try: + let payload = if ws.peerClosed: ws.closePayload else: "\x03\xe8" + result = ws.sendFrame(8, payload, deadline, ipCleanup) + if result.kind != wsReady: return + if ws.peerClosed: return + while true: + result = ws.receiveMessage(deadline, ipCleanup, true) + case result.kind + of wsMessage: discard + of wsClosed: + return WebSocketResult(kind: (if ws.peerClosed: wsReady else: wsFailure), + error: (if ws.peerClosed: "" else: "Peer ended without a WebSocket close frame")) + else: return + finally: + ws.closeNativeWebSocket() + proc connectNativeWebSocket*(url: string, deadline: MonoTime, - maxMessageBytes: Positive): WebSocketConnection = + maxMessageBytes: Positive, + received: Option[ReceiveObserver] = none(ReceiveObserver)): WebSocketConnection = + ## The receive owner invokes the observer before buffering each raw chunk, + ## including upgrade bytes and incomplete frames. Observers must not reenter IO. if interruptionRequested(): return WebSocketConnection(kind: wsInterrupted) if getMonoTime() >= deadline: return WebSocketConnection(kind: wsDeadline) let parsed = parseUri(url) @@ -318,7 +349,7 @@ proc connectNativeWebSocket*(url: string, deadline: MonoTime, return WebSocketConnection(kind: wsFailure, error: "Invalid native WebSocket URL") var transport = parsed transport.scheme = if parsed.scheme == "wss": "https" else: "http" - let ws = NativeWebSocket(handle: easy_init(), maximum: maxMessageBytes) + let ws = NativeWebSocket(handle: easy_init(), maximum: maxMessageBytes, received: received) doAssert ws.handle != nil initLock(ws.ioLock) initLock(ws.sendLock) diff --git a/tests/support/native_websocket_probe.nim b/tests/support/native_websocket_probe.nim index 7942109..c9c7037 100644 --- a/tests/support/native_websocket_probe.nim +++ b/tests/support/native_websocket_probe.nim @@ -12,16 +12,24 @@ proc main() = installNativeStopHandlers() let mode = paramStr(2) let started = getMonoTime() - let connected = connectNativeWebSocket(paramStr(1), started + initDuration(milliseconds = 600), 16 * 1024 * 1024) + var observed = "" + let observer: ReceiveObserver = proc(data: string) = observed.add(data) + let connected = connectNativeWebSocket(paramStr(1), started + initDuration(milliseconds = 600), 16 * 1024 * 1024, some(observer)) echo $(%*{"connect": $connected.kind, "error": connected.error, "curl_version": $libcurl.version(), "elapsed_ms": (getMonoTime() - started).inMilliseconds}) if connected.kind != wsReady: quit(0) let socket = connected.socket - defer: closeNativeWebSocket(socket) + defer: + if mode notin ["close", "close-stop", "close-eof", "close-stall"]: closeNativeWebSocket(socket) proc report(value: WebSocketResult) = echo $(%*{"kind": $value.kind, "data": value.data, "elapsed_ms": (getMonoTime() - started).inMilliseconds}) - if mode == "mixed-binary": + if mode in ["close", "close-stop", "close-eof", "close-stall"]: + if mode == "close-stop": + report(socket.receiveNativeMessage(getMonoTime() + initDuration(seconds = 5))) + report(socket.finishNativeWebSocket(getMonoTime() + initDuration(milliseconds = 600))) + echo $(%*{"received_b64": encode(observed)}) + elif mode == "mixed-binary": let deadline = getMonoTime() + initDuration(milliseconds = 600) for index in 0 ..< 3: let received = socket.receiveNativeMessage(deadline) diff --git a/tests/test_native_websocket.py b/tests/test_native_websocket.py index ef627d4..d850cdb 100644 --- a/tests/test_native_websocket.py +++ b/tests/test_native_websocket.py @@ -9,7 +9,6 @@ import socket import ssl import tempfile -import unittest import struct import subprocess import sys @@ -73,7 +72,8 @@ def serve(): with connection: connection.settimeout(5) if rejected_tls: - assert connection.recv(1) == b'' + # The client result below verifies rejection. TLS backends may + # signal that rejection with either EOF or a TCP reset. return upgrade(connection, valid, stall) operation(connection) @@ -94,9 +94,9 @@ def serve(): output, error = child.stdout.read(), child.stderr.read() assert child.returncode == 0, (name, error) rows = [first] + [json.loads(line) for line in output.splitlines()] - if rejected_tls and not trust: - with unittest.TestCase().assertRaises(ssl.SSLError): - future.result(timeout=5) + if rejected_tls: + error = future.exception(timeout=5) + assert error is None or isinstance(error, ssl.SSLError), error else: future.result(timeout=5) print(json.dumps({'case': name, 'rows': rows}), flush=True) @@ -240,3 +240,26 @@ def stopped_binary(connection): assert rows[1]['kind'] == 'wsInterrupted' assert base64.b64decode(rows[2]['body_b64'], validate=True) == b'\x00\xff\xe2\x82' print('native websocket transport gates passed') + +# Real peer acknowledgement is required; EOF and a silent peer cannot qualify. +def finish_close(connection): + assert client_frame(connection) == (8, struct.pack('!H', 1000)) + connection.sendall(frame(struct.pack('!H', 1000), opcode=8)) +rows = run_case('owned-close-handshake', 'close', finish_close) +assert rows[1]['kind'] == 'wsReady' +assert base64.b64decode(rows[2]['received_b64']).endswith(frame(struct.pack('!H', 1000), opcode=8)) + +def stop_partial_close(connection): + partial = frame(b'partial-before-stop') + connection.sendall(partial[:4]) + assert client_frame(connection) == (8, struct.pack('!H', 1000)) + connection.sendall(partial[4:] + frame(struct.pack('!H', 1000), opcode=8)) +for sig in (signal.SIGTERM, signal.SIGINT): + rows = run_case('close-after-' + sig.name, 'close-stop', stop_partial_close, sig) + assert rows[1]['kind'] == 'wsInterrupted' and rows[2]['kind'] == 'wsReady' + wire = base64.b64decode(rows[3]['received_b64']) + assert wire.split(b'\r\n\r\n', 1)[1] == frame(b'partial-before-stop') + frame(struct.pack('!H', 1000), opcode=8) +rows = run_case('close-without-peer-frame', 'close-eof', lambda connection: client_frame(connection)) +assert rows[1]['kind'] == 'wsFailure' +rows = run_case('close-stalled-peer', 'close-stall', lambda connection: (client_frame(connection), time.sleep(.8))) +assert rows[1]['kind'] == 'wsDeadline' and rows[1]['elapsed_ms'] < 1000