feat: Phase 4 — WebSocket real-time subscriptions

- Full WebSocket server (RFC 6455) with frame encode/decode
- HTTP upgrade handshake with SHA1 accept key
- SUBSCRIBE/UNSUBSCRIBE table messages
- Per-client subscription tracking via HashSet
- Broadcast to subscribers on data changes
- Ping/pong keepalive
- CORS header on upgrade response
- All 216 tests pass
This commit is contained in:
2026-05-06 11:35:35 +03:00
parent 775703e008
commit 6e6ed73cfd
+299
View File
@@ -0,0 +1,299 @@
## BaraDB WebSocket Server — real-time subscriptions
import std/asyncdispatch
import std/asyncnet
import std/strutils
import std/tables
import std/base64
import std/sets
type
WsFrame = object
fin: bool
opcode: uint8
masked: bool
payloadLen: uint64
maskKey: array[4, byte]
payload: string
WsClient* = ref object
socket: AsyncSocket
id: int
subscriptions: HashSet[string]
WsServer* = ref object
clients*: Table[int, WsClient]
nextId: int
running: bool
onInsert*: proc (table, key, value: string) {.closure.}
onDelete*: proc (table, key: string) {.closure.}
proc newWsServer*(): WsServer =
WsServer(clients: initTable[int, WsClient](), nextId: 1, running: false)
# ----------------------------------------------------------------------
# WebSocket frame encoding/decoding (RFC 6455)
# ----------------------------------------------------------------------
proc encodeFrame(opcode: uint8, payload: string): string =
result = ""
let isMasked = false
var b0 = 0x80'u8 or opcode
result.add(char(b0))
var b1 = 0'u8
if not isMasked:
if payload.len < 126:
b1 = uint8(payload.len)
elif payload.len <= 65535:
b1 = 126
else:
b1 = 127
result.add(char(b1))
if payload.len >= 126 and payload.len <= 65535:
var len16 = uint16(payload.len)
result.add(char((len16 shr 8) and 0xFF))
result.add(char(len16 and 0xFF))
elif payload.len > 65535:
var len64 = uint64(payload.len)
for i in countdown(7, 0):
result.add(char((len64 shr (i * 8)) and 0xFF))
result.add(payload)
proc decodeFrame(data: string): (WsFrame, int) =
if data.len < 2:
return (WsFrame(), 0)
var frame = WsFrame()
let b0 = uint8(data[0])
let b1 = uint8(data[1])
frame.fin = (b0 and 0x80) != 0
frame.opcode = b0 and 0x0F
frame.masked = (b1 and 0x80) != 0
var len = uint64(b1 and 0x7F)
var offset = 2
if len == 126:
if data.len < 4: return (WsFrame(), 0)
len = (uint64(uint8(data[2])) shl 8) or uint64(uint8(data[3]))
offset = 4
elif len == 127:
if data.len < 10: return (WsFrame(), 0)
len = 0
for i in 0..7:
len = (len shl 8) or uint64(uint8(data[2 + i]))
offset = 10
if frame.masked:
if data.len < offset + 4: return (WsFrame(), 0)
for i in 0..3:
frame.maskKey[i] = byte(data[offset + i])
offset += 4
if uint64(data.len) < uint64(offset) + len:
return (Wsframe(), 0)
let plen = int(len)
if frame.masked:
for i in 0..<plen:
frame.payload.add(char(byte(data[offset + i]) xor frame.maskKey[i mod 4]))
else:
frame.payload = data[offset..offset + plen - 1]
return (frame, offset + plen)
# ----------------------------------------------------------------------
# HTTP upgrade handshake
# ----------------------------------------------------------------------
proc computeAcceptKey(key: string): string =
let magic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
# Minimal SHA1 implementation for WebSocket handshake
proc rotateLeft(x: uint32, n: int): uint32 {.inline.} =
(x shl uint32(n)) or (x shr uint32(32 - n))
var h: array[5, uint32] = [0x67452301'u32, 0xEFCDAB89'u32, 0x98BADCFE'u32, 0x10325476'u32, 0xC3D2E1F0'u32]
let input = key & magic
var msg = newSeq[byte](input.len)
for i, c in input:
msg[i] = byte(c)
let origLen = msg.len
msg.add(0x80'u8)
while (msg.len + 8) mod 64 != 0:
msg.add(0'u8)
let bitLen = uint64(origLen) * 8
for i in countdown(7, 0):
msg.add(byte((bitLen shr (i * 8)) and 0xFF))
var w: array[80, uint32]
for chunkStart in countup(0, msg.len - 64, 64):
for i in 0..15:
w[i] = (uint32(msg[chunkStart + i*4]) shl 24) or
(uint32(msg[chunkStart + i*4 + 1]) shl 16) or
(uint32(msg[chunkStart + i*4 + 2]) shl 8) or
uint32(msg[chunkStart + i*4 + 3])
for i in 16..79:
w[i] = rotateLeft(w[i-3] xor w[i-8] xor w[i-14] xor w[i-16], 1)
var a = h[0]; var b = h[1]; var c = h[2]; var d = h[3]; var e = h[4]
for i in 0..79:
var f, k: uint32
if i < 20:
f = (b and c) or ((not b) and d); k = 0x5A827999'u32
elif i < 40:
f = b xor c xor d; k = 0x6ED9EBA1'u32
elif i < 60:
f = (b and c) or (b and d) or (c and d); k = 0x8F1BBCDC'u32
else:
f = b xor c xor d; k = 0xCA62C1D6'u32
let temp = rotateLeft(a, 5) + f + e + k + w[i]
e = d; d = c; c = rotateLeft(b, 30); b = a; a = temp
h[0] = h[0] + a; h[1] = h[1] + b; h[2] = h[2] + c; h[3] = h[3] + d; h[4] = h[4] + e
var digest = newString(20)
for i in 0..4:
digest[i*4] = char((h[i] shr 24) and 0xFF)
digest[i*4 + 1] = char((h[i] shr 16) and 0xFF)
digest[i*4 + 2] = char((h[i] shr 8) and 0xFF)
digest[i*4 + 3] = char(h[i] and 0xFF)
return encode(digest)
# ----------------------------------------------------------------------
# Subscription management
# ----------------------------------------------------------------------
proc subscribe*(client: WsClient, table: string) =
client.subscriptions.incl(table)
proc unsubscribe*(client: WsClient, table: string) =
client.subscriptions.excl(table)
proc notifyClient(client: WsClient, msg: string) {.async.} =
try:
let frame = encodeFrame(0x1, msg)
await client.socket.send(frame)
except:
discard
proc broadcastToTable*(server: WsServer, table: string, msg: string) {.async.} =
for id, client in server.clients:
if table in client.subscriptions:
asyncCheck client.notifyClient(msg)
# ----------------------------------------------------------------------
# WebSocket client handler
# ----------------------------------------------------------------------
proc handleWsClient(server: WsServer, client: AsyncSocket, id: int) {.async.} =
echo "WebSocket client ", id, " connected"
var wsClient = WsClient(socket: client, id: id, subscriptions: initHashSet[string]())
server.clients[id] = wsClient
var buf = ""
try:
while true:
let chunk = await client.recv(4096)
if chunk.len == 0:
break
buf.add(chunk)
while buf.len >= 2:
let (frame, consumed) = decodeFrame(buf)
if consumed == 0:
break
case frame.opcode
of 0x8: # close
client.close()
server.clients.del(id)
return
of 0x9: # ping
let pong = encodeFrame(0xA, frame.payload)
await client.send(pong)
of 0x1: # text
let msg = frame.payload
if msg.startsWith("SUBSCRIBE "):
let table = msg[10..^1].strip()
wsClient.subscribe(table)
let ack = encodeFrame(0x1, "OK subscribed to " & table)
await client.send(ack)
elif msg.startsWith("UNSUBSCRIBE "):
let table = msg[12..^1].strip()
wsClient.unsubscribe(table)
let ack = encodeFrame(0x1, "OK unsubscribed from " & table)
await client.send(ack)
else:
let echo = encodeFrame(0x1, "ECHO: " & msg)
await client.send(echo)
else:
discard
buf = buf[consumed..^1]
except:
discard
finally:
echo "WebSocket client ", id, " disconnected"
server.clients.del(id)
client.close()
# ----------------------------------------------------------------------
# HTTP upgrade + WebSocket handoff
# ----------------------------------------------------------------------
proc handleConnection(server: WsServer, client: AsyncSocket) {.async.} =
let firstLine = await client.recvLine()
if firstLine.len == 0:
client.close()
return
var headers = initTable[string, string]()
var wsKey = ""
while true:
let line = await client.recvLine()
if line == "\r" or line == "":
break
let parts = line.split(":", maxSplit = 1)
if parts.len >= 2:
let key = parts[0].strip().toLower()
let val = parts[1].strip()
headers[key] = val
if key == "sec-websocket-key":
wsKey = val
if wsKey.len == 0:
await client.send("HTTP/1.1 400 Bad Request\r\n\r\n")
client.close()
return
let acceptKey = computeAcceptKey(wsKey)
var response = "HTTP/1.1 101 Switching Protocols\r\n"
response &= "Upgrade: websocket\r\n"
response &= "Connection: Upgrade\r\n"
response &= "Sec-WebSocket-Accept: " & acceptKey & "\r\n"
response &= "Access-Control-Allow-Origin: *\r\n"
response &= "\r\n"
await client.send(response)
inc server.nextId
asyncCheck server.handleWsClient(client, server.nextId)
proc run*(server: WsServer, port: int = 8081) {.async.} =
server.running = true
let sock = newAsyncSocket()
sock.setSockOpt(OptReuseAddr, true)
sock.bindAddr(Port(port))
sock.listen()
echo "BaraDB WebSocket listening on port ", port
while server.running:
let client = await sock.accept()
asyncCheck server.handleConnection(client)
proc stop*(server: WsServer) =
server.running = false