Skip to main content

rpi_cli/remote/
framing.rs

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