Add `read_line_pedantic` for non-empty lines

server/src/protocol/v2/mod.rs

@@ -128,8 +128,25 @@ impl<'a> Parser<'a> {
             }
         }
     }
+    /// Attempt to read a line, **rejecting an empty payload**
+    fn read_line_pedantic(&mut self) -> ParseResult<UnsafeSlice> {
+        let start_ptr = self.cursor_ptr();
+        unsafe {
+            while self.not_exhausted() && self.get_byte_at_cursor() != b'\n' {
+                self.incr_cursor();
+            }
+            let len = self.cursor_ptr() as usize - start_ptr as usize;
+            if self.not_exhausted() && len != 0 && self.get_byte_at_cursor() == b'\n' {
+                self.incr_cursor(); // skip LF
+                Ok(UnsafeSlice::new(start_ptr, len))
+            } else {
+                Err(ParseError::NotEnough)
+            }
+        }
+    }
+    /// Attempt to read an `usize` from the buffer
     fn read_usize(&mut self) -> ParseResult<usize> {
-        let line = self.read_line()?;
+        let line = self.read_line_pedantic()?;
         let bytes = unsafe {
             // UNSAFE(@ohsayan): We just extracted the slice
             line.as_slice()

server/src/protocol/v2/tests.rs

@@ -419,3 +419,91 @@ fn read_line_subsequent_lf() {
     }
     ensure_exhausted(&p1);
 }
+
+#[test]
+fn read_line_pedantic_okay() {
+    for (len, src) in slices_lf_with_len() {
+        let mut parser = Parser::new(&src);
+        if len == 0 {
+            // should be empty, so NotEnough
+            assert_eq!(
+                parser.read_line_pedantic().unwrap_err(),
+                ParseError::NotEnough
+            );
+        } else {
+            // should work
+            unsafe {
+                assert_eq!(
+                    parser.read_line_pedantic().unwrap().as_slice(),
+                    &src.as_slice()[..len - 1]
+                );
+            }
+            // now, we attempt to read which should work
+            ensure_zero_reads(&mut parser);
+        }
+        // ensure it is exhausted
+        ensure_exhausted(&parser);
+        // now, we attempt to read another line which should fail
+        assert_eq!(
+            parser.read_line_pedantic().unwrap_err(),
+            ParseError::NotEnough
+        );
+        // ensure that cursor is at end
+        unsafe {
+            assert_eq!(parser.cursor_ptr(), src.as_ptr().add(len));
+        }
+    }
+}
+
+#[test]
+fn read_line_pedantic_fail_empty() {
+    let payload = v!(b"");
+    assert_eq!(
+        Parser::new(&payload).read_line_pedantic().unwrap_err(),
+        ParseError::NotEnough
+    );
+}
+
+#[test]
+fn read_line_pedantic_fail_only_lf() {
+    let payload = v!(b"\n");
+    assert_eq!(
+        Parser::new(&payload).read_line_pedantic().unwrap_err(),
+        ParseError::NotEnough
+    );
+}
+
+#[test]
+fn read_line_pedantic_fail_only_lf_extra_data() {
+    let payload = v!(b"\n1");
+    assert_eq!(
+        Parser::new(&payload).read_line_pedantic().unwrap_err(),
+        ParseError::NotEnough
+    );
+}
+
+#[test]
+fn read_usize_fail_empty() {
+    let payload = v!(b"");
+    assert_eq!(
+        Parser::new(&payload).read_usize().unwrap_err(),
+        ParseError::NotEnough
+    );
+}
+
+#[test]
+fn read_usize_fail_no_lf() {
+    let payload = v!(b"1");
+    assert_eq!(
+        Parser::new(&payload).read_usize().unwrap_err(),
+        ParseError::NotEnough
+    );
+}
+
+#[test]
+fn read_usize_okay() {
+    let payload = v!(b"1\n");
+    assert_eq!(Parser::new(&payload).read_usize().unwrap(), 1);
+    let payload = v!(b"1234\n");
+    assert_eq!(Parser::new(&payload).read_usize().unwrap(), 1234);
+}