clawft_kernel/
mesh_tcp.rs

1//! TCP transport backend for mesh networking (K6).
2//!
3//! Uses tokio TCP for node-to-node communication. Production Cloud/Edge
4//! deployments prefer QUIC (`mesh_quic` / ADR-026) with optional Noise
5//! (`mesh_noise` / snow). Plain TCP remains the default and is useful
6//! when UDP is blocked.
7//!
8//! # Cancel-safety (ADR-010 / WEFT-18)
9//!
10//! [`TcpMeshStream::recv`] keeps partial length/body progress on `self` so
11//! it is safe to race inside `tokio::select!` (mesh peer loop in `boot.rs`).
12//! A naive `read_exact(len)` then `read_exact(body)` pair is **not**
13//! cancel-safe: dropping mid-body desyncs the next length prefix.
14
15use std::net::SocketAddr;
16
17use async_trait::async_trait;
18use tokio::io::{AsyncReadExt, AsyncWriteExt};
19use tokio::net::{TcpListener, TcpStream};
20
21use crate::mesh::{MAX_MESSAGE_SIZE, MeshError, MeshStream, MeshTransport, TransportListener};
22
23// ── TcpMeshStream ───────────────────────────────────────────────
24
25/// A [`MeshStream`] backed by a single TCP connection.
26///
27/// Messages are length-prefixed on the wire (4-byte big-endian length
28/// followed by that many bytes of payload).
29///
30/// `recv` is **cancel-safe** (WEFT-18): partial prefix/body progress is
31/// retained on `self` across cancellations so `tokio::select!` peer loops
32/// do not desync the frame stream.
33pub struct TcpMeshStream {
34    stream: TcpStream,
35    remote: SocketAddr,
36    /// Staging for the 4-byte length prefix (always partially or fully filled).
37    len_buf: [u8; 4],
38    /// How many of `len_buf` bytes are valid.
39    len_filled: usize,
40    /// In-progress body once the length prefix is complete (`None` = still
41    /// reading the length).
42    body: Option<Vec<u8>>,
43    /// Bytes of `body` filled so far.
44    body_filled: usize,
45}
46
47impl TcpMeshStream {
48    fn new(stream: TcpStream, remote: SocketAddr) -> Self {
49        Self {
50            stream,
51            remote,
52            len_buf: [0u8; 4],
53            len_filled: 0,
54            body: None,
55            body_filled: 0,
56        }
57    }
58
59    /// Reset framing state after a complete message (or hard error).
60    fn reset_frame(&mut self) {
61        self.len_filled = 0;
62        self.body = None;
63        self.body_filled = 0;
64    }
65}
66
67#[async_trait]
68impl MeshStream for TcpMeshStream {
69    async fn send(&mut self, data: &[u8]) -> Result<(), MeshError> {
70        let len = (data.len() as u32).to_be_bytes();
71        self.stream
72            .write_all(&len)
73            .await
74            .map_err(|e| MeshError::Io(e.to_string()))?;
75        self.stream
76            .write_all(data)
77            .await
78            .map_err(|e| MeshError::Io(e.to_string()))?;
79        self.stream
80            .flush()
81            .await
82            .map_err(|e| MeshError::Io(e.to_string()))?;
83        Ok(())
84    }
85
86    async fn recv(&mut self) -> Result<Vec<u8>, MeshError> {
87        // Cancel-safe framed read (ADR-010 / WEFT-18).
88        //
89        // Progress counters and buffers live on `self`. Each successful
90        // `read` updates those fields before the next `.await`. Dropping
91        // this future while Pending does not discard already-consumed
92        // bytes, so a later `recv` resumes the same frame.
93        //
94        // Field destructuring lets us mutably borrow the socket and the
95        // staging buffers together without fighting the borrow checker
96        // across await points.
97
98        // ── length prefix ──────────────────────────────────────────
99        while self.len_filled < 4 {
100            let Self {
101                stream,
102                len_buf,
103                len_filled,
104                ..
105            } = self;
106            let n = stream
107                .read(&mut len_buf[*len_filled..])
108                .await
109                .map_err(|e| MeshError::Io(e.to_string()))?;
110            if n == 0 {
111                return Err(MeshError::ConnectionClosed);
112            }
113            *len_filled += n;
114        }
115
116        // ── allocate body once length is known ─────────────────────
117        if self.body.is_none() {
118            let len = u32::from_be_bytes(self.len_buf) as usize;
119            if len > MAX_MESSAGE_SIZE {
120                self.reset_frame();
121                return Err(MeshError::MessageTooLarge {
122                    size: len,
123                    max: MAX_MESSAGE_SIZE,
124                });
125            }
126            self.body = Some(vec![0u8; len]);
127            self.body_filled = 0;
128        }
129
130        // ── body ───────────────────────────────────────────────────
131        let expected = self.body.as_ref().map(|b| b.len()).unwrap_or(0);
132        while self.body_filled < expected {
133            let Self {
134                stream,
135                body,
136                body_filled,
137                ..
138            } = self;
139            let buf = body
140                .as_mut()
141                .expect("body allocated before body-fill loop");
142            let n = stream
143                .read(&mut buf[*body_filled..])
144                .await
145                .map_err(|e| MeshError::Io(e.to_string()))?;
146            if n == 0 {
147                return Err(MeshError::ConnectionClosed);
148            }
149            *body_filled += n;
150        }
151
152        let out = self.body.take().expect("body complete");
153        self.reset_frame();
154        Ok(out)
155    }
156
157    async fn close(&mut self) -> Result<(), MeshError> {
158        self.stream
159            .shutdown()
160            .await
161            .map_err(|e| MeshError::Io(e.to_string()))
162    }
163
164    fn remote_addr(&self) -> Option<SocketAddr> {
165        Some(self.remote)
166    }
167}
168
169// ── TcpTransportListener ────────────────────────────────────────
170
171/// A [`TransportListener`] backed by a tokio [`TcpListener`].
172pub struct TcpTransportListener {
173    listener: TcpListener,
174}
175
176#[async_trait]
177impl TransportListener for TcpTransportListener {
178    async fn accept(&mut self) -> Result<(Box<dyn MeshStream>, SocketAddr), MeshError> {
179        let (stream, addr) = self
180            .listener
181            .accept()
182            .await
183            .map_err(|e| MeshError::Io(e.to_string()))?;
184        Ok((Box::new(TcpMeshStream::new(stream, addr)), addr))
185    }
186
187    fn local_addr(&self) -> Result<SocketAddr, MeshError> {
188        self.listener
189            .local_addr()
190            .map_err(|e| MeshError::Io(e.to_string()))
191    }
192}
193
194// ── TcpTransport ────────────────────────────────────────────────
195
196/// TCP mesh transport.
197///
198/// Implements [`MeshTransport`] using plain TCP.  Addresses may be
199/// bare `ip:port` or prefixed with `tcp://`.
200pub struct TcpTransport;
201
202#[async_trait]
203impl MeshTransport for TcpTransport {
204    fn name(&self) -> &str {
205        "tcp"
206    }
207
208    async fn listen(&self, addr: &str) -> Result<Box<dyn TransportListener>, MeshError> {
209        let bind_addr = addr.strip_prefix("tcp://").unwrap_or(addr);
210        let listener = TcpListener::bind(bind_addr)
211            .await
212            .map_err(|e| MeshError::Io(e.to_string()))?;
213        Ok(Box::new(TcpTransportListener { listener }))
214    }
215
216    async fn connect(&self, addr: &str) -> Result<Box<dyn MeshStream>, MeshError> {
217        let connect_addr = addr.strip_prefix("tcp://").unwrap_or(addr);
218        let stream = TcpStream::connect(connect_addr)
219            .await
220            .map_err(|e| MeshError::Io(e.to_string()))?;
221        let remote = stream
222            .peer_addr()
223            .map_err(|e| MeshError::Io(e.to_string()))?;
224        Ok(Box::new(TcpMeshStream::new(stream, remote)))
225    }
226
227    fn supports(&self, addr: &str) -> bool {
228        // Supports bare IP:port or tcp:// scheme.
229        !addr.contains("://") || addr.starts_with("tcp://")
230    }
231}
232
233// ── Tests ───────────────────────────────────────────────────────
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    #[tokio::test]
240    async fn tcp_transport_connect_send_recv() {
241        let transport = TcpTransport;
242
243        // Bind to an OS-assigned port.
244        let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
245        let addr = listener.local_addr().unwrap();
246
247        // "Client" node sends, then reads response.
248        let send_task = tokio::spawn(async move {
249            let mut stream = TcpTransport.connect(&addr.to_string()).await.unwrap();
250            stream.send(b"hello from node A").await.unwrap();
251            let response = stream.recv().await.unwrap();
252            assert_eq!(response, b"hello from node B");
253            stream.close().await.unwrap();
254        });
255
256        // "Server" node accepts, reads, responds.
257        let (mut server_stream, _client_addr) = listener.accept().await.unwrap();
258        let received = server_stream.recv().await.unwrap();
259        assert_eq!(received, b"hello from node A");
260        server_stream.send(b"hello from node B").await.unwrap();
261
262        send_task.await.unwrap();
263    }
264
265    #[tokio::test]
266    async fn two_node_cluster_forms() {
267        use crate::mesh::WeftHandshake;
268        use crate::mesh_framing::{FrameType, MeshFrame};
269
270        let transport = TcpTransport;
271        let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
272        let addr = listener.local_addr().unwrap();
273
274        // Node A connects to Node B and sends a handshake.
275        let node_a = tokio::spawn(async move {
276            let mut stream = TcpTransport.connect(&addr.to_string()).await.unwrap();
277
278            let handshake = WeftHandshake {
279                node_id: "node-a".into(),
280                governance_genesis_hash: [0u8; 32],
281                governance_version: "1.0".into(),
282                capabilities: 0xFF,
283                kem_supported: false,
284                chain_seq: 0,
285                supported_sync_streams: vec![1, 2, 3],
286            };
287            let payload = serde_json::to_vec(&handshake).unwrap();
288            let frame = MeshFrame {
289                frame_type: FrameType::Handshake,
290                payload,
291            };
292            let encoded = frame.encode().unwrap();
293            // Send the wire bytes *after* the 4-byte length prefix through
294            // the transport layer (which adds its own length prefix).
295            stream.send(&encoded).await.unwrap();
296
297            // Receive handshake response.
298            let data = stream.recv().await.unwrap();
299            // The transport-level length prefix is consumed by recv();
300            // the returned bytes are a full encoded MeshFrame (length + type + payload).
301            let len = u32::from_be_bytes([data[0], data[1], data[2], data[3]]) as usize;
302            let response = MeshFrame::decode(&data[4..4 + len]).unwrap();
303            assert_eq!(response.frame_type, FrameType::Handshake);
304            let remote: WeftHandshake = serde_json::from_slice(&response.payload).unwrap();
305            assert_eq!(remote.node_id, "node-b");
306
307            stream.close().await.unwrap();
308        });
309
310        // Node B accepts and exchanges handshake.
311        let (mut stream, _) = listener.accept().await.unwrap();
312        let data = stream.recv().await.unwrap();
313        let len = u32::from_be_bytes([data[0], data[1], data[2], data[3]]) as usize;
314        let frame = MeshFrame::decode(&data[4..4 + len]).unwrap();
315        assert_eq!(frame.frame_type, FrameType::Handshake);
316        let remote: WeftHandshake = serde_json::from_slice(&frame.payload).unwrap();
317        assert_eq!(remote.node_id, "node-a");
318
319        // Send our handshake back.
320        let handshake = WeftHandshake {
321            node_id: "node-b".into(),
322            governance_genesis_hash: [0u8; 32],
323            governance_version: "1.0".into(),
324            capabilities: 0xFF,
325            kem_supported: false,
326            chain_seq: 0,
327            supported_sync_streams: vec![1, 2, 3],
328        };
329        let payload = serde_json::to_vec(&handshake).unwrap();
330        let response = MeshFrame {
331            frame_type: FrameType::Handshake,
332            payload,
333        };
334        stream.send(&response.encode().unwrap()).await.unwrap();
335
336        node_a.await.unwrap();
337    }
338
339    #[tokio::test]
340    async fn cross_node_ipc_delivers() {
341        use crate::ipc::{KernelMessage, MessageTarget};
342        use crate::mesh_ipc::MeshIpcEnvelope;
343
344        let transport = TcpTransport;
345        let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
346        let addr = listener.local_addr().unwrap();
347
348        // Node A sends an IPC message to Node B.
349        let node_a = tokio::spawn(async move {
350            let mut stream = TcpTransport.connect(&addr.to_string()).await.unwrap();
351
352            let msg = KernelMessage::text(1, MessageTarget::Service("health".into()), "ping");
353            let envelope = MeshIpcEnvelope::new("node-a".into(), "node-b".into(), msg);
354            let bytes = envelope.to_bytes().unwrap();
355            stream.send(&bytes).await.unwrap();
356
357            // Receive response.
358            let response_bytes = stream.recv().await.unwrap();
359            let response = MeshIpcEnvelope::from_bytes(&response_bytes).unwrap();
360            assert_eq!(response.source_node, "node-b");
361            assert_eq!(response.dest_node, "node-a");
362
363            stream.close().await.unwrap();
364        });
365
366        // Node B receives and responds.
367        let (mut stream, _) = listener.accept().await.unwrap();
368        let data = stream.recv().await.unwrap();
369        let envelope = MeshIpcEnvelope::from_bytes(&data).unwrap();
370        assert_eq!(envelope.source_node, "node-a");
371        assert_eq!(envelope.dest_node, "node-b");
372
373        // Send response.
374        let response_msg = KernelMessage::text(0, MessageTarget::Kernel, "pong");
375        let response = MeshIpcEnvelope::new("node-b".into(), "node-a".into(), response_msg);
376        stream.send(&response.to_bytes().unwrap()).await.unwrap();
377
378        node_a.await.unwrap();
379    }
380
381    #[test]
382    fn tcp_transport_supports() {
383        let t = TcpTransport;
384        assert!(t.supports("127.0.0.1:9489"));
385        assert!(t.supports("tcp://127.0.0.1:9489"));
386        assert!(!t.supports("quic://127.0.0.1:9489"));
387    }
388
389    /// WEFT-18: race `recv` against another arm; the peer must still decode
390    /// subsequent frames after many mid-frame cancellations (simulates the
391    /// mesh bidirectional `select!` in `boot.rs`).
392    #[tokio::test]
393    async fn recv_survives_select_cancellation() {
394        use std::time::Duration;
395
396        let transport = TcpTransport;
397        let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
398        let addr = listener.local_addr().unwrap();
399
400        let payload_a = vec![0xAAu8; 64];
401        let payload_b = vec![0xBBu8; 128];
402        let payload_a2 = payload_a.clone();
403        let payload_b2 = payload_b.clone();
404
405        let writer = tokio::spawn(async move {
406            let mut stream = TcpTransport.connect(&addr.to_string()).await.unwrap();
407            // Pace writes so the reader often parks mid-frame.
408            stream.send(&payload_a2).await.unwrap();
409            tokio::time::sleep(Duration::from_millis(5)).await;
410            stream.send(&payload_b2).await.unwrap();
411            stream.close().await.unwrap();
412        });
413
414        let (mut server, _) = listener.accept().await.unwrap();
415
416        // Hammer select! cancellations while the first frame arrives.
417        let mut got_first = None;
418        for _ in 0..200 {
419            tokio::select! {
420                biased;
421                result = server.recv() => {
422                    got_first = Some(result.expect("recv after cancel must succeed"));
423                    break;
424                }
425                _ = tokio::time::sleep(Duration::from_micros(50)) => {
426                    // Cancel recv mid-frame; framing state must survive.
427                }
428            }
429        }
430        let first = got_first.expect("should eventually complete first frame");
431        assert_eq!(first, payload_a);
432
433        // Second frame must still align (no length-prefix desync).
434        let second = server.recv().await.expect("second frame");
435        assert_eq!(second, payload_b);
436
437        writer.await.unwrap();
438    }
439
440    /// WEFT-18: length-prefix progress stored on self after partial fill.
441    #[tokio::test]
442    async fn recv_resumes_after_manual_partial_len_state() {
443        let transport = TcpTransport;
444        let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
445        let addr = listener.local_addr().unwrap();
446
447        let writer = tokio::spawn(async move {
448            let mut stream = TcpTransport.connect(&addr.to_string()).await.unwrap();
449            stream.send(b"resume-ok").await.unwrap();
450            stream.close().await.unwrap();
451        });
452
453        let (mut server, _) = listener.accept().await.unwrap();
454
455        // Force a cancel race a few times, then a clean recv.
456        for _ in 0..20 {
457            tokio::select! {
458                biased;
459                r = server.recv() => {
460                    assert_eq!(r.unwrap(), b"resume-ok");
461                    writer.await.unwrap();
462                    return;
463                }
464                _ = std::future::ready(()) => {}
465            }
466        }
467        let msg = server.recv().await.unwrap();
468        assert_eq!(msg, b"resume-ok");
469        writer.await.unwrap();
470    }
471}
clawft_kernel/mesh_tcp.rs

clawft_kernel/
mesh_tcp.rs
