segment_rs/
frame.rs

1use atoi::atoi;
2use bytes::Buf;
3use bytes::Bytes;
4use std::io::Cursor;
5use std::str;
6use thiserror::Error;
7
8/// Identifier for string type
9pub const STRING_IDENT: u8 = b'$';
10/// Identifier for integer type
11pub const INTEGER_IDENT: u8 = b'%';
12/// Identifier for array type
13pub const ARRAY_IDENT: u8 = b'*';
14/// Identifier for boolean type
15pub const BOOLEAN_IDENT: u8 = b'^';
16/// Identifier for null type
17pub const NULL_IDENT: u8 = b'-';
18/// Identifier for map type
19pub const MAP_IDENT: u8 = b'#';
20/// Identifier for double type
21pub const DOUBLE_IDENT: u8 = b'.';
22/// Identifier for error type
23pub const ERROR_IDENT: u8 = b'!';
24
25/// Represents a Segment protocol frame
26#[derive(Debug, PartialEq)]
27pub enum Frame {
28    /// Represents a string frame which is a binary safe string
29    String(Bytes),
30    /// Represents an integer frame which is a signed 64 bit integer
31    Integer(i64),
32    /// Represents an array frame which holds a Vec of frames
33    Array(Vec<Frame>),
34    /// Represents a boolean frame
35    Boolean(bool),
36    /// Represents a null frame
37    Null,
38    /// Represents a map frame which is a collection of key value pairs. A Vec is used to represnt a map
39    Map(Vec<Frame>),
40    /// Represents an double frame which is a 64 bit floating point
41    Double(f64),
42    /// Represents a error frame which is a binary safe string
43    Error(Bytes),
44}
45
46/// Represents frame parsing error
47#[derive(Debug, Error, PartialEq)]
48pub enum ParseFrameError {
49    /// Occurs when there is not enough data to parse the frame. This indicates that more data needs to be buffered
50    #[error("incomplete frame")]
51    Incomplete,
52
53    /// Occurs when we receive a malformed frame
54    #[error("invalid frame format")]
55    InvalidFormat,
56}
57
58impl Frame {
59    /// Returns the name of the enum as a &'static str
60    pub fn as_str(&self) -> &'static str {
61        match self {
62            Frame::String(_) => {"frame::String"}
63            Frame::Integer(_) => {"frame::Integer"}
64            Frame::Array(_) => {"frame::Array"}
65            Frame::Boolean(_) => {"frame::Boolean"}
66            Frame::Null => {"frame::StrNull"}
67            Frame::Map(_) => {"frame::Map"}
68            Frame::Double(_) => {"frame::Double"}
69            Frame::Error(_) => {"frame::Error"}
70        }
71    }
72}
73
74/// Parses the buffered data into frames
75pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Frame, ParseFrameError> {
76    let line = get_line(buf)?;
77    if line.is_empty() {
78        return Err(ParseFrameError::InvalidFormat);
79    }
80    let frame_type = line[0];
81    let line = &line[1..];
82    match frame_type {
83        STRING_IDENT => parse_string(buf, line),
84        INTEGER_IDENT => parse_integer(line),
85        ARRAY_IDENT => parse_array(buf, line),
86        BOOLEAN_IDENT => parse_boolean(line),
87        NULL_IDENT => parse_null(line),
88        MAP_IDENT => parse_map(buf, line),
89        DOUBLE_IDENT => parse_double(line),
90        ERROR_IDENT => parse_error(buf, line),
91        _ => Err(ParseFrameError::InvalidFormat),
92    }
93}
94
95fn get_line<'a>(buf: &mut Cursor<&'a [u8]>) -> Result<&'a [u8], ParseFrameError> {
96    if !buf.has_remaining() {
97        return Err(ParseFrameError::Incomplete);
98    }
99
100    let start = buf.position() as usize;
101    let end = buf.get_ref().len() - 1;
102
103    for i in start..end {
104        if buf.get_ref()[i] == b'\r' && buf.get_ref()[i + 1] == b'\n' {
105            buf.set_position((i + 2) as u64);
106            return Ok(&buf.get_ref()[start..i]);
107        }
108    }
109
110    Err(ParseFrameError::Incomplete)
111}
112
113fn skip(buf: &mut Cursor<&[u8]>, n: usize) -> Result<(), ParseFrameError> {
114    if buf.remaining() < n {
115        return Err(ParseFrameError::Incomplete);
116    }
117    buf.advance(n);
118    Ok(())
119}
120
121fn parse_string(buf: &mut Cursor<&[u8]>, line: &[u8]) -> Result<Frame, ParseFrameError> {
122    let len = atoi::<usize>(line).ok_or(ParseFrameError::InvalidFormat)?;
123    let n = len + 2;
124
125    if buf.remaining() < n {
126        return Err(ParseFrameError::Incomplete);
127    }
128
129    let data = Bytes::copy_from_slice(&buf.chunk()[..len]);
130
131    skip(buf, n)?;
132
133    Ok(Frame::String(data))
134}
135
136fn parse_integer(line: &[u8]) -> Result<Frame, ParseFrameError> {
137    let int = atoi::<i64>(line).ok_or(ParseFrameError::InvalidFormat)?;
138    Ok(Frame::Integer(int))
139}
140
141fn parse_array(buf: &mut Cursor<&[u8]>, line: &[u8]) -> Result<Frame, ParseFrameError> {
142    let len = atoi::<usize>(line).ok_or(ParseFrameError::InvalidFormat)?;
143    let mut vec = Vec::with_capacity(len);
144    for _ in 0..len {
145        vec.push(parse(buf)?);
146    }
147
148    Ok(Frame::Array(vec))
149}
150
151fn parse_boolean(line: &[u8]) -> Result<Frame, ParseFrameError> {
152    if line.len() > 1 {
153        return Err(ParseFrameError::InvalidFormat);
154    }
155
156    let val = line[0];
157
158    match val {
159        b'0' => Ok(Frame::Boolean(false)),
160        b'1' => Ok(Frame::Boolean(true)),
161        _ => Err(ParseFrameError::InvalidFormat),
162    }
163}
164
165fn parse_null(line: &[u8]) -> Result<Frame, ParseFrameError> {
166    if !line.is_empty() {
167        return Err(ParseFrameError::InvalidFormat);
168    }
169    Ok(Frame::Null)
170}
171
172fn parse_map(buf: &mut Cursor<&[u8]>, line: &[u8]) -> Result<Frame, ParseFrameError> {
173    let len = atoi::<usize>(line).ok_or(ParseFrameError::InvalidFormat)?;
174    let mut map = Vec::with_capacity(2 * len);
175    for _ in 0..len {
176        let key = parse(buf)?;
177        let value = parse(buf)?;
178        map.push(key);
179        map.push(value);
180    }
181
182    Ok(Frame::Map(map))
183}
184
185fn parse_double(line: &[u8]) -> Result<Frame, ParseFrameError> {
186    let double = str::from_utf8(line)
187        .map_err(|_| ParseFrameError::InvalidFormat)?
188        .parse::<f64>()
189        .map_err(|_| ParseFrameError::InvalidFormat)?;
190    Ok(Frame::Double(double))
191}
192
193fn parse_error(buf: &mut Cursor<&[u8]>, line: &[u8]) -> Result<Frame, ParseFrameError> {
194    let len = atoi::<usize>(line).ok_or(ParseFrameError::InvalidFormat)?;
195    let n = len + 2;
196
197    if buf.remaining() < n {
198        return Err(ParseFrameError::Incomplete);
199    }
200
201    let data = Bytes::copy_from_slice(&buf.chunk()[..len]);
202
203    skip(buf, n)?;
204
205    Ok(Frame::Error(data))
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use bytes::{BufMut, BytesMut};
212    use std::fs;
213    use std::io::Cursor;
214    use std::path::{Path, PathBuf};
215
216    fn get_cursor_from_bytes(bytes: &[u8]) -> Cursor<&[u8]> {
217        Cursor::new(bytes)
218    }
219
220    fn read_file(path: PathBuf) -> Vec<u8> {
221        fs::read(path).unwrap()
222    }
223
224    fn get_frame_from_file(data: &[u8], ident: u8) -> Bytes {
225        let mut frame = BytesMut::new();
226        frame.put_u8(ident);
227        frame.put_slice(format!("{}", data.len()).as_bytes());
228        frame.put_u8(b'\r');
229        frame.put_u8(b'\n');
230        frame.put(data);
231        frame.put_u8(b'\r');
232        frame.put_u8(b'\n');
233
234        frame.copy_to_bytes(frame.len())
235    }
236
237    #[test]
238    fn parse_given_empty_line_returns_invalid_format_error() {
239        let mut buf = get_cursor_from_bytes(b"\r\n");
240        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
241    }
242
243    #[test]
244    fn parse_given_unknown_type_returns_invalid_format_error() {
245        let mut buf = get_cursor_from_bytes(b"foo\r\n");
246        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
247    }
248
249    #[test]
250    fn parse_given_string_with_no_length_returns_invalid_format_error() {
251        let mut buf = get_cursor_from_bytes(b"$\r\n");
252        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
253    }
254
255    #[test]
256    fn parse_given_string_with_invalid_length_returns_invalid_format_error() {
257        let mut buf = get_cursor_from_bytes(b"$abc\r\n");
258        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
259    }
260
261    #[test]
262    fn parse_given_incomplete_string_with_zero_length_returns_incomplete_error() {
263        let mut buf = get_cursor_from_bytes(b"$0\r\n");
264        assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete))
265    }
266
267    #[test]
268    fn parse_given_incomplete_string_with_non_zero_length_returns_incomplete_error() {
269        let mut buf = get_cursor_from_bytes(b"$1\r\n");
270        assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete))
271    }
272
273    #[test]
274    fn parse_given_string_with_zero_length_returns_empty_string() {
275        let mut buf = get_cursor_from_bytes(b"$0\r\n\r\n");
276        assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from(""))))
277    }
278
279    #[test]
280    fn parse_given_string_with_length_less_than_length_of_data_returns_data_upto_given_length() {
281        let mut buf = get_cursor_from_bytes(b"$1\r\nfoo\r\n");
282        assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from("f"))))
283    }
284
285    #[test]
286    fn parse_given_string_with_length_greater_than_length_of_data_returns_incomplete_error() {
287        let mut buf = get_cursor_from_bytes(b"$100\r\nfoo\r\n");
288        assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete))
289    }
290
291    #[test]
292    fn parse_given_string_returns_string() {
293        let mut buf = get_cursor_from_bytes(b"$3\r\nfoo\r\n");
294        assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from("foo"))))
295    }
296
297    #[test]
298    fn parse_given_string_with_delimiter_in_data_returns_string() {
299        let mut buf = get_cursor_from_bytes(b"$5\r\nfoo\r\n\r\n");
300        assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from("foo\r\n"))))
301    }
302
303    #[test]
304    fn parse_given_string_with_pdf_data_returns_pdf_data() {
305        let file_data = read_file(Path::new("test_data").join("test.pdf"));
306        let frame = get_frame_from_file(file_data.as_slice(), STRING_IDENT);
307        let mut buf = get_cursor_from_bytes(&frame);
308        assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from(file_data))))
309    }
310
311    #[test]
312    fn parse_given_string_with_png_data_returns_png_data() {
313        let file_data = read_file(Path::new("test_data").join("test.png"));
314        let frame = get_frame_from_file(file_data.as_slice(), STRING_IDENT);
315        let mut buf = get_cursor_from_bytes(&frame);
316        assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from(file_data))))
317    }
318
319    #[test]
320    fn parse_given_string_with_jpg_data_returns_jpg_data() {
321        let file_data = read_file(Path::new("test_data").join("test.jpg"));
322        let frame = get_frame_from_file(file_data.as_slice(), STRING_IDENT);
323        let mut buf = get_cursor_from_bytes(&frame);
324        assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from(file_data))))
325    }
326
327    #[test]
328    fn parse_given_string_with_html_data_returns_html_data() {
329        let file_data = read_file(Path::new("test_data").join("test.html"));
330        let frame = get_frame_from_file(file_data.as_slice(), STRING_IDENT);
331        let mut buf = get_cursor_from_bytes(&frame);
332        assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from(file_data))))
333    }
334
335    #[test]
336    fn parse_given_invalid_integer_returns_invalid_format_error() {
337        let mut buf = get_cursor_from_bytes(b"%abc\r\n");
338        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
339    }
340
341    #[test]
342    fn parse_given_empty_integer_returns_invalid_format_error() {
343        let mut buf = get_cursor_from_bytes(b"%\r\n");
344        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
345    }
346
347    #[test]
348    fn parse_given_negative_integer_returns_given_integer() {
349        let mut buf = get_cursor_from_bytes(b"%-1\r\n");
350        assert_eq!(parse(&mut buf), Ok(Frame::Integer(-1)))
351    }
352
353    #[test]
354    fn parse_given_zero_returns_zero() {
355        let mut buf = get_cursor_from_bytes(b"%0\r\n");
356        assert_eq!(parse(&mut buf), Ok(Frame::Integer(0)))
357    }
358
359    #[test]
360    fn parse_given_positive_integer_returns_given_integer() {
361        let mut buf = get_cursor_from_bytes(b"%1000\r\n");
362        assert_eq!(parse(&mut buf), Ok(Frame::Integer(1000)))
363    }
364
365    #[test]
366    fn parse_given_out_of_range_integer_returns_format_error() {
367        let mut buf = get_cursor_from_bytes(b"%9223372036854775808\r\n");
368        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
369    }
370
371    #[test]
372    fn parse_given_false_returns_false() {
373        let mut buf = get_cursor_from_bytes(b"^0\r\n");
374        assert_eq!(parse(&mut buf), Ok(Frame::Boolean(false)))
375    }
376
377    #[test]
378    fn parse_given_true_returns_true() {
379        let mut buf = get_cursor_from_bytes(b"^1\r\n");
380        assert_eq!(parse(&mut buf), Ok(Frame::Boolean(true)))
381    }
382
383    #[test]
384    fn parse_given_invalid_boolean_returns_invalid_format_error() {
385        let mut buf = get_cursor_from_bytes(b"^foo\r\n");
386        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
387    }
388
389    #[test]
390    fn parse_given_null_returns_null() {
391        let mut buf = get_cursor_from_bytes(b"-\r\n");
392        assert_eq!(parse(&mut buf), Ok(Frame::Null))
393    }
394
395    #[test]
396    fn parse_given_invalid_null_returns_invalid_format_error() {
397        let mut buf = get_cursor_from_bytes(b"-foo\r\n");
398        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
399    }
400
401    #[test]
402    fn parse_given_double_with_invalid_decimal_part_returns_invalid_format_error() {
403        let mut buf = get_cursor_from_bytes(b".20.foo\r\n");
404        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
405    }
406
407    #[test]
408    fn parse_given_double_with_invalid_integer_part_returns_invalid_format_error() {
409        let mut buf = get_cursor_from_bytes(b".foo.90\r\n");
410        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
411    }
412
413    #[test]
414    fn parse_given_double_with_zero_decimal_part_returns_double() {
415        let mut buf = get_cursor_from_bytes(b".10.000\r\n");
416        assert_eq!(parse(&mut buf), Ok(Frame::Double(10.0)))
417    }
418
419    #[test]
420    fn parse_given_double_with_trailing_zeroes_returns_double() {
421        let mut buf = get_cursor_from_bytes(b".10.100\r\n");
422        assert_eq!(parse(&mut buf), Ok(Frame::Double(10.1)))
423    }
424
425    #[test]
426    fn parse_given_double_with_leading_zeroes_returns_double() {
427        let mut buf = get_cursor_from_bytes(b".000010.100\r\n");
428        assert_eq!(parse(&mut buf), Ok(Frame::Double(10.1)))
429    }
430
431    #[test]
432    fn parse_given_invalid_double_returns_invalid_format_error() {
433        let mut buf = get_cursor_from_bytes(b".abc\r\n");
434        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
435    }
436
437    #[test]
438    fn parse_given_error_with_no_length_returns_invalid_format_error() {
439        let mut buf = get_cursor_from_bytes(b"!\r\n");
440        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
441    }
442
443    #[test]
444    fn parse_given_error_with_invalid_length_returns_invalid_format_error() {
445        let mut buf = get_cursor_from_bytes(b"!abc\r\n");
446        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
447    }
448
449    #[test]
450    fn parse_given_incomplete_error_with_zero_length_returns_incomplete_error() {
451        let mut buf = get_cursor_from_bytes(b"!0\r\n");
452        assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete))
453    }
454
455    #[test]
456    fn parse_given_incomplete_error_with_non_zero_length_returns_incomplete_error() {
457        let mut buf = get_cursor_from_bytes(b"!1\r\n");
458        assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete))
459    }
460
461    #[test]
462    fn parse_given_error_with_zero_length_returns_empty_error_frame() {
463        let mut buf = get_cursor_from_bytes(b"!0\r\n\r\n");
464        assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from(""))))
465    }
466
467    #[test]
468    fn parse_given_error_with_length_less_than_length_of_data_returns_data_upto_given_length() {
469        let mut buf = get_cursor_from_bytes(b"!1\r\nfoo\r\n");
470        assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from("f"))))
471    }
472
473    #[test]
474    fn parse_given_error_with_length_greater_than_length_of_data_returns_incomplete_error() {
475        let mut buf = get_cursor_from_bytes(b"!100\r\nfoo\r\n");
476        assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete))
477    }
478
479    #[test]
480    fn parse_given_error_returns_error_frame() {
481        let mut buf = get_cursor_from_bytes(b"!3\r\nfoo\r\n");
482        assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from("foo"))))
483    }
484
485    #[test]
486    fn parse_given_error_with_delimiter_in_data_returns_error_frame() {
487        let mut buf = get_cursor_from_bytes(b"!5\r\nfoo\r\n\r\n");
488        assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from("foo\r\n"))))
489    }
490
491    #[test]
492    fn parse_given_error_with_pdf_data_returns_pdf_data() {
493        let file_data = read_file(Path::new("test_data").join("test.pdf"));
494        let frame = get_frame_from_file(file_data.as_slice(), ERROR_IDENT);
495        let mut buf = get_cursor_from_bytes(&frame);
496        assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from(file_data))))
497    }
498
499    #[test]
500    fn parse_given_error_with_png_data_returns_png_data() {
501        let file_data = read_file(Path::new("test_data").join("test.png"));
502        let frame = get_frame_from_file(file_data.as_slice(), ERROR_IDENT);
503        let mut buf = get_cursor_from_bytes(&frame);
504        assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from(file_data))))
505    }
506
507    #[test]
508    fn parse_given_error_with_jpg_data_returns_jpg_data() {
509        let file_data = read_file(Path::new("test_data").join("test.jpg"));
510        let frame = get_frame_from_file(file_data.as_slice(), ERROR_IDENT);
511        let mut buf = get_cursor_from_bytes(&frame);
512        assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from(file_data))))
513    }
514
515    #[test]
516    fn parse_given_error_with_html_data_returns_html_data() {
517        let file_data = read_file(Path::new("test_data").join("test.html"));
518        let frame = get_frame_from_file(file_data.as_slice(), ERROR_IDENT);
519        let mut buf = get_cursor_from_bytes(&frame);
520        assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from(file_data))))
521    }
522
523    #[test]
524    fn parse_given_array_with_no_length_returns_invalid_format_error() {
525        let mut buf = get_cursor_from_bytes(b"*\r\n");
526        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
527    }
528
529    #[test]
530    fn parse_given_array_with_zero_length_returns_empty_array() {
531        let mut buf = get_cursor_from_bytes(b"*0\r\n");
532        assert_eq!(parse(&mut buf), Ok(Frame::Array(Vec::new())))
533    }
534
535    #[test]
536    fn parse_given_array_with_invalid_length_invalid_format_error() {
537        let mut buf = get_cursor_from_bytes(b"*abc\r\n");
538        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
539    }
540
541    #[test]
542    fn parse_given_map_with_no_length_returns_invalid_format_error() {
543        let mut buf = get_cursor_from_bytes(b"#\r\n");
544        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
545    }
546
547    #[test]
548    fn parse_given_map_with_zero_length_returns_empty_map() {
549        let mut buf = get_cursor_from_bytes(b"#0\r\n");
550        assert_eq!(parse(&mut buf), Ok(Frame::Map(Vec::new())))
551    }
552
553    #[test]
554    fn parse_given_map_with_invalid_length_returns_invalid_format_error() {
555        let mut buf = get_cursor_from_bytes(b"#abc\r\n");
556        assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat))
557    }
558
559    #[test]
560    fn parse_given_incomplete_map_return_incomplete_error() {
561        let mut buf = get_cursor_from_bytes(b"#2\r\n$3\r\nfoo\r\n");
562        assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete))
563    }
564}
segment_rs/frame.rs

segment_rs/
frame.rs
