use anyhow::Result; use tokio::io::{AsyncRead, AsyncWrite, DuplexStream}; #[cfg(feature = "transport-traits")] pub use crate::transport::{Transport, TransportAcceptor, TransportInfo, TransportKind}; #[cfg(not(feature = "transport-traits"))] pub use local_traits::{Transport, TransportAcceptor, TransportInfo, TransportKind}; #[cfg(not(feature = "transport-traits"))] mod local_traits { use anyhow::Result; use async_trait::async_trait; use std::net::SocketAddr; use tokio::io::{AsyncRead, AsyncWrite}; #[async_trait] pub trait Transport: Send + Sync + 'static { type Stream: AsyncRead + AsyncWrite + Unpin + Send + 'static; async fn connect(&self) -> Result; fn describe(&self) -> String; } #[async_trait] pub trait TransportAcceptor: Send + Sync + 'static { type Stream: AsyncRead + AsyncWrite + Unpin + Send + 'static; async fn accept(&self) -> Result<(Self::Stream, TransportInfo)>; } #[derive(Debug, Clone)] pub struct TransportInfo { pub remote_addr: Option, pub transport_kind: TransportKind, } #[derive(Debug, Clone)] pub enum TransportKind { Tcp, Tls { server_name: Option }, Iroh { endpoint_id: String }, } } pub struct MockStream { inner: DuplexStream, } impl MockStream { pub fn new(inner: DuplexStream) -> Self { Self { inner } } } impl AsyncRead for MockStream { fn poll_read( self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, buf: &mut tokio::io::ReadBuf<'_>, ) -> std::task::Poll> { std::pin::Pin::new(&mut self.get_mut().inner).poll_read(cx, buf) } } impl AsyncWrite for MockStream { fn poll_write( self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, buf: &[u8], ) -> std::task::Poll> { std::pin::Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) } fn poll_flush( self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> std::task::Poll> { std::pin::Pin::new(&mut self.get_mut().inner).poll_flush(cx) } fn poll_shutdown( self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> std::task::Poll> { std::pin::Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) } } impl Unpin for MockStream {} pub struct MockTransport { buf_size: usize, } impl MockTransport { pub fn new(buf_size: usize) -> Self { Self { buf_size } } } #[async_trait::async_trait] impl Transport for MockTransport { type Stream = MockStream; async fn connect(&self) -> Result { let (client, _) = tokio::io::duplex(self.buf_size); Ok(MockStream::new(client)) } fn describe(&self) -> String { "mock".to_string() } } pub struct MockTransportAcceptor { buf_size: usize, } impl MockTransportAcceptor { pub fn new(buf_size: usize) -> Self { Self { buf_size } } } #[async_trait::async_trait] impl TransportAcceptor for MockTransportAcceptor { type Stream = MockStream; async fn accept(&self) -> Result<(Self::Stream, TransportInfo)> { let (_, server) = tokio::io::duplex(self.buf_size); let info = TransportInfo { remote_addr: None, transport_kind: TransportKind::Tcp, }; Ok((MockStream::new(server), info)) } } pub fn mock_pair(buf_size: usize) -> (MockStream, MockStream) { let (client, server) = tokio::io::duplex(buf_size); (MockStream::new(client), MockStream::new(server)) }