nnrp_runtime/
transport.rs

1use async_trait::async_trait;
2use nnrp_core::{CommonHeader, COMMON_HEADER_LEN};
3use tokio::{
4    io::{AsyncReadExt, AsyncWriteExt},
5    net::{TcpListener, TcpStream, ToSocketAddrs},
6};
7
8use crate::{RuntimeError, RuntimePacket};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub struct RuntimeFrameLimits {
12    pub max_packet_bytes: usize,
13}
14
15impl RuntimeFrameLimits {
16    pub const DEFAULT_MAX_PACKET_BYTES: usize = 64 * 1024 * 1024;
17
18    pub const fn new(max_packet_bytes: usize) -> Self {
19        Self { max_packet_bytes }
20    }
21
22    pub fn validate_packet_len(&self, declared: usize) -> Result<(), RuntimeError> {
23        if declared > self.max_packet_bytes {
24            Err(RuntimeError::FrameTooLarge {
25                declared,
26                max: self.max_packet_bytes,
27            })
28        } else {
29            Ok(())
30        }
31    }
32}
33
34impl Default for RuntimeFrameLimits {
35    fn default() -> Self {
36        Self::new(Self::DEFAULT_MAX_PACKET_BYTES)
37    }
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum RuntimeTransportKind {
42    Tcp,
43    Quic,
44    Ipc,
45    WebSocket,
46}
47
48impl RuntimeTransportKind {
49    pub fn transport_id(self) -> nnrp_core::TransportId {
50        match self {
51            Self::Tcp => nnrp_core::TransportId::Tcp,
52            Self::Quic => nnrp_core::TransportId::Quic,
53            Self::Ipc => nnrp_core::TransportId::Ipc,
54            Self::WebSocket => nnrp_core::TransportId::WebSocket,
55        }
56    }
57}
58
59pub type BoxedFramedTransport = Box<dyn FramedTransport>;
60pub type BoxedFramedListener = Box<dyn FramedListener>;
61
62#[async_trait]
63pub trait FramedTransport: Send {
64    fn transport_kind(&self) -> RuntimeTransportKind;
65    async fn read_packet(&mut self) -> Result<RuntimePacket, RuntimeError>;
66    async fn write_packet(&mut self, packet: &RuntimePacket) -> Result<(), RuntimeError>;
67    async fn close(&mut self) -> Result<(), RuntimeError>;
68}
69
70#[async_trait]
71pub trait FramedListener: Send + Sync {
72    fn transport_kind(&self) -> RuntimeTransportKind;
73    fn local_addr(&self) -> Result<std::net::SocketAddr, RuntimeError>;
74    async fn accept(&self) -> Result<BoxedFramedTransport, RuntimeError>;
75}
76
77#[derive(Debug)]
78pub struct TcpTransport {
79    stream: TcpStream,
80    limits: RuntimeFrameLimits,
81}
82
83impl TcpTransport {
84    pub fn new(stream: TcpStream) -> Self {
85        Self::new_with_limits(stream, RuntimeFrameLimits::default())
86    }
87
88    pub fn new_with_limits(stream: TcpStream, limits: RuntimeFrameLimits) -> Self {
89        Self { stream, limits }
90    }
91
92    pub async fn connect(addr: impl ToSocketAddrs) -> Result<Self, RuntimeError> {
93        Self::connect_with_limits(addr, RuntimeFrameLimits::default()).await
94    }
95
96    pub async fn connect_with_limits(
97        addr: impl ToSocketAddrs,
98        limits: RuntimeFrameLimits,
99    ) -> Result<Self, RuntimeError> {
100        Ok(Self {
101            stream: TcpStream::connect(addr).await?,
102            limits,
103        })
104    }
105}
106
107#[async_trait]
108impl FramedTransport for TcpTransport {
109    fn transport_kind(&self) -> RuntimeTransportKind {
110        RuntimeTransportKind::Tcp
111    }
112
113    async fn read_packet(&mut self) -> Result<RuntimePacket, RuntimeError> {
114        let mut header_bytes = [0u8; COMMON_HEADER_LEN];
115        self.stream.read_exact(&mut header_bytes).await?;
116        let header = CommonHeader::parse(&header_bytes)?;
117        self.limits.validate_packet_len(header.packet_len()?)?;
118
119        let mut metadata = vec![0u8; header.meta_len as usize];
120        if !metadata.is_empty() {
121            self.stream.read_exact(&mut metadata).await?;
122        }
123
124        let mut body = vec![0u8; header.body_len as usize];
125        if !body.is_empty() {
126            self.stream.read_exact(&mut body).await?;
127        }
128
129        Ok(RuntimePacket::from_parts(header, metadata, body)?)
130    }
131
132    async fn write_packet(&mut self, packet: &RuntimePacket) -> Result<(), RuntimeError> {
133        let bytes = packet.to_bytes()?;
134        self.limits.validate_packet_len(bytes.len())?;
135        self.stream.write_all(&bytes).await?;
136        self.stream.flush().await?;
137        Ok(())
138    }
139
140    async fn close(&mut self) -> Result<(), RuntimeError> {
141        self.stream.shutdown().await?;
142        Ok(())
143    }
144}
145
146#[derive(Debug)]
147pub struct TcpFramedListener {
148    listener: TcpListener,
149    limits: RuntimeFrameLimits,
150}
151
152impl TcpFramedListener {
153    pub fn new(listener: TcpListener) -> Self {
154        Self::new_with_limits(listener, RuntimeFrameLimits::default())
155    }
156
157    pub fn new_with_limits(listener: TcpListener, limits: RuntimeFrameLimits) -> Self {
158        Self { listener, limits }
159    }
160
161    pub async fn bind(addr: impl ToSocketAddrs) -> Result<Self, RuntimeError> {
162        Self::bind_with_limits(addr, RuntimeFrameLimits::default()).await
163    }
164
165    pub async fn bind_with_limits(
166        addr: impl ToSocketAddrs,
167        limits: RuntimeFrameLimits,
168    ) -> Result<Self, RuntimeError> {
169        Ok(Self {
170            listener: TcpListener::bind(addr).await?,
171            limits,
172        })
173    }
174}
175
176#[async_trait]
177impl FramedListener for TcpFramedListener {
178    fn transport_kind(&self) -> RuntimeTransportKind {
179        RuntimeTransportKind::Tcp
180    }
181
182    fn local_addr(&self) -> Result<std::net::SocketAddr, RuntimeError> {
183        Ok(self.listener.local_addr()?)
184    }
185
186    async fn accept(&self) -> Result<BoxedFramedTransport, RuntimeError> {
187        let (stream, _) = self.listener.accept().await?;
188        Ok(Box::new(TcpTransport::new_with_limits(stream, self.limits)))
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use nnrp_core::TransportId;
196
197    #[test]
198    fn runtime_transport_kinds_map_to_frozen_transport_ids() {
199        assert_eq!(RuntimeTransportKind::Tcp.transport_id(), TransportId::Tcp);
200        assert_eq!(RuntimeTransportKind::Quic.transport_id(), TransportId::Quic);
201        assert_eq!(RuntimeTransportKind::Ipc.transport_id(), TransportId::Ipc);
202        assert_eq!(
203            RuntimeTransportKind::WebSocket.transport_id(),
204            TransportId::WebSocket
205        );
206    }
207}
nnrp_runtime/transport.rs

nnrp_runtime/
transport.rs
