Skip to main content

clawft_kernel/
mesh_test_support.rs

1//! Multi-node mesh test fixtures (WEFT-112).
2//!
3//! Provides in-process transport and peer mocks so K6 integration tests
4//! do not need real TCP/QUIC sockets.
5//!
6//! # Fixtures
7//!
8//! | Type | Role |
9//! |------|------|
10//! | [`InMemoryTransport`] | `MeshTransport` over tokio channels (`mem://` URLs) |
11//! | [`FaultyTransport`] | Wraps in-memory transport with drop/fail injection |
12//! | [`MockPeer`] | Lightweight simulated mesh node with inbox |
13//! | [`MockClock`] | Deterministic clock (re-export of [`crate::mesh_clock::MockClock`]) |
14//!
15//! # Downstream authors
16//!
17//! ```ignore
18//! use clawft_kernel::mesh_test_support::{
19//!     connected_pair, FaultyTransport, InMemoryTransport, MockClock, MockPeer,
20//! };
21//!
22//! #[tokio::test]
23//! async fn two_node_exchange() {
24//!     let (mut a, mut b) = connected_pair().await.unwrap();
25//!     a.send(b"hello").await.unwrap();
26//!     assert_eq!(b.recv().await.unwrap(), b"hello");
27//! }
28//! ```
29//!
30//! Prefer `mem://node-a` style addresses. Share one
31//! [`InMemoryTransport`] (or clone via [`InMemoryTransport::clone`])
32//! across peers so `listen`/`connect` meet in the same registry.
33
34use std::collections::HashMap;
35use std::net::{Ipv4Addr, SocketAddr};
36use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
37use std::sync::{Arc, Mutex};
38
39use async_trait::async_trait;
40use tokio::sync::{mpsc, oneshot};
41
42use crate::mesh::{
43    MeshError, MeshPeer, MeshStream, MeshTransport, TransportListener, WeftHandshake,
44};
45use crate::mesh_framing::FrameType;
46
47// Re-export for a single import path (WEFT-112 AC).
48pub use crate::mesh_clock::{Clock, MockClock, RealClock};
49
50// ── In-memory stream ─────────────────────────────────────────────
51
52/// Bidirectional in-memory byte stream (tokio channels).
53pub struct InMemoryStream {
54    tx: Option<mpsc::Sender<Vec<u8>>>,
55    rx: mpsc::Receiver<Vec<u8>>,
56    remote: Option<SocketAddr>,
57    closed: bool,
58}
59
60impl InMemoryStream {
61    fn pair(buffer: usize) -> (Self, Self) {
62        let (a_tx, b_rx) = mpsc::channel(buffer);
63        let (b_tx, a_rx) = mpsc::channel(buffer);
64        let left = Self {
65            tx: Some(a_tx),
66            rx: a_rx,
67            remote: None,
68            closed: false,
69        };
70        let right = Self {
71            tx: Some(b_tx),
72            rx: b_rx,
73            remote: None,
74            closed: false,
75        };
76        (left, right)
77    }
78
79    fn with_remote(mut self, addr: SocketAddr) -> Self {
80        self.remote = Some(addr);
81        self
82    }
83}
84
85#[async_trait]
86impl MeshStream for InMemoryStream {
87    async fn send(&mut self, data: &[u8]) -> Result<(), MeshError> {
88        if self.closed {
89            return Err(MeshError::ConnectionClosed);
90        }
91        let tx = self.tx.as_ref().ok_or(MeshError::ConnectionClosed)?;
92        tx.send(data.to_vec())
93            .await
94            .map_err(|_| MeshError::ConnectionClosed)
95    }
96
97    async fn recv(&mut self) -> Result<Vec<u8>, MeshError> {
98        if self.closed {
99            return Err(MeshError::ConnectionClosed);
100        }
101        self.rx.recv().await.ok_or(MeshError::ConnectionClosed)
102    }
103
104    async fn close(&mut self) -> Result<(), MeshError> {
105        self.closed = true;
106        self.tx.take();
107        Ok(())
108    }
109
110    fn remote_addr(&self) -> Option<SocketAddr> {
111        self.remote
112    }
113}
114
115// ── In-memory listener ───────────────────────────────────────────
116
117struct PendingConn {
118    stream: InMemoryStream,
119    remote: SocketAddr,
120}
121
122/// Listener that yields accepted in-memory streams.
123pub struct InMemoryListener {
124    rx: mpsc::Receiver<PendingConn>,
125    local: SocketAddr,
126}
127
128#[async_trait]
129impl TransportListener for InMemoryListener {
130    async fn accept(&mut self) -> Result<(Box<dyn MeshStream>, SocketAddr), MeshError> {
131        let pending = self
132            .rx
133            .recv()
134            .await
135            .ok_or_else(|| MeshError::Transport("listener closed".into()))?;
136        Ok((Box::new(pending.stream), pending.remote))
137    }
138
139    fn local_addr(&self) -> Result<SocketAddr, MeshError> {
140        Ok(self.local)
141    }
142}
143
144// ── InMemoryTransport ────────────────────────────────────────────
145
146struct ListenerEntry {
147    tx: mpsc::Sender<PendingConn>,
148}
149
150/// Shared registry for in-memory `mem://` endpoints.
151#[derive(Clone, Default)]
152pub struct InMemoryTransport {
153    inner: Arc<Mutex<HashMap<String, ListenerEntry>>>,
154    next_port: Arc<AtomicU32>,
155    buffer: usize,
156}
157
158impl InMemoryTransport {
159    /// Create an empty in-memory transport fabric.
160    pub fn new() -> Self {
161        Self {
162            inner: Arc::new(Mutex::new(HashMap::new())),
163            next_port: Arc::new(AtomicU32::new(10_000)),
164            buffer: 256,
165        }
166    }
167
168    /// Channel buffer depth for new stream pairs.
169    pub fn with_buffer(mut self, buffer: usize) -> Self {
170        self.buffer = buffer.max(1);
171        self
172    }
173
174    fn normalize(addr: &str) -> String {
175        addr.trim().to_string()
176    }
177
178    fn alloc_local_addr(&self) -> SocketAddr {
179        let port = self.next_port.fetch_add(1, Ordering::Relaxed) as u16;
180        SocketAddr::from((Ipv4Addr::LOCALHOST, port))
181    }
182
183    fn synthetic_remote(&self) -> SocketAddr {
184        self.alloc_local_addr()
185    }
186}
187
188#[async_trait]
189impl MeshTransport for InMemoryTransport {
190    fn name(&self) -> &str {
191        "in-memory"
192    }
193
194    async fn listen(&self, addr: &str) -> Result<Box<dyn TransportListener>, MeshError> {
195        let key = Self::normalize(addr);
196        if !self.supports(&key) {
197            return Err(MeshError::Transport(format!(
198                "in-memory transport requires mem:// address, got {addr}"
199            )));
200        }
201        let (tx, rx) = mpsc::channel(self.buffer);
202        let local = self.alloc_local_addr();
203        {
204            let mut map = self
205                .inner
206                .lock()
207                .map_err(|_| MeshError::Transport("lock poisoned".into()))?;
208            if map.contains_key(&key) {
209                return Err(MeshError::Transport(format!("already listening on {key}")));
210            }
211            map.insert(key, ListenerEntry { tx });
212        }
213        Ok(Box::new(InMemoryListener { rx, local }))
214    }
215
216    async fn connect(&self, addr: &str) -> Result<Box<dyn MeshStream>, MeshError> {
217        let key = Self::normalize(addr);
218        if !self.supports(&key) {
219            return Err(MeshError::Transport(format!(
220                "in-memory transport requires mem:// address, got {addr}"
221            )));
222        }
223        let (client, server) = InMemoryStream::pair(self.buffer);
224        let remote = self.synthetic_remote();
225        let server = server.with_remote(remote);
226        let client = client.with_remote(remote);
227
228        let tx = {
229            let map = self
230                .inner
231                .lock()
232                .map_err(|_| MeshError::Transport("lock poisoned".into()))?;
233            map.get(&key)
234                .map(|e| e.tx.clone())
235                .ok_or_else(|| MeshError::Transport(format!("no listener for {key}")))?
236        };
237
238        tx.send(PendingConn {
239            stream: server,
240            remote,
241        })
242        .await
243        .map_err(|_| MeshError::Transport("listener closed".into()))?;
244
245        Ok(Box::new(client))
246    }
247
248    fn supports(&self, addr: &str) -> bool {
249        addr.starts_with("mem://")
250    }
251}
252
253// ── FaultyTransport ──────────────────────────────────────────────
254
255/// Fault-injection wrapper around [`InMemoryTransport`].
256///
257/// Useful for SWIM / circuit-breaker / reconnect tests.
258pub struct FaultyTransport {
259    inner: InMemoryTransport,
260    /// Drop every Nth *message* on streams created via connect (0 = never).
261    drop_every_n: AtomicU32,
262    /// Fail the next N `connect` attempts.
263    fail_next_connects: AtomicU32,
264    /// Artificial delay applied to connect (milliseconds). 0 = none.
265    latency_ms: AtomicU64,
266}
267
268impl FaultyTransport {
269    /// Wrap an in-memory transport with default (healthy) fault settings.
270    pub fn new(inner: InMemoryTransport) -> Self {
271        Self {
272            inner,
273            drop_every_n: AtomicU32::new(0),
274            fail_next_connects: AtomicU32::new(0),
275            latency_ms: AtomicU64::new(0),
276        }
277    }
278
279    /// Drop every Nth sent message (`0` disables).
280    pub fn set_drop_every_n(&self, n: u32) {
281        self.drop_every_n.store(n, Ordering::SeqCst);
282    }
283
284    /// Fail the next `n` connect attempts with a transport error.
285    pub fn set_fail_next_connects(&self, n: u32) {
286        self.fail_next_connects.store(n, Ordering::SeqCst);
287    }
288
289    /// Artificial connect latency in milliseconds.
290    pub fn set_latency_ms(&self, ms: u64) {
291        self.latency_ms.store(ms, Ordering::SeqCst);
292    }
293
294    /// Access the underlying fabric (e.g. to share listeners).
295    pub fn inner(&self) -> &InMemoryTransport {
296        &self.inner
297    }
298}
299
300/// Stream that optionally drops outbound messages.
301struct FaultyStream {
302    inner: Box<dyn MeshStream>,
303    drop_every_n: Arc<AtomicU32>,
304    msg_counter: Arc<AtomicU64>,
305}
306
307#[async_trait]
308impl MeshStream for FaultyStream {
309    async fn send(&mut self, data: &[u8]) -> Result<(), MeshError> {
310        let n = self.drop_every_n.load(Ordering::SeqCst);
311        if n > 0 {
312            let c = self.msg_counter.fetch_add(1, Ordering::SeqCst) + 1;
313            if c.is_multiple_of(u64::from(n)) {
314                // Silently drop.
315                return Ok(());
316            }
317        }
318        self.inner.send(data).await
319    }
320
321    async fn recv(&mut self) -> Result<Vec<u8>, MeshError> {
322        self.inner.recv().await
323    }
324
325    async fn close(&mut self) -> Result<(), MeshError> {
326        self.inner.close().await
327    }
328
329    fn remote_addr(&self) -> Option<SocketAddr> {
330        self.inner.remote_addr()
331    }
332}
333
334#[async_trait]
335impl MeshTransport for FaultyTransport {
336    fn name(&self) -> &str {
337        "faulty-in-memory"
338    }
339
340    async fn listen(&self, addr: &str) -> Result<Box<dyn TransportListener>, MeshError> {
341        self.inner.listen(addr).await
342    }
343
344    async fn connect(&self, addr: &str) -> Result<Box<dyn MeshStream>, MeshError> {
345        let remaining = self.fail_next_connects.load(Ordering::SeqCst);
346        if remaining > 0 {
347            self.fail_next_connects
348                .store(remaining.saturating_sub(1), Ordering::SeqCst);
349            return Err(MeshError::Transport("injected connect failure".into()));
350        }
351        let latency = self.latency_ms.load(Ordering::SeqCst);
352        if latency > 0 {
353            tokio::time::sleep(std::time::Duration::from_millis(latency)).await;
354        }
355        let stream = self.inner.connect(addr).await?;
356        Ok(Box::new(FaultyStream {
357            inner: stream,
358            drop_every_n: Arc::new(AtomicU32::new(self.drop_every_n.load(Ordering::SeqCst))),
359            // Share counter across streams created from this transport.
360            msg_counter: Arc::new(AtomicU64::new(0)),
361        }))
362    }
363
364    fn supports(&self, addr: &str) -> bool {
365        self.inner.supports(addr)
366    }
367}
368
369// ── MockPeer ─────────────────────────────────────────────────────
370
371/// Lightweight simulated mesh node for multi-node tests.
372///
373/// Holds identity, handshake metadata, an optional shared transport, and
374/// an inbox of received raw frames / payloads for assertions.
375pub struct MockPeer {
376    /// Node identity string.
377    pub node_id: String,
378    /// Listen / dial address (`mem://...`).
379    pub address: String,
380    /// Governance genesis for cluster isolation tests.
381    pub governance_genesis: [u8; 32],
382    /// Handshake payload this peer would present.
383    pub handshake: WeftHandshake,
384    /// Received payloads (raw bytes).
385    pub inbox: Arc<Mutex<Vec<Vec<u8>>>>,
386    transport: InMemoryTransport,
387}
388
389impl MockPeer {
390    /// Create a peer with a fresh identity in the given cluster.
391    pub fn new(node_id: impl Into<String>, genesis: [u8; 32]) -> Self {
392        let node_id = node_id.into();
393        let address = format!("mem://{node_id}");
394        let handshake = WeftHandshake {
395            node_id: node_id.clone(),
396            governance_genesis_hash: genesis,
397            governance_version: "1.0.0".into(),
398            capabilities: 0,
399            kem_supported: false,
400            chain_seq: 0,
401            supported_sync_streams: FrameType::ALL.iter().map(|t| *t as u8).collect(),
402        };
403        Self {
404            node_id,
405            address,
406            governance_genesis: genesis,
407            handshake,
408            inbox: Arc::new(Mutex::new(Vec::new())),
409            transport: InMemoryTransport::new(),
410        }
411    }
412
413    /// Attach a shared transport fabric (required for multi-peer meshes).
414    pub fn with_transport(mut self, transport: InMemoryTransport) -> Self {
415        self.transport = transport;
416        self
417    }
418
419    /// Peer that rejects foreign clusters (different genesis).
420    pub fn foreign_cluster(node_id: impl Into<String>) -> Self {
421        Self::new(node_id, [0xFF; 32])
422    }
423
424    /// Build a [`MeshPeer`] view of this mock.
425    pub fn as_mesh_peer(&self) -> MeshPeer {
426        MeshPeer {
427            node_id: self.node_id.clone(),
428            handshake: self.handshake.clone(),
429            address: "127.0.0.1:0"
430                .parse()
431                .expect("static socket addr"),
432        }
433    }
434
435    /// Shared transport handle.
436    pub fn transport(&self) -> &InMemoryTransport {
437        &self.transport
438    }
439
440    /// Push a payload into the inbox (for tests that inject delivery).
441    pub fn push_inbox(&self, data: Vec<u8>) {
442        self.inbox
443            .lock()
444            .expect("inbox lock")
445            .push(data);
446    }
447
448    /// Snapshot of inbox contents.
449    pub fn inbox_snapshot(&self) -> Vec<Vec<u8>> {
450        self.inbox.lock().expect("inbox lock").clone()
451    }
452
453    /// Start accepting one connection and reading frames into the inbox.
454    ///
455    /// Returns a oneshot that fires when the accept loop has bound the
456    /// listener (so dialers can race-free connect).
457    pub async fn listen_once(
458        &self,
459    ) -> Result<(tokio::task::JoinHandle<()>, oneshot::Receiver<()>), MeshError> {
460        let mut listener = self.transport.listen(&self.address).await?;
461        let (ready_tx, ready_rx) = oneshot::channel();
462        let inbox = Arc::clone(&self.inbox);
463        let handle = tokio::spawn(async move {
464            let _ = ready_tx.send(());
465            if let Ok((mut stream, _)) = listener.accept().await {
466                while let Ok(data) = stream.recv().await {
467                    if let Ok(mut guard) = inbox.lock() {
468                        guard.push(data);
469                    }
470                }
471            }
472        });
473        Ok((handle, ready_rx))
474    }
475}
476
477// ── Helpers ──────────────────────────────────────────────────────
478
479/// Create a connected pair of in-memory streams (no peers).
480pub async fn connected_pair() -> Result<(InMemoryStream, InMemoryStream), MeshError> {
481    // Direct pair without registry — fastest unit path.
482    Ok(InMemoryStream::pair(256))
483}
484
485/// Two mock peers sharing a fabric, with an established stream pair.
486///
487/// Returns `(peer_a, peer_b, stream_from_a, stream_from_b)` after A
488/// listens and B connects.
489pub async fn two_node_streams(
490    genesis: [u8; 32],
491) -> Result<
492    (
493        MockPeer,
494        MockPeer,
495        Box<dyn MeshStream>,
496        Box<dyn MeshStream>,
497    ),
498    MeshError,
499> {
500    let fabric = InMemoryTransport::new();
501    let peer_a = MockPeer::new("node-a", genesis).with_transport(fabric.clone());
502    let peer_b = MockPeer::new("node-b", genesis).with_transport(fabric.clone());
503
504    let mut listener = fabric.listen(&peer_a.address).await?;
505    let connect_fut = fabric.connect(&peer_a.address);
506    let accept_fut = listener.accept();
507    let (client, (server, _)) = tokio::try_join!(connect_fut, accept_fut)?;
508    Ok((peer_a, peer_b, client, server))
509}
510
511#[cfg(test)]
512mod tests {
513    use super::*;
514    use crate::mesh_clock::MockClock;
515    use crate::mesh_framing::{FrameType, MeshFrame};
516    use crate::mesh_heartbeat::{HeartbeatConfig, HeartbeatState, HeartbeatTracker};
517    use std::time::Duration;
518
519    #[tokio::test]
520    async fn in_memory_two_node_exchange() {
521        let fabric = InMemoryTransport::new();
522        let mut listener = fabric.listen("mem://alpha").await.unwrap();
523        let connect = fabric.connect("mem://alpha");
524        let accept = listener.accept();
525        let (mut client, (mut server, addr)) = tokio::try_join!(connect, accept).unwrap();
526        assert_eq!(addr.ip(), Ipv4Addr::LOCALHOST);
527
528        client.send(b"ping").await.unwrap();
529        assert_eq!(server.recv().await.unwrap(), b"ping");
530        server.send(b"pong").await.unwrap();
531        assert_eq!(client.recv().await.unwrap(), b"pong");
532    }
533
534    #[tokio::test]
535    async fn mock_peer_inbox_receives_frames() {
536        let genesis = [0x11; 32];
537        let fabric = InMemoryTransport::new();
538        let peer = MockPeer::new("peer-a", genesis).with_transport(fabric.clone());
539        let (handle, ready) = peer.listen_once().await.unwrap();
540        ready.await.unwrap();
541
542        let mut stream = fabric.connect(&peer.address).await.unwrap();
543        let frame = MeshFrame {
544            frame_type: FrameType::Heartbeat,
545            payload: b"tick".to_vec(),
546        };
547        stream.send(&frame.encode().unwrap()).await.unwrap();
548        stream.close().await.unwrap();
549        let _ = handle.await;
550
551        let inbox = peer.inbox_snapshot();
552        assert_eq!(inbox.len(), 1);
553        let decoded = MeshFrame::decode(&inbox[0][4..]).unwrap();
554        assert_eq!(decoded.frame_type, FrameType::Heartbeat);
555        assert_eq!(decoded.payload, b"tick");
556    }
557
558    #[tokio::test]
559    async fn faulty_transport_fails_connect() {
560        let fabric = InMemoryTransport::new();
561        let _listener = fabric.listen("mem://fault").await.unwrap();
562        let faulty = FaultyTransport::new(fabric);
563        faulty.set_fail_next_connects(1);
564        match faulty.connect("mem://fault").await {
565            Ok(_) => panic!("expected injected connect failure"),
566            Err(err) => assert!(err.to_string().contains("injected connect failure")),
567        }
568        // Second connect succeeds.
569        let _ok = faulty.connect("mem://fault").await.expect("connect after fail budget");
570    }
571
572    #[tokio::test]
573    async fn faulty_transport_drops_messages() {
574        let fabric = InMemoryTransport::new();
575        let mut listener = fabric.listen("mem://drop").await.unwrap();
576        let faulty = FaultyTransport::new(fabric);
577        // Drop every message (N=1).
578        faulty.set_drop_every_n(1);
579        let connect = faulty.connect("mem://drop");
580        let accept = listener.accept();
581        let (mut client, (mut server, _)) = tokio::try_join!(connect, accept).unwrap();
582
583        client.send(b"lost").await.unwrap();
584        // Server should not receive; use timeout.
585        let result = tokio::time::timeout(Duration::from_millis(50), server.recv()).await;
586        assert!(result.is_err(), "dropped message must not arrive");
587    }
588
589    #[test]
590    fn mock_clock_drives_heartbeat() {
591        let clock = MockClock::new();
592        let config = HeartbeatConfig {
593            suspect_timeout: Duration::from_millis(100),
594            ..HeartbeatConfig::default()
595        };
596        let mut tracker = HeartbeatTracker::with_clock(config, Arc::new(clock.clone()));
597        tracker.add_peer("n1".into());
598        tracker.record_miss("n1");
599        assert_eq!(tracker.peer_state("n1"), Some(HeartbeatState::Suspect));
600        clock.advance(Duration::from_millis(150));
601        tracker.record_miss("n1");
602        assert_eq!(tracker.peer_state("n1"), Some(HeartbeatState::Dead));
603    }
604
605    #[tokio::test]
606    async fn connected_pair_helper() {
607        let (mut a, mut b) = connected_pair().await.unwrap();
608        a.send(b"x").await.unwrap();
609        assert_eq!(b.recv().await.unwrap(), b"x");
610    }
611
612    #[tokio::test]
613    async fn two_node_streams_helper() {
614        let (_a, _b, mut sa, mut sb) = two_node_streams([7; 32]).await.unwrap();
615        sa.send(b"hello").await.unwrap();
616        assert_eq!(sb.recv().await.unwrap(), b"hello");
617    }
618
619    #[test]
620    fn mock_peer_foreign_cluster() {
621        let p = MockPeer::foreign_cluster("x");
622        assert_eq!(p.governance_genesis, [0xFF; 32]);
623    }
624}