use alloc::boxed::Box;
use alloc::string::String;
use alloc::vec::Vec;
use core::future::Future;
use super::{BoxFut, Connection, Dialer, Listener, PeerInfo, TransportError, TransportResult};
pub type IoError = TransportError;
pub trait ByteRead {
fn read<'a>(
&'a mut self,
buf: &'a mut [u8],
) -> impl Future<Output = TransportResult<usize>> + Send + 'a;
}
pub trait ByteWrite {
fn write_all<'a>(
&'a mut self,
buf: &'a [u8],
) -> impl Future<Output = TransportResult<()>> + Send + 'a;
fn flush(&mut self) -> impl Future<Output = TransportResult<()>> + Send + '_;
}
pub trait ByteStream {
fn read<'a>(
&'a mut self,
buf: &'a mut [u8],
) -> impl Future<Output = TransportResult<usize>> + Send + 'a;
fn write_all<'a>(
&'a mut self,
buf: &'a [u8],
) -> impl Future<Output = TransportResult<()>> + Send + 'a;
fn flush(&mut self) -> impl Future<Output = TransportResult<()>> + Send + '_;
fn split(&mut self) -> (impl ByteRead + Send + '_, impl ByteWrite + Send + '_);
}
pub trait StreamDialer {
type Stream: ByteStream + Send;
fn connect<'a>(
&'a self,
host: &'a str,
port: u16,
) -> impl Future<Output = TransportResult<Self::Stream>> + Send + 'a;
}
pub trait StreamListener {
type Stream: ByteStream + Send;
fn accept(
&mut self,
) -> impl Future<Output = TransportResult<(Self::Stream, PeerInfo)>> + Send + '_;
}
pub trait Datagram {
fn send_to<'a>(
&'a mut self,
buf: &'a [u8],
to: core::net::SocketAddr,
) -> impl Future<Output = TransportResult<()>> + Send + 'a;
fn recv_from<'a>(
&'a mut self,
buf: &'a mut [u8],
) -> impl Future<Output = TransportResult<(usize, core::net::SocketAddr)>> + Send + 'a;
fn local_addr(&self) -> Option<core::net::SocketAddr>;
}
pub trait DatagramBinder {
type Socket: Datagram + Send;
fn bind(&self, port: u16) -> impl Future<Output = TransportResult<Self::Socket>> + Send + '_;
}
pub trait Delay {
fn sleep(&self, d: core::time::Duration) -> impl Future<Output = ()> + Send;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FrameFault {
Recoverable,
Fatal,
}
pub trait Framer {
fn encode(&self, frame: &[u8], out: &mut Vec<u8>) -> Result<(), FrameFault>;
fn push_bytes(&mut self, bytes: &[u8]);
fn next_frame(&mut self) -> Option<Result<Vec<u8>, FrameFault>>;
}
pub trait FramerFactory {
type Framer: Framer + Send;
fn framer(&self) -> Self::Framer;
}
impl<F, T> FramerFactory for F
where
F: Fn() -> T,
T: Framer + Send,
{
type Framer = T;
fn framer(&self) -> T {
self()
}
}
pub struct FramedConnection<S, F, const RC: usize = 256, const WC: usize = 256> {
stream: S,
framer: F,
peer: PeerInfo,
}
impl<S, F, const RC: usize, const WC: usize> FramedConnection<S, F, RC, WC> {
pub fn new(stream: S, framer: F) -> Self {
Self {
stream,
framer,
peer: PeerInfo::default(),
}
}
pub fn with_peer(stream: S, framer: F, peer: PeerInfo) -> Self {
Self {
stream,
framer,
peer,
}
}
}
impl<S, F, const RC: usize, const WC: usize> Connection for FramedConnection<S, F, RC, WC>
where
S: ByteStream + Send,
F: Framer + Send,
{
fn recv(&mut self) -> BoxFut<'_, TransportResult<Option<Vec<u8>>>> {
Box::pin(async move {
loop {
match self.framer.next_frame() {
Some(Ok(frame)) => return Ok(Some(frame)),
Some(Err(FrameFault::Recoverable)) => continue,
Some(Err(FrameFault::Fatal)) => {
log_warn!("framed recv: unrecoverable framing error, closing connection");
return Err(TransportError::Framing);
}
None => {}
}
let mut chunk = [0u8; RC];
match self.stream.read(&mut chunk).await {
Ok(0) => return Ok(None), Ok(n) => self.framer.push_bytes(&chunk[..n]),
Err(e) => return Err(e),
}
}
})
}
fn send<'a>(&'a mut self, frame: &'a [u8]) -> BoxFut<'a, TransportResult<()>> {
Box::pin(async move {
let mut out = Vec::new();
if let Err(_fault) = self.framer.encode(frame, &mut out) {
log_warn!(
"framed send: framer rejected a {}-byte frame ({:?}), closing connection",
frame.len(),
_fault
);
return Err(TransportError::Framing);
}
for chunk in out.chunks(WC) {
self.stream.write_all(chunk).await?;
}
self.stream.flush().await
})
}
fn peer(&self) -> &PeerInfo {
&self.peer
}
}
#[derive(Clone)]
pub struct FramingDialer<D, FF, const RC: usize = 256, const WC: usize = 256> {
dialer: D,
framers: FF,
host: String,
port: u16,
}
impl<D, FF, const RC: usize, const WC: usize> FramingDialer<D, FF, RC, WC> {
pub fn new(dialer: D, framers: FF, host: impl Into<String>, port: u16) -> Self {
Self {
dialer,
framers,
host: host.into(),
port,
}
}
}
impl<D, FF, const RC: usize, const WC: usize> Dialer for FramingDialer<D, FF, RC, WC>
where
D: StreamDialer + Send + Sync,
FF: FramerFactory + Send + Sync,
D::Stream: 'static,
FF::Framer: 'static,
{
fn connect(&self) -> BoxFut<'_, TransportResult<Box<dyn Connection>>> {
Box::pin(async move {
let stream = self.dialer.connect(&self.host, self.port).await?;
let conn: FramedConnection<D::Stream, FF::Framer, RC, WC> =
FramedConnection::new(stream, self.framers.framer());
Ok(Box::new(conn) as Box<dyn Connection>)
})
}
}
pub struct FramingListener<L, FF, const RC: usize = 256, const WC: usize = 256> {
listener: L,
framers: FF,
}
impl<L, FF, const RC: usize, const WC: usize> FramingListener<L, FF, RC, WC> {
pub fn new(listener: L, framers: FF) -> Self {
Self { listener, framers }
}
}
impl<L, FF, const RC: usize, const WC: usize> Listener for FramingListener<L, FF, RC, WC>
where
L: StreamListener + Send,
FF: FramerFactory + Send,
L::Stream: 'static,
FF::Framer: 'static,
{
fn accept(&mut self) -> BoxFut<'_, TransportResult<Box<dyn Connection>>> {
Box::pin(async move {
let (stream, peer) = self.listener.accept().await?;
let conn: FramedConnection<L::Stream, FF::Framer, RC, WC> =
FramedConnection::with_peer(stream, self.framers.framer(), peer);
Ok(Box::new(conn) as Box<dyn Connection>)
})
}
}
pub struct OneShot<T> {
inner: spin::Mutex<Option<T>>,
}
impl<T> OneShot<T> {
pub fn new(value: T) -> Self {
Self {
inner: spin::Mutex::new(Some(value)),
}
}
pub fn take(&self) -> Option<T> {
self.inner.lock().take()
}
}
impl<T> Default for OneShot<T> {
fn default() -> Self {
Self {
inner: spin::Mutex::new(None),
}
}
}
impl<T> From<T> for OneShot<T> {
fn from(value: T) -> Self {
Self::new(value)
}
}
pub struct OneShotDialer<C> {
conn: OneShot<C>,
}
impl<C> OneShotDialer<C> {
pub fn new(conn: C) -> Self {
Self {
conn: OneShot::new(conn),
}
}
}
impl<C: Connection + 'static> Dialer for OneShotDialer<C> {
fn connect(&self) -> BoxFut<'_, TransportResult<Box<dyn Connection>>> {
Box::pin(async move {
self.conn
.take()
.map(|c| Box::new(c) as Box<dyn Connection>)
.ok_or(TransportError::Closed)
})
}
}
pub struct OneShotListener<C> {
conn: OneShot<C>,
}
impl<C> OneShotListener<C> {
pub fn new(conn: C) -> Self {
Self {
conn: OneShot::new(conn),
}
}
}
impl<C: Connection + 'static> Listener for OneShotListener<C> {
fn accept(&mut self) -> BoxFut<'_, TransportResult<Box<dyn Connection>>> {
Box::pin(async move {
match self.conn.take() {
Some(c) => Ok(Box::new(c) as Box<dyn Connection>),
None => core::future::pending().await,
}
})
}
}
#[allow(dead_code)]
fn _generic_code_over_the_traits_is_send<D, L, B, T>(
dialer: D,
mut listener: L,
binder: B,
delay: T,
) -> BoxFut<'static, ()>
where
D: StreamDialer + Send + 'static,
L: StreamListener + Send + 'static,
B: DatagramBinder + Send + 'static,
T: Delay + Send + 'static,
{
Box::pin(async move {
let mut buf = [0u8; 16];
if let Ok(mut stream) = dialer.connect("example.test", 1883).await {
let _ = stream.read(&mut buf).await;
let _ = stream.write_all(&buf).await;
let _ = stream.flush().await;
}
if let Ok((mut stream, _peer)) = listener.accept().await {
let _ = stream.read(&mut buf).await;
}
if let Ok(mut socket) = binder.bind(0).await {
let _ = socket.recv_from(&mut buf).await;
let _ = socket.local_addr();
}
delay.sleep(core::time::Duration::from_millis(1)).await;
})
}
#[allow(dead_code)]
fn _one_shot_is_send_sync_without_unsafe() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<OneShot<Vec<u8>>>();
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::sync::Arc;
use alloc::vec;
#[derive(Default)]
struct LenFramer {
buf: Vec<u8>,
}
impl Framer for LenFramer {
fn encode(&self, frame: &[u8], out: &mut Vec<u8>) -> Result<(), FrameFault> {
if frame.len() >= 0xFE {
return Err(FrameFault::Recoverable);
}
out.push(frame.len() as u8);
out.extend_from_slice(frame);
Ok(())
}
fn push_bytes(&mut self, bytes: &[u8]) {
self.buf.extend_from_slice(bytes);
}
fn next_frame(&mut self) -> Option<Result<Vec<u8>, FrameFault>> {
let len = *self.buf.first()? as usize;
if len == 0xFF {
self.buf.remove(0);
return Some(Err(FrameFault::Recoverable));
}
if len == 0xFE {
self.buf.clear();
return Some(Err(FrameFault::Fatal));
}
if self.buf.len() < len + 1 {
return None;
}
let frame = self.buf[1..len + 1].to_vec();
self.buf.drain(..len + 1);
Some(Ok(frame))
}
}
#[derive(Default)]
struct StreamState {
reads: Vec<Vec<u8>>,
written: Vec<u8>,
write_sizes: Vec<usize>,
flushes: usize,
}
#[derive(Clone, Default)]
struct MockStream(Arc<spin::Mutex<StreamState>>);
impl MockStream {
fn with_reads(reads: Vec<Vec<u8>>) -> Self {
Self(Arc::new(spin::Mutex::new(StreamState {
reads,
..Default::default()
})))
}
}
impl ByteStream for MockStream {
async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> TransportResult<usize> {
let mut st = self.0.lock();
if st.reads.is_empty() {
return Ok(0);
}
let chunk = st.reads.remove(0);
let n = chunk.len().min(buf.len());
buf[..n].copy_from_slice(&chunk[..n]);
Ok(n)
}
async fn write_all<'a>(&'a mut self, buf: &'a [u8]) -> TransportResult<()> {
let mut st = self.0.lock();
st.written.extend_from_slice(buf);
st.write_sizes.push(buf.len());
Ok(())
}
async fn flush(&mut self) -> TransportResult<()> {
self.0.lock().flushes += 1;
Ok(())
}
fn split(&mut self) -> (impl ByteRead + Send + '_, impl ByteWrite + Send + '_) {
(MockHalf(self.clone()), MockHalf(self.clone()))
}
}
struct MockHalf(MockStream);
impl ByteRead for MockHalf {
async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> TransportResult<usize> {
ByteStream::read(&mut self.0, buf).await
}
}
impl ByteWrite for MockHalf {
async fn write_all<'a>(&'a mut self, buf: &'a [u8]) -> TransportResult<()> {
ByteStream::write_all(&mut self.0, buf).await
}
async fn flush(&mut self) -> TransportResult<()> {
ByteStream::flush(&mut self.0).await
}
}
struct FailingStream;
impl ByteStream for FailingStream {
async fn read<'a>(&'a mut self, _buf: &'a mut [u8]) -> TransportResult<usize> {
Err(TransportError::Closed)
}
async fn write_all<'a>(&'a mut self, _buf: &'a [u8]) -> TransportResult<()> {
Err(TransportError::Closed)
}
async fn flush(&mut self) -> TransportResult<()> {
Ok(())
}
fn split(&mut self) -> (impl ByteRead + Send + '_, impl ByteWrite + Send + '_) {
(FailingHalf, FailingHalf)
}
}
struct FailingHalf;
impl ByteRead for FailingHalf {
async fn read<'a>(&'a mut self, _buf: &'a mut [u8]) -> TransportResult<usize> {
Err(TransportError::Closed)
}
}
impl ByteWrite for FailingHalf {
async fn write_all<'a>(&'a mut self, _buf: &'a [u8]) -> TransportResult<()> {
Err(TransportError::Closed)
}
async fn flush(&mut self) -> TransportResult<()> {
Ok(())
}
}
#[derive(Clone, Default)]
struct DuplexMock {
written: Arc<spin::Mutex<Vec<u8>>>,
wrote: Arc<tokio::sync::Notify>,
}
impl ByteStream for DuplexMock {
async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> TransportResult<usize> {
self.wrote.notified().await;
let written = self.written.lock();
let n = written.len().min(buf.len());
buf[..n].copy_from_slice(&written[..n]);
Ok(n)
}
async fn write_all<'a>(&'a mut self, buf: &'a [u8]) -> TransportResult<()> {
self.written.lock().extend_from_slice(buf);
self.wrote.notify_one();
Ok(())
}
async fn flush(&mut self) -> TransportResult<()> {
Ok(())
}
fn split(&mut self) -> (impl ByteRead + Send + '_, impl ByteWrite + Send + '_) {
(DuplexHalf(self.clone()), DuplexHalf(self.clone()))
}
}
struct DuplexHalf(DuplexMock);
impl ByteRead for DuplexHalf {
async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> TransportResult<usize> {
ByteStream::read(&mut self.0, buf).await
}
}
impl ByteWrite for DuplexHalf {
async fn write_all<'a>(&'a mut self, buf: &'a [u8]) -> TransportResult<()> {
ByteStream::write_all(&mut self.0, buf).await
}
async fn flush(&mut self) -> TransportResult<()> {
ByteStream::flush(&mut self.0).await
}
}
struct MockDialer(MockStream);
impl StreamDialer for MockDialer {
type Stream = MockStream;
async fn connect<'a>(&'a self, _host: &'a str, _port: u16) -> TransportResult<MockStream> {
Ok(self.0.clone())
}
}
struct MockListener(Option<MockStream>);
impl StreamListener for MockListener {
type Stream = MockStream;
async fn accept(&mut self) -> TransportResult<(MockStream, PeerInfo)> {
let s = self.0.take().ok_or(TransportError::Closed)?;
let peer = PeerInfo {
peer_addr: Some("10.0.0.1:5555".into()),
..Default::default()
};
Ok((s, peer))
}
}
fn framed(stream: MockStream) -> FramedConnection<MockStream, LenFramer, 256, 256> {
FramedConnection::new(stream, LenFramer::default())
}
#[tokio::test]
async fn recv_yields_each_frame_then_eof() {
let mut conn = framed(MockStream::with_reads(vec![vec![
2, b'h', b'i', 3, b'y', b'e', b's',
]]));
assert_eq!(conn.recv().await.unwrap(), Some(b"hi".to_vec()));
assert_eq!(conn.recv().await.unwrap(), Some(b"yes".to_vec()));
assert_eq!(conn.recv().await.unwrap(), None, "closed peer is Ok(None)");
}
#[tokio::test]
async fn recv_reassembles_a_frame_split_across_reads() {
let mut conn = framed(MockStream::with_reads(vec![
vec![3, b'a'],
vec![b'b'],
vec![b'c'],
]));
assert_eq!(conn.recv().await.unwrap(), Some(b"abc".to_vec()));
}
#[tokio::test]
async fn recv_skips_a_corrupt_run_and_resyncs() {
let mut conn = framed(MockStream::with_reads(vec![vec![
0xFF, 0xFF, 2, b'o', b'k',
]]));
assert_eq!(
conn.recv().await.unwrap(),
Some(b"ok".to_vec()),
"a run that fails to decode is skipped, not fatal"
);
}
#[tokio::test]
async fn recv_closes_on_an_unrecoverable_framing_error() {
let mut conn = framed(MockStream::with_reads(vec![vec![0xFE, 2, b'o', b'k']]));
assert_eq!(
conn.recv().await,
Err(TransportError::Framing),
"a fatal fault ends the connection, it is not skipped"
);
}
#[tokio::test]
async fn send_reports_a_frame_the_framer_rejects() {
let stream = MockStream::default();
let mut conn = framed(stream.clone());
let oversized = vec![b'x'; 0xFE];
assert_eq!(
conn.send(&oversized).await,
Err(TransportError::Framing),
"a dropped frame is an error, never a silent Ok"
);
let st = stream.0.lock();
assert!(
st.written.is_empty(),
"nothing half-encoded reaches the wire"
);
assert_eq!(st.flushes, 0);
}
#[tokio::test]
async fn recv_propagates_the_streams_own_error() {
let mut conn = FramedConnection::<_, _, 256, 256>::new(FailingStream, LenFramer::default());
assert_eq!(
conn.recv().await,
Err(TransportError::Closed),
"the stream classifies the failure; framing must not flatten it"
);
}
#[tokio::test]
async fn send_encodes_then_flushes() {
let stream = MockStream::default();
let mut conn = framed(stream.clone());
conn.send(b"hi").await.unwrap();
let st = stream.0.lock();
assert_eq!(st.written, vec![2, b'h', b'i']);
assert_eq!(st.flushes, 1);
}
#[tokio::test]
async fn send_splits_a_frame_larger_than_the_write_chunk() {
let stream = MockStream::default();
let mut conn: FramedConnection<MockStream, LenFramer, 256, 4> =
FramedConnection::new(stream.clone(), LenFramer::default());
conn.send(b"0123456789").await.unwrap();
let st = stream.0.lock();
assert_eq!(st.write_sizes, vec![4, 4, 3], "11 encoded bytes at WC = 4");
assert_eq!(st.written.len(), 11);
}
#[tokio::test]
async fn framing_dialer_produces_a_working_connection() {
let stream = MockStream::with_reads(vec![vec![2, b'h', b'i']]);
let dialer: FramingDialer<_, _, 256, 256> =
FramingDialer::new(MockDialer(stream), LenFramer::default, "host.test", 1883);
let mut conn = dialer.connect().await.unwrap();
assert_eq!(conn.recv().await.unwrap(), Some(b"hi".to_vec()));
}
#[tokio::test]
async fn framing_listener_carries_peer_metadata_through() {
let stream = MockStream::with_reads(vec![vec![2, b'h', b'i']]);
let mut listener: FramingListener<_, _, 256, 256> =
FramingListener::new(MockListener(Some(stream)), LenFramer::default);
let mut conn = listener.accept().await.unwrap();
assert_eq!(conn.peer().peer_addr.as_deref(), Some("10.0.0.1:5555"));
assert_eq!(conn.recv().await.unwrap(), Some(b"hi".to_vec()));
}
#[test]
fn a_framed_connection_is_boxable_as_dyn_connection() {
let conn = framed(MockStream::default());
let _boxed: Box<dyn Connection> = Box::new(conn);
}
#[tokio::test]
async fn split_halves_run_concurrently() {
let mut stream = DuplexMock::default();
let (mut rx, mut tx) = stream.split();
let mut buf = [0u8; 4];
let (read, write) = tokio::join!(rx.read(&mut buf), tx.write_all(b"ping"));
write.expect("write half");
assert_eq!(read.expect("read half"), 4);
assert_eq!(&buf, b"ping");
}
#[tokio::test]
async fn split_halves_address_the_same_stream() {
let stream = MockStream::with_reads(vec![b"hi".to_vec()]);
let mut split_me = stream.clone();
let (mut rx, mut tx) = split_me.split();
tx.write_all(b"out").await.expect("write half");
let mut buf = [0u8; 8];
let n = rx.read(&mut buf).await.expect("read half");
assert_eq!(&buf[..n], b"hi", "the read half drains the stream's reads");
assert_eq!(
stream.0.lock().written,
b"out",
"the write half reaches the stream the halves came from"
);
}
#[test]
fn one_shot_yields_its_value_once() {
let cell = OneShot::new(7u32);
assert_eq!(cell.take(), Some(7));
assert_eq!(cell.take(), None, "a second take must not get the resource");
}
#[test]
fn default_one_shot_is_empty() {
assert_eq!(OneShot::<u32>::default().take(), None);
}
#[test]
fn framer_factory_is_implemented_for_closures() {
struct Noop;
impl Framer for Noop {
fn encode(&self, _frame: &[u8], _out: &mut Vec<u8>) -> Result<(), FrameFault> {
Ok(())
}
fn push_bytes(&mut self, _bytes: &[u8]) {}
fn next_frame(&mut self) -> Option<Result<Vec<u8>, FrameFault>> {
None
}
}
fn takes_factory<FF: FramerFactory>(ff: FF) -> FF::Framer {
ff.framer()
}
let _framer = takes_factory(|| Noop);
}
}