sysknife_daemon/transport/
framing.rs

1use std::io;
2use tokio::io::{AsyncReadExt, AsyncWriteExt};
3
4/// Maximum body size accepted by `recv`. Connections sending a larger
5/// length header are terminated immediately.
6pub const MAX_MESSAGE_BYTES: usize = 4 * 1024 * 1024;
7
8#[derive(Debug, thiserror::Error)]
9pub enum FramingError {
10    #[error("message too large: {0} bytes (max {MAX_MESSAGE_BYTES})")]
11    MessageTooLarge(usize),
12    #[error("io error: {0}")]
13    Io(#[from] io::Error),
14}
15
16/// Wraps an async stream with 4-byte little-endian length-prefix framing.
17///
18/// Each message is sent as `[len: u32 LE][body: len bytes]`. The maximum
19/// body length is [`MAX_MESSAGE_BYTES`]; larger messages are rejected at
20/// both the send and receive side.
21pub struct FramedStream<S> {
22    inner: S,
23}
24
25impl<S> FramedStream<S> {
26    pub fn new(inner: S) -> Self {
27        Self { inner }
28    }
29
30    pub fn into_inner(self) -> S {
31        self.inner
32    }
33}
34
35impl<S: AsyncReadExt + AsyncWriteExt + Unpin> FramedStream<S> {
36    pub async fn send(&mut self, data: &[u8]) -> Result<(), FramingError> {
37        if data.len() > MAX_MESSAGE_BYTES {
38            return Err(FramingError::MessageTooLarge(data.len()));
39        }
40        self.inner
41            .write_all(&(data.len() as u32).to_le_bytes())
42            .await?;
43        self.inner.write_all(data).await?;
44        Ok(())
45    }
46
47    pub async fn recv(&mut self) -> Result<Vec<u8>, FramingError> {
48        let mut header = [0u8; 4];
49        self.inner.read_exact(&mut header).await?;
50        let len = u32::from_le_bytes(header) as usize;
51        if len > MAX_MESSAGE_BYTES {
52            return Err(FramingError::MessageTooLarge(len));
53        }
54        let mut body = vec![0u8; len];
55        self.inner.read_exact(&mut body).await?;
56        Ok(body)
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63    use tokio::io::duplex;
64
65    /// Send `data` from one half of a duplex pair, receive from the other.
66    async fn round_trip(data: &[u8]) -> Vec<u8> {
67        let (a, b) = duplex(MAX_MESSAGE_BYTES + 8);
68        let mut sender = FramedStream::new(a);
69        let mut recvr = FramedStream::new(b);
70        sender.send(data).await.expect("send failed");
71        recvr.recv().await.expect("recv failed")
72    }
73
74    #[tokio::test]
75    async fn round_trip_empty_message() {
76        assert_eq!(round_trip(b"").await, b"");
77    }
78
79    #[tokio::test]
80    async fn round_trip_single_byte() {
81        assert_eq!(round_trip(b"x").await, b"x");
82    }
83
84    #[tokio::test]
85    async fn round_trip_4095_bytes() {
86        let data = vec![0xABu8; 4095];
87        assert_eq!(round_trip(&data).await, data);
88    }
89
90    #[tokio::test]
91    async fn round_trip_4096_bytes() {
92        let data = vec![0xCDu8; 4096];
93        assert_eq!(round_trip(&data).await, data);
94    }
95
96    #[tokio::test]
97    async fn round_trip_json_payload() {
98        let msg = br#"{"type":"preview","action_name":"GetSystemState"}"#;
99        assert_eq!(round_trip(msg).await, msg);
100    }
101
102    #[tokio::test]
103    async fn send_rejects_message_over_4mib() {
104        let (a, _b) = duplex(8);
105        let mut sender = FramedStream::new(a);
106        let oversized = vec![0u8; MAX_MESSAGE_BYTES + 1];
107        let err = sender.send(&oversized).await.unwrap_err();
108        assert!(
109            matches!(err, FramingError::MessageTooLarge(n) if n == MAX_MESSAGE_BYTES + 1),
110            "expected MessageTooLarge, got: {err}"
111        );
112    }
113
114    #[tokio::test]
115    async fn recv_rejects_header_claiming_over_4mib() {
116        let (a, b) = duplex(16);
117        let mut raw_sender = a;
118        let mut recvr = FramedStream::new(b);
119        // Write a header claiming MAX + 1 bytes
120        let oversized_len = (MAX_MESSAGE_BYTES + 1) as u32;
121        raw_sender
122            .write_all(&oversized_len.to_le_bytes())
123            .await
124            .unwrap();
125        let err = recvr.recv().await.unwrap_err();
126        assert!(
127            matches!(err, FramingError::MessageTooLarge(n) if n == MAX_MESSAGE_BYTES + 1),
128            "expected MessageTooLarge, got: {err}"
129        );
130    }
131
132    #[tokio::test]
133    async fn multiple_messages_on_same_stream() {
134        let (a, b) = duplex(MAX_MESSAGE_BYTES + 8);
135        let mut sender = FramedStream::new(a);
136        let mut recvr = FramedStream::new(b);
137        sender.send(b"first").await.unwrap();
138        sender.send(b"second").await.unwrap();
139        assert_eq!(recvr.recv().await.unwrap(), b"first");
140        assert_eq!(recvr.recv().await.unwrap(), b"second");
141    }
142}
sysknife_daemon/transport/framing.rs

sysknife_daemon/transport/
framing.rs
