1use crate::error::DecodeError;
2use crate::leb128;
3
4pub struct FrameWriter {
34 buf: Vec<u8>,
35}
36
37impl FrameWriter {
38 pub fn new() -> Self {
40 Self {
41 buf: Vec::with_capacity(256),
42 }
43 }
44
45 pub fn write_frame(&mut self, payload: &[u8]) {
48 leb128::encode(&mut self.buf, payload.len() as u64);
50 self.buf.extend_from_slice(payload);
52 }
53
54 pub fn as_bytes(&self) -> &[u8] {
56 &self.buf
57 }
58
59 pub fn finish(self) -> Vec<u8> {
61 self.buf
62 }
63
64 pub fn reset(&mut self) {
66 self.buf.clear();
67 }
68}
69
70impl Default for FrameWriter {
71 fn default() -> Self {
72 Self::new()
73 }
74}
75
76pub struct FrameReader<'a> {
103 data: &'a [u8],
104 pos: usize,
105}
106
107impl<'a> FrameReader<'a> {
108 pub fn new(data: &'a [u8]) -> Self {
110 Self { data, pos: 0 }
111 }
112
113 pub fn read_frame(&mut self) -> Option<Result<&'a [u8], DecodeError>> {
115 if self.pos >= self.data.len() {
116 return None;
117 }
118
119 let (len, consumed) = match leb128::decode(&self.data[self.pos..], 10) {
120 Ok(v) => v,
121 Err(e) => return Some(Err(e)),
122 };
123 self.pos += consumed;
124
125 let len = len as usize;
126 if self.pos + len > self.data.len() {
127 return Some(Err(DecodeError::UnexpectedEof));
128 }
129
130 let payload = &self.data[self.pos..self.pos + len];
131 self.pos += len;
132 Some(Ok(payload))
133 }
134
135 pub fn is_empty(&self) -> bool {
137 self.pos >= self.data.len()
138 }
139
140 pub fn remaining(&self) -> &'a [u8] {
142 &self.data[self.pos..]
143 }
144}
145
146#[cfg(test)]
147mod tests {
148 use super::*;
149 use crate::{BitReader, BitWriter};
150
151 #[test]
152 fn roundtrip_single_frame() {
153 let mut w = BitWriter::new();
154 w.write_u32(42);
155 let payload = w.finish();
156
157 let mut fw = FrameWriter::new();
158 fw.write_frame(&payload);
159 let stream = fw.finish();
160
161 let mut fr = FrameReader::new(&stream);
162 let frame = fr.read_frame().unwrap().unwrap();
163 let mut r = BitReader::new(frame);
164 assert_eq!(r.read_u32().unwrap(), 42);
165 assert!(fr.read_frame().is_none());
166 }
167
168 #[test]
169 fn roundtrip_multiple_frames() {
170 let mut fw = FrameWriter::new();
171
172 let mut w = BitWriter::new();
173 w.write_u32(1);
174 fw.write_frame(&w.finish());
175
176 let mut w = BitWriter::new();
177 w.write_string("hello");
178 fw.write_frame(&w.finish());
179
180 let mut w = BitWriter::new();
181 w.write_u64(999);
182 fw.write_frame(&w.finish());
183
184 let stream = fw.finish();
185 let mut fr = FrameReader::new(&stream);
186
187 let f1 = fr.read_frame().unwrap().unwrap();
188 assert_eq!(BitReader::new(f1).read_u32().unwrap(), 1);
189
190 let f2 = fr.read_frame().unwrap().unwrap();
191 assert_eq!(BitReader::new(f2).read_string().unwrap(), "hello");
192
193 let f3 = fr.read_frame().unwrap().unwrap();
194 assert_eq!(BitReader::new(f3).read_u64().unwrap(), 999);
195
196 assert!(fr.read_frame().is_none());
197 }
198
199 #[test]
200 fn writer_reset_reuse() {
201 let mut fw = FrameWriter::new();
202
203 let mut w = BitWriter::new();
204 w.write_u32(1);
205 fw.write_frame(&w.finish());
206
207 let bytes1 = fw.finish();
208
209 let mut fw = FrameWriter::new();
210 let mut w = BitWriter::new();
211 w.write_u32(1);
212 fw.write_frame(&w.finish());
213 let bytes2 = fw.finish();
214
215 assert_eq!(bytes1, bytes2);
216 }
217
218 #[test]
219 fn empty_stream() {
220 let mut fr = FrameReader::new(&[]);
221 assert!(fr.read_frame().is_none());
222 assert!(fr.is_empty());
223 }
224
225 #[test]
226 fn corrupted_length() {
227 let stream = [0x80]; let mut fr = FrameReader::new(&stream);
230 match fr.read_frame() {
231 Some(Err(_)) => {} other => panic!("expected error, got {:?}", other),
233 }
234 }
235
236 #[test]
237 fn truncated_payload() {
238 let stream = [10, 0x01, 0x02, 0x03];
240 let mut fr = FrameReader::new(&stream);
241 match fr.read_frame() {
242 Some(Err(DecodeError::UnexpectedEof)) => {} other => panic!("expected EOF error, got {:?}", other),
244 }
245 }
246}