diff --git a/src/barabadb/core/raft.nim b/src/barabadb/core/raft.nim index 55ae65e..e7f6ff3 100644 --- a/src/barabadb/core/raft.nim +++ b/src/barabadb/core/raft.nim @@ -605,15 +605,26 @@ proc processMessage*(net: RaftNetwork, msg: RaftMessage) {.async.} = of rmkAppendEntriesReply: net.node.handleAppendReply(msg.senderId, msg) +proc recvExact*(client: AsyncSocket, size: int): Future[string] {.async.} = + ## Reads exactly `size` bytes from `client`. A short return means the peer + ## disconnected mid-frame (EOF); callers must treat it as end of stream. + var buf = "" + while buf.len < size: + let chunk = await client.recv(size - buf.len) + if chunk.len == 0: + break + buf.add(chunk) + return buf + proc receiveLoop(net: RaftNetwork, client: AsyncSocket) {.async.} = try: while net.running: - let lenData = await client.recv(4) + let lenData = await recvExact(client, 4) if lenData.len < 4: break var pos = 0 let payloadLen = int(readUint32(cast[seq[byte]](lenData), pos)) - let payloadStr = await client.recv(payloadLen) + let payloadStr = await recvExact(client, payloadLen) if payloadStr.len < payloadLen: break var payload = newSeq[byte](payloadLen) diff --git a/tests/test_all.nim b/tests/test_all.nim index f398247..d55b290 100644 --- a/tests/test_all.nim +++ b/tests/test_all.nim @@ -4,6 +4,7 @@ import std/tables import std/strutils import std/os import std/asyncdispatch +import std/asyncnet import std/monotimes import std/base64 import std/json @@ -2459,6 +2460,153 @@ suite "Raft Network Transport": waitFor netA.processMessage(staleMsg) check timer.checkTimeout() + test "recvExact reassembles fragmented frames": + const port = 29025 + let voteReq = RaftMessage(kind: rmkRequestVote, term: 7, senderId: "peer-x", + lastLogIndex: 3, lastLogTerm: 2) + let data = serialize(voteReq) + var frame = newSeq[byte](4 + data.len) + frame[0] = byte(data.len shr 24) + frame[1] = byte(data.len shr 16) + frame[2] = byte(data.len shr 8) + frame[3] = byte(data.len) + for i in 0 ..< data.len: + frame[4 + i] = data[i] + + var hdrLen = -1 + var trailingEofLen = -1 + var midFrameEofLen = -1 + var gotMsg: RaftMessage + proc reader() {.async.} = + let l = newAsyncSocket() + l.setSockOpt(OptReuseAddr, true) + l.bindAddr(Port(port)) + l.listen() + # Connection 1: one full fragmented frame, then EOF. + let c1 = await l.accept() + let hdr = await recvExact(c1, 4) + hdrLen = hdr.len + let payloadLen = (int(uint8(hdr[0])) shl 24) or (int(uint8(hdr[1])) shl 16) or + (int(uint8(hdr[2])) shl 8) or int(uint8(hdr[3])) + let payloadStr = await recvExact(c1, payloadLen) + let trailing = await recvExact(c1, 1) # peer closed -> EOF + trailingEofLen = trailing.len + var payload = newSeq[byte](payloadLen) + for i in 0 ..< payloadLen: + payload[i] = byte(payloadStr[i]) + gotMsg = deserializeRaftMessage(payload) + c1.close() + # Connection 2: EOF in the middle of the 4-byte header. + let c2 = await l.accept() + let shortHdr = await recvExact(c2, 4) + midFrameEofLen = shortHdr.len + c2.close() + l.close() + asyncCheck reader() + waitFor sleepAsync(50) + + let s1 = newAsyncSocket() + waitFor s1.connect("127.0.0.1", Port(port)) + waitFor s1.send(cast[string](frame[0 ..< 1])) + waitFor sleepAsync(20) + waitFor s1.send(cast[string](frame[1 ..< 4 + data.len div 2])) + waitFor sleepAsync(20) + waitFor s1.send(cast[string](frame[4 + data.len div 2 .. ^1])) + s1.close() + + waitFor sleepAsync(50) + let s2 = newAsyncSocket() + waitFor s2.connect("127.0.0.1", Port(port)) + waitFor s2.send(cast[string](frame[0 ..< 2])) + s2.close() + + var waited = 0 + while midFrameEofLen < 0 and waited < 2000: + waitFor sleepAsync(50) + waited += 50 + + check hdrLen == 4 + check gotMsg.kind == rmkRequestVote + check gotMsg.term == 7 + check gotMsg.senderId == "peer-x" + check gotMsg.lastLogIndex == 3 + check gotMsg.lastLogTerm == 2 + check trailingEofLen == 0 + check midFrameEofLen == 2 + + test "framing reassembles chunked messages": + const srvPort = 29021 + const cliPort = 29023 + var srv = newRaftNode("srv", @["cli"], raftPort = srvPort) + srv.electionTimeout = 60000 # keep the server passive during the test + srv.peerAddrs["cli"] = ("127.0.0.1", cliPort) + let net = newRaftNetwork(srv) + asyncCheck net.run() + waitFor sleepAsync(50) + + # Test-side listener stands in for the "cli" peer and captures the reply. + let replyListener = newAsyncSocket() + replyListener.setSockOpt(OptReuseAddr, true) + replyListener.bindAddr(Port(cliPort)) + replyListener.listen() + + var replyMsg: RaftMessage + var gotReply = false + proc acceptor() {.async.} = + let conn = await replyListener.accept() + var hdr = "" + while hdr.len < 4: + let c = await conn.recv(4 - hdr.len) + if c.len == 0: return + hdr.add(c) + let payloadLen = (int(uint8(hdr[0])) shl 24) or (int(uint8(hdr[1])) shl 16) or + (int(uint8(hdr[2])) shl 8) or int(uint8(hdr[3])) + var body = "" + while body.len < payloadLen: + let c = await conn.recv(payloadLen - body.len) + if c.len == 0: return + body.add(c) + var payload = newSeq[byte](payloadLen) + for i in 0 ..< payloadLen: + payload[i] = byte(body[i]) + replyMsg = deserializeRaftMessage(payload) + gotReply = true + asyncCheck acceptor() + + let voteReq = RaftMessage(kind: rmkRequestVote, term: 1, senderId: "cli") + let data = serialize(voteReq) + var frame = newSeq[byte](4 + data.len) + frame[0] = byte(data.len shr 24) + frame[1] = byte(data.len shr 16) + frame[2] = byte(data.len shr 8) + frame[3] = byte(data.len) + for i in 0 ..< data.len: + frame[4 + i] = data[i] + + let client = newAsyncSocket() + waitFor client.connect("127.0.0.1", Port(srvPort)) + # Fragment the frame mid-header and mid-payload with real delays. + waitFor client.send(cast[string](frame[0 ..< 2])) + waitFor sleepAsync(30) + waitFor client.send(cast[string](frame[2 ..< 4 + data.len div 2])) + waitFor sleepAsync(30) + waitFor client.send(cast[string](frame[4 + data.len div 2 .. ^1])) + + var waited = 0 + while not gotReply and waited < 2000: + waitFor sleepAsync(50) + waited += 50 + + client.close() + net.stop() + replyListener.close() + waitFor sleepAsync(50) + + check gotReply + if gotReply: + check replyMsg.kind == rmkRequestVoteReply + check replyMsg.success + suite "CLI Autocomplete": test "Autocomplete commands": let res = autocomplete("he")