use std::net::SocketAddr;
use async_trait::async_trait;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use crate::mesh::{MeshError, MeshStream, MeshTransport, TransportListener, MAX_MESSAGE_SIZE};
pub struct TcpMeshStream {
stream: TcpStream,
remote: SocketAddr,
}
#[async_trait]
impl MeshStream for TcpMeshStream {
async fn send(&mut self, data: &[u8]) -> Result<(), MeshError> {
let len = (data.len() as u32).to_be_bytes();
self.stream
.write_all(&len)
.await
.map_err(|e| MeshError::Io(e.to_string()))?;
self.stream
.write_all(data)
.await
.map_err(|e| MeshError::Io(e.to_string()))?;
self.stream
.flush()
.await
.map_err(|e| MeshError::Io(e.to_string()))?;
Ok(())
}
async fn recv(&mut self) -> Result<Vec<u8>, MeshError> {
let mut len_buf = [0u8; 4];
self.stream
.read_exact(&mut len_buf)
.await
.map_err(|e| MeshError::Io(e.to_string()))?;
let len = u32::from_be_bytes(len_buf) as usize;
if len > MAX_MESSAGE_SIZE {
return Err(MeshError::MessageTooLarge {
size: len,
max: MAX_MESSAGE_SIZE,
});
}
let mut buf = vec![0u8; len];
self.stream
.read_exact(&mut buf)
.await
.map_err(|e| MeshError::Io(e.to_string()))?;
Ok(buf)
}
async fn close(&mut self) -> Result<(), MeshError> {
self.stream
.shutdown()
.await
.map_err(|e| MeshError::Io(e.to_string()))
}
fn remote_addr(&self) -> Option<SocketAddr> {
Some(self.remote)
}
}
pub struct TcpTransportListener {
listener: TcpListener,
}
#[async_trait]
impl TransportListener for TcpTransportListener {
async fn accept(&mut self) -> Result<(Box<dyn MeshStream>, SocketAddr), MeshError> {
let (stream, addr) = self
.listener
.accept()
.await
.map_err(|e| MeshError::Io(e.to_string()))?;
Ok((
Box::new(TcpMeshStream {
stream,
remote: addr,
}),
addr,
))
}
fn local_addr(&self) -> Result<SocketAddr, MeshError> {
self.listener
.local_addr()
.map_err(|e| MeshError::Io(e.to_string()))
}
}
pub struct TcpTransport;
#[async_trait]
impl MeshTransport for TcpTransport {
fn name(&self) -> &str {
"tcp"
}
async fn listen(&self, addr: &str) -> Result<Box<dyn TransportListener>, MeshError> {
let bind_addr = addr.strip_prefix("tcp://").unwrap_or(addr);
let listener = TcpListener::bind(bind_addr)
.await
.map_err(|e| MeshError::Io(e.to_string()))?;
Ok(Box::new(TcpTransportListener { listener }))
}
async fn connect(&self, addr: &str) -> Result<Box<dyn MeshStream>, MeshError> {
let connect_addr = addr.strip_prefix("tcp://").unwrap_or(addr);
let stream = TcpStream::connect(connect_addr)
.await
.map_err(|e| MeshError::Io(e.to_string()))?;
let remote = stream
.peer_addr()
.map_err(|e| MeshError::Io(e.to_string()))?;
Ok(Box::new(TcpMeshStream { stream, remote }))
}
fn supports(&self, addr: &str) -> bool {
!addr.contains("://") || addr.starts_with("tcp://")
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn tcp_transport_connect_send_recv() {
let transport = TcpTransport;
let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let send_task = tokio::spawn(async move {
let mut stream = TcpTransport.connect(&addr.to_string()).await.unwrap();
stream.send(b"hello from node A").await.unwrap();
let response = stream.recv().await.unwrap();
assert_eq!(response, b"hello from node B");
stream.close().await.unwrap();
});
let (mut server_stream, _client_addr) = listener.accept().await.unwrap();
let received = server_stream.recv().await.unwrap();
assert_eq!(received, b"hello from node A");
server_stream.send(b"hello from node B").await.unwrap();
send_task.await.unwrap();
}
#[tokio::test]
async fn two_node_cluster_forms() {
use crate::mesh::WeftHandshake;
use crate::mesh_framing::{FrameType, MeshFrame};
let transport = TcpTransport;
let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let node_a = tokio::spawn(async move {
let mut stream = TcpTransport.connect(&addr.to_string()).await.unwrap();
let handshake = WeftHandshake {
node_id: "node-a".into(),
governance_genesis_hash: [0u8; 32],
governance_version: "1.0".into(),
capabilities: 0xFF,
kem_supported: false,
chain_seq: 0,
supported_sync_streams: vec![1, 2, 3],
};
let payload = serde_json::to_vec(&handshake).unwrap();
let frame = MeshFrame {
frame_type: FrameType::Handshake,
payload,
};
let encoded = frame.encode().unwrap();
stream.send(&encoded).await.unwrap();
let data = stream.recv().await.unwrap();
let len = u32::from_be_bytes([data[0], data[1], data[2], data[3]]) as usize;
let response = MeshFrame::decode(&data[4..4 + len]).unwrap();
assert_eq!(response.frame_type, FrameType::Handshake);
let remote: WeftHandshake = serde_json::from_slice(&response.payload).unwrap();
assert_eq!(remote.node_id, "node-b");
stream.close().await.unwrap();
});
let (mut stream, _) = listener.accept().await.unwrap();
let data = stream.recv().await.unwrap();
let len = u32::from_be_bytes([data[0], data[1], data[2], data[3]]) as usize;
let frame = MeshFrame::decode(&data[4..4 + len]).unwrap();
assert_eq!(frame.frame_type, FrameType::Handshake);
let remote: WeftHandshake = serde_json::from_slice(&frame.payload).unwrap();
assert_eq!(remote.node_id, "node-a");
let handshake = WeftHandshake {
node_id: "node-b".into(),
governance_genesis_hash: [0u8; 32],
governance_version: "1.0".into(),
capabilities: 0xFF,
kem_supported: false,
chain_seq: 0,
supported_sync_streams: vec![1, 2, 3],
};
let payload = serde_json::to_vec(&handshake).unwrap();
let response = MeshFrame {
frame_type: FrameType::Handshake,
payload,
};
stream.send(&response.encode().unwrap()).await.unwrap();
node_a.await.unwrap();
}
#[tokio::test]
async fn cross_node_ipc_delivers() {
use crate::ipc::{KernelMessage, MessageTarget};
use crate::mesh_ipc::MeshIpcEnvelope;
let transport = TcpTransport;
let mut listener = transport.listen("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let node_a = tokio::spawn(async move {
let mut stream = TcpTransport.connect(&addr.to_string()).await.unwrap();
let msg = KernelMessage::text(1, MessageTarget::Service("health".into()), "ping");
let envelope = MeshIpcEnvelope::new("node-a".into(), "node-b".into(), msg);
let bytes = envelope.to_bytes().unwrap();
stream.send(&bytes).await.unwrap();
let response_bytes = stream.recv().await.unwrap();
let response = MeshIpcEnvelope::from_bytes(&response_bytes).unwrap();
assert_eq!(response.source_node, "node-b");
assert_eq!(response.dest_node, "node-a");
stream.close().await.unwrap();
});
let (mut stream, _) = listener.accept().await.unwrap();
let data = stream.recv().await.unwrap();
let envelope = MeshIpcEnvelope::from_bytes(&data).unwrap();
assert_eq!(envelope.source_node, "node-a");
assert_eq!(envelope.dest_node, "node-b");
let response_msg = KernelMessage::text(0, MessageTarget::Kernel, "pong");
let response = MeshIpcEnvelope::new("node-b".into(), "node-a".into(), response_msg);
stream.send(&response.to_bytes().unwrap()).await.unwrap();
node_a.await.unwrap();
}
#[test]
fn tcp_transport_supports() {
let t = TcpTransport;
assert!(t.supports("127.0.0.1:9470"));
assert!(t.supports("tcp://127.0.0.1:9470"));
assert!(!t.supports("quic://127.0.0.1:9470"));
}
}