Complete owned WebSocket shutdown and retain partial received bytes

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)

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)

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