rpi_cli/remote/
framing.rs

1//! Length-prefixed framing for the binary (CBOR) protocol.
2//!
3//! Port of native Pi's `packages/protocol/src/framing.ts`. Each frame is a
4//! 4-byte big-endian unsigned payload length followed by the payload bytes.
5//! [`FrameDecoder`] incrementally splits arbitrary byte chunks into complete
6//! payloads and enforces a maximum frame length so a corrupt/hostile stream
7//! cannot make the peer allocate unbounded memory.
8
9use std::fmt;
10
11/// Native `FRAME_HEADER_LENGTH`.
12pub const FRAME_HEADER_LENGTH: usize = 4;
13/// Native `DEFAULT_MAX_FRAME_LENGTH` (16 MiB).
14pub const DEFAULT_MAX_FRAME_LENGTH: usize = 16 * 1024 * 1024;
15/// Native `PAYLOAD_BLOCK_SIZE`.
16pub const PAYLOAD_BLOCK_SIZE: usize = 64 * 1024;
17
18/// Framing error.
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct FrameError(pub String);
21
22impl fmt::Display for FrameError {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        f.write_str(&self.0)
25    }
26}
27
28impl std::error::Error for FrameError {}
29
30/// Prefix `payload` with its 4-byte big-endian length.
31pub fn encode_frame(payload: &[u8]) -> Result<Vec<u8>, FrameError> {
32    if payload.len() > u32::MAX as usize {
33        return Err(FrameError("Frame payload exceeds the unsigned 32-bit length limit".into()));
34    }
35    let mut frame = Vec::with_capacity(FRAME_HEADER_LENGTH + payload.len());
36    frame.extend_from_slice(&(payload.len() as u32).to_be_bytes());
37    frame.extend_from_slice(payload);
38    Ok(frame)
39}
40
41/// Validate that `frame` contains exactly one complete frame within `max`.
42pub fn assert_complete_frame(frame: &[u8], max_frame_length: usize) -> Result<(), FrameError> {
43    if frame.len() < FRAME_HEADER_LENGTH {
44        return Err(FrameError("Frame does not contain a complete length prefix".into()));
45    }
46    let length = u32::from_be_bytes([frame[0], frame[1], frame[2], frame[3]]) as usize;
47    if length > max_frame_length {
48        return Err(FrameError(format!(
49            "Frame length {length} exceeds configured limit of {max_frame_length}"
50        )));
51    }
52    if frame.len() != FRAME_HEADER_LENGTH + length {
53        return Err(FrameError("Frame must contain exactly one complete payload".into()));
54    }
55    Ok(())
56}
57
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59enum State {
60    Open,
61    Ended,
62    Failed,
63}
64
65/// Incremental decoder: feed byte chunks, get whole payloads back.
66pub struct FrameDecoder {
67    header: [u8; FRAME_HEADER_LENGTH],
68    header_len: usize,
69    max_frame_length: usize,
70    payload: Vec<u8>,
71    expected_payload_len: Option<usize>,
72    state: State,
73}
74
75impl FrameDecoder {
76    pub fn new(max_frame_length: usize) -> Self {
77        Self {
78            header: [0; FRAME_HEADER_LENGTH],
79            header_len: 0,
80            max_frame_length,
81            payload: Vec::new(),
82            expected_payload_len: None,
83            state: State::Open,
84        }
85    }
86
87    pub fn with_default_max() -> Self {
88        Self::new(DEFAULT_MAX_FRAME_LENGTH)
89    }
90
91    /// Feed `chunk`, returning every complete payload it completed.
92    pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<Vec<u8>>, FrameError> {
93        match self.state {
94            State::Ended => return Err(FrameError("Frame decoder has ended".into())),
95            State::Failed => return Err(FrameError("Frame decoder has failed".into())),
96            State::Open => {}
97        }
98
99        let mut frames = Vec::new();
100        let mut offset = 0usize;
101        while offset < chunk.len() {
102            if self.expected_payload_len.is_none() {
103                let take = (FRAME_HEADER_LENGTH - self.header_len).min(chunk.len() - offset);
104                self.header[self.header_len..self.header_len + take]
105                    .copy_from_slice(&chunk[offset..offset + take]);
106                self.header_len += take;
107                offset += take;
108                if self.header_len < FRAME_HEADER_LENGTH {
109                    continue;
110                }
111                let frame_length =
112                    u32::from_be_bytes(self.header) as usize;
113                self.header_len = 0;
114                if frame_length > self.max_frame_length {
115                    self.state = State::Failed;
116                    return Err(FrameError(format!(
117                        "Frame length {frame_length} exceeds configured limit of {}",
118                        self.max_frame_length
119                    )));
120                }
121                if frame_length == 0 {
122                    frames.push(Vec::new());
123                    continue;
124                }
125                self.expected_payload_len = Some(frame_length);
126                self.payload = Vec::with_capacity(frame_length.min(PAYLOAD_BLOCK_SIZE));
127            }
128
129            let expected = self.expected_payload_len.unwrap();
130            let want = expected - self.payload.len();
131            let take = want.min(chunk.len() - offset);
132            self.payload.extend_from_slice(&chunk[offset..offset + take]);
133            offset += take;
134
135            if self.payload.len() == expected {
136                frames.push(std::mem::take(&mut self.payload));
137                self.expected_payload_len = None;
138            }
139        }
140        Ok(frames)
141    }
142
143    /// Mark the stream ended. Any buffered partial frame is an error.
144    pub fn end(&mut self) -> Result<(), FrameError> {
145        if self.header_len != 0 || self.expected_payload_len.is_some() {
146            self.state = State::Failed;
147            return Err(FrameError("Frame decoder ended mid-frame".into()));
148        }
149        self.state = State::Ended;
150        Ok(())
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn encode_round_trip() {
160        let frame = encode_frame(b"hello").unwrap();
161        assert_eq!(&frame[..4], &[0, 0, 0, 5]);
162        assert_complete_frame(&frame, DEFAULT_MAX_FRAME_LENGTH).unwrap();
163    }
164
165    #[test]
166    fn decoder_splits_multiple_and_partial_chunks() {
167        let mut decoder = FrameDecoder::with_default_max();
168        let a = encode_frame(b"one").unwrap();
169        let b = encode_frame(b"two").unwrap();
170        let mut stream = Vec::new();
171        stream.extend_from_slice(&a);
172        stream.extend_from_slice(&b);
173
174        // Deliberately split mid-header and mid-payload.
175        let out1 = decoder.push(&stream[..2]).unwrap();
176        assert!(out1.is_empty());
177        let out2 = decoder.push(&stream[2..5]).unwrap();
178        assert!(out2.is_empty());
179        let out3 = decoder.push(&stream[5..7]).unwrap();
180        assert_eq!(out3, vec![b"one".to_vec()]);
181        let out4 = decoder.push(&stream[7..]).unwrap();
182        assert_eq!(out4, vec![b"two".to_vec()]);
183        decoder.end().unwrap();
184    }
185
186    #[test]
187    fn empty_frame() {
188        let frame = encode_frame(b"").unwrap();
189        let mut decoder = FrameDecoder::with_default_max();
190        assert_eq!(decoder.push(&frame).unwrap(), vec![Vec::<u8>::new()]);
191    }
192
193    #[test]
194    fn rejects_oversized() {
195        let frame = encode_frame(&vec![0u8; 32]).unwrap();
196        assert!(assert_complete_frame(&frame, 16).is_err());
197        let mut decoder = FrameDecoder::new(16);
198        assert!(decoder.push(&frame).is_err());
199    }
200
201    #[test]
202    fn rejects_trailing_bytes() {
203        let mut frame = encode_frame(b"x").unwrap();
204        frame.push(0);
205        assert!(assert_complete_frame(&frame, DEFAULT_MAX_FRAME_LENGTH).is_err());
206    }
207}
rpi_cli/remote/framing.rs

rpi_cli/remote/
framing.rs
