rpi_cli/remote/
framing.rs1use std::fmt;
10
11pub const FRAME_HEADER_LENGTH: usize = 4;
13pub const DEFAULT_MAX_FRAME_LENGTH: usize = 16 * 1024 * 1024;
15pub const PAYLOAD_BLOCK_SIZE: usize = 64 * 1024;
17
18#[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
30pub 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
43pub 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
71pub 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 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 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 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}