smb_server_transport/
tcp.rs

1//! NBSS-over-TCP transport ([RFC 1002] session service, direct hosting on
2//! port 445), built on `tokio_uring` for zero-copy io_uring networking.
3//!
4//! io_uring socket reads deliver *owned* buffers of arbitrary length, so the
5//! reader keeps a carry-over buffer and hands out exact-length header/body
6//! slices for NBSS framing. Because `TcpStream::read`/`write_all` take `&self`,
7//! the read and write halves share one [`Rc<TcpStream>`]; concurrent read and
8//! write on the same fd is safe under io_uring.
9
10use async_trait::async_trait;
11use std::io;
12use std::rc::Rc;
13use tokio_uring::net::TcpStream;
14
15use crate::{Frame, FrameSink, FrameSource, Transport, TransportError};
16
17/// NetBIOS session-service message types we care about ([RFC 1002] §4).
18mod nb_type {
19    /// Session message carrying an SMB PDU.
20    pub const SESSION_MESSAGE: u8 = 0x00;
21    /// Session request handshake (port 139 style; answered positively).
22    pub const SESSION_REQUEST: u8 = 0x81;
23    /// Positive session response.
24    pub const POSITIVE_RESPONSE: u8 = 0x82;
25}
26
27/// Largest NBSS session-message payload we accept (2 MiB).
28const MAX_FRAME: usize = 0x20_0000;
29/// NBSS header length: message type (1 byte) + 24-bit length (3 bytes).
30const NBSS_HEADER_LEN: usize = 4;
31/// Per-read socket scratch size handed to io_uring.
32const READ_CHUNK: usize = 64 * 1024;
33
34/// Buffered NBSS reader over an io_uring TCP stream.
35struct NbssReader {
36    stream: Rc<TcpStream>,
37    /// Bytes already read from the socket but not yet consumed by framing.
38    carry: Vec<u8>,
39}
40
41impl std::fmt::Debug for NbssReader {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        f.debug_struct("NbssReader").field("carry", &self.carry.len()).finish()
44    }
45}
46
47impl NbssReader {
48    fn new(stream: Rc<TcpStream>) -> Self {
49        Self { stream, carry: Vec::new() }
50    }
51
52    /// Pull one more chunk from the socket into `carry`. Returns `false` at EOF.
53    async fn fill(&mut self) -> Result<bool, TransportError> {
54        let (res, chunk) = self.stream.read(vec![0u8; READ_CHUNK]).await;
55        let n = res?;
56        if n == 0 {
57            return Ok(false);
58        }
59        self.carry.extend_from_slice(&chunk[..n]);
60        Ok(true)
61    }
62
63    /// Take exactly `n` bytes, reading more as needed. `Ok(None)` at clean EOF.
64    async fn take(&mut self, n: usize) -> Result<Option<Vec<u8>>, TransportError> {
65        while self.carry.len() < n {
66            if !self.fill().await? {
67                return Ok(None);
68            }
69        }
70        let rest = self.carry.split_off(n);
71        Ok(Some(std::mem::replace(&mut self.carry, rest)))
72    }
73
74    /// Read one NBSS session-message frame, skipping control frames. When
75    /// `answer` is `Some`, a legacy SESSION_REQUEST is answered positively
76    /// (unified transport only; the split read half passes `None`).
77    async fn next_frame(
78        &mut self,
79        mut answer: Option<&Rc<TcpStream>>,
80    ) -> Result<Option<Frame>, TransportError> {
81        loop {
82            let hdr = match self.take(NBSS_HEADER_LEN).await? {
83                Some(h) => h,
84                None => return Ok(None),
85            };
86            let len = decode_nbss_len(&hdr);
87            match hdr[0] {
88                nb_type::SESSION_MESSAGE => {
89                    if len > MAX_FRAME {
90                        return Err(io::Error::new(
91                            io::ErrorKind::InvalidData,
92                            "oversized NBSS frame",
93                        )
94                        .into());
95                    }
96                    return Ok(self.take(len).await?.map(Frame));
97                }
98                nb_type::SESSION_REQUEST => {
99                    if len > 0 && self.take(len).await?.is_none() {
100                        return Ok(None);
101                    }
102                    if let Some(s) = answer.take() {
103                        write_nbss(s, &[nb_type::POSITIVE_RESPONSE, 0, 0, 0]).await?;
104                    }
105                }
106                _ => {
107                    // Keepalive / unknown control: drain payload and continue.
108                    if len > 0 && self.take(len).await?.is_none() {
109                        return Ok(None);
110                    }
111                }
112            }
113        }
114    }
115}
116
117/// Decode the 24-bit NBSS session-message length from a 4-byte header.
118///
119/// SMB-over-TCP uses the full 24-bit length ([MS-SMB2] §2.1), so all of byte 1
120/// participates. Masking to the RFC 1002 17-bit limit would truncate any frame
121/// >= 128 KiB (multi-credit reads/writes, sealed transforms) and desync the
122/// > stream; the length must decode symmetrically with [`encode_nbss_len`].
123#[cfg_attr(dylint_lib = "no_magic_numbers", allow(no_magic_numbers))]
124fn decode_nbss_len(hdr: &[u8]) -> usize {
125    ((hdr[1] as usize) << 16) | ((hdr[2] as usize) << 8) | hdr[3] as usize
126}
127
128/// Encode a 24-bit NBSS session-message length into the 3 length bytes.
129#[cfg_attr(dylint_lib = "no_magic_numbers", allow(no_magic_numbers))]
130fn encode_nbss_len(len: usize) -> [u8; 3] {
131    [(len >> 16) as u8, (len >> 8) as u8, (len & 0xff) as u8]
132}
133
134/// Frame `data` in an NBSS session message and write it via io_uring.
135async fn write_nbss(stream: &TcpStream, data: &[u8]) -> Result<(), TransportError> {
136    // Frames above 64 KiB (multi-credit reads/writes, sealed transform
137    // payloads) need the middle byte, and above 128 KiB the high byte.
138    let len = encode_nbss_len(data.len());
139    let mut frame = Vec::with_capacity(NBSS_HEADER_LEN + data.len());
140    frame.push(nb_type::SESSION_MESSAGE);
141    frame.extend_from_slice(&len);
142    frame.extend_from_slice(data);
143    let (res, _buf) = stream.write_all(frame).await;
144    res.map_err(Into::into)
145}
146
147/// TCP transport speaking RFC 1002 session protocol over io_uring.
148pub struct TcpTransport {
149    reader: NbssReader,
150    stream: Rc<TcpStream>,
151    peer: String,
152}
153
154impl std::fmt::Debug for TcpTransport {
155    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156        f.debug_struct("TcpTransport").field("peer", &self.peer).finish()
157    }
158}
159
160impl TcpTransport {
161    /// Wrap an already-accepted connection. `peer` is the remote address
162    /// reported by `accept` (io_uring `TcpStream` exposes no `peer_addr`).
163    pub fn new(stream: TcpStream, peer: String) -> Self {
164        let stream = Rc::new(stream);
165        Self {
166            reader: NbssReader::new(stream.clone()),
167            stream,
168            peer,
169        }
170    }
171}
172
173#[async_trait(?Send)]
174impl Transport for TcpTransport {
175    async fn recv(&mut self) -> Result<Option<Frame>, TransportError> {
176        self.reader.next_frame(Some(&self.stream)).await
177    }
178
179    async fn send(&mut self, data: &[u8]) -> Result<(), TransportError> {
180        write_nbss(&self.stream, data).await
181    }
182
183    fn peer(&self) -> String {
184        self.peer.clone()
185    }
186
187    fn split(self: Box<Self>) -> (Box<dyn FrameSource>, Box<dyn FrameSink>) {
188        (
189            Box::new(TcpSource {
190                reader: self.reader,
191                peer: self.peer,
192            }),
193            Box::new(TcpSink {
194                stream: self.stream,
195            }),
196        )
197    }
198}
199
200/// Read half of a split [`TcpTransport`].
201pub struct TcpSource {
202    reader: NbssReader,
203    peer: String,
204}
205
206impl std::fmt::Debug for TcpSource {
207    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
208        f.debug_struct("TcpSource").field("peer", &self.peer).finish()
209    }
210}
211
212#[async_trait(?Send)]
213impl FrameSource for TcpSource {
214    async fn recv(&mut self) -> Result<Option<Frame>, TransportError> {
215        // Split read half cannot write, so legacy SESSION_REQUEST is skipped.
216        self.reader.next_frame(None).await
217    }
218
219    fn peer(&self) -> String {
220        self.peer.clone()
221    }
222}
223
224/// Write half of a split [`TcpTransport`]. Owned by the per-connection writer.
225pub struct TcpSink {
226    stream: Rc<TcpStream>,
227}
228
229impl std::fmt::Debug for TcpSink {
230    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
231        f.debug_struct("TcpSink").finish()
232    }
233}
234
235#[async_trait(?Send)]
236impl FrameSink for TcpSink {
237    async fn send(&mut self, data: &[u8]) -> Result<(), TransportError> {
238        write_nbss(&self.stream, data).await
239    }
240}
241
242#[cfg(test)]
243mod nbss_len_tests {
244    use super::{decode_nbss_len, encode_nbss_len, MAX_FRAME};
245
246    /// Encoding then decoding must round-trip across the whole accepted range,
247    /// including sizes that need the high length byte (>= 128 KiB). This guards
248    /// the framing bug where masking byte 1 to 1 bit truncated large frames.
249    #[test]
250    fn round_trips_across_full_range() {
251        for len in [0usize, 1, 63, 64 * 1024, 128 * 1024, 256 * 1024, MAX_FRAME] {
252            let enc = encode_nbss_len(len);
253            let hdr = [0u8, enc[0], enc[1], enc[2]];
254            assert_eq!(decode_nbss_len(&hdr), len, "round-trip failed for {len}");
255        }
256    }
257
258    /// A 256 KiB frame sets the high length byte; a 17-bit decode would drop it.
259    #[test]
260    fn decodes_high_byte() {
261        let enc = encode_nbss_len(256 * 1024);
262        assert_eq!(enc, [0x04, 0x00, 0x00]);
263        assert_eq!(decode_nbss_len(&[0x00, 0x04, 0x00, 0x00]), 256 * 1024);
264    }
265}
smb_server_transport/tcp.rs

smb_server_transport/
tcp.rs
