1use 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
47pub use crate::mesh_clock::{Clock, MockClock, RealClock};
49
50pub 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
115struct PendingConn {
118 stream: InMemoryStream,
119 remote: SocketAddr,
120}
121
122pub 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
144struct ListenerEntry {
147 tx: mpsc::Sender<PendingConn>,
148}
149
150#[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 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 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
253pub struct FaultyTransport {
259 inner: InMemoryTransport,
260 drop_every_n: AtomicU32,
262 fail_next_connects: AtomicU32,
264 latency_ms: AtomicU64,
266}
267
268impl FaultyTransport {
269 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 pub fn set_drop_every_n(&self, n: u32) {
281 self.drop_every_n.store(n, Ordering::SeqCst);
282 }
283
284 pub fn set_fail_next_connects(&self, n: u32) {
286 self.fail_next_connects.store(n, Ordering::SeqCst);
287 }
288
289 pub fn set_latency_ms(&self, ms: u64) {
291 self.latency_ms.store(ms, Ordering::SeqCst);
292 }
293
294 pub fn inner(&self) -> &InMemoryTransport {
296 &self.inner
297 }
298}
299
300struct 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 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 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
369pub struct MockPeer {
376 pub node_id: String,
378 pub address: String,
380 pub governance_genesis: [u8; 32],
382 pub handshake: WeftHandshake,
384 pub inbox: Arc<Mutex<Vec<Vec<u8>>>>,
386 transport: InMemoryTransport,
387}
388
389impl MockPeer {
390 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 pub fn with_transport(mut self, transport: InMemoryTransport) -> Self {
415 self.transport = transport;
416 self
417 }
418
419 pub fn foreign_cluster(node_id: impl Into<String>) -> Self {
421 Self::new(node_id, [0xFF; 32])
422 }
423
424 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 pub fn transport(&self) -> &InMemoryTransport {
437 &self.transport
438 }
439
440 pub fn push_inbox(&self, data: Vec<u8>) {
442 self.inbox
443 .lock()
444 .expect("inbox lock")
445 .push(data);
446 }
447
448 pub fn inbox_snapshot(&self) -> Vec<Vec<u8>> {
450 self.inbox.lock().expect("inbox lock").clone()
451 }
452
453 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
477pub async fn connected_pair() -> Result<(InMemoryStream, InMemoryStream), MeshError> {
481 Ok(InMemoryStream::pair(256))
483}
484
485pub 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 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 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 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}