fix(proto): emitting DecoderError::HeadTooLong

Author: paolobarbolini

watermelon-proto/src/proto/decoder/mod.rs

@@ -83,8 +83,14 @@ pub(super) fn decode(
                 }
 
                 let Some(control_line_len) = memchr::memmem::find(read_buf, b"\r\n") else {
-                    *last_bytes_read = read_buf.len();
-                    return Ok(None);
+                    return if read_buf.len() < MAX_HEAD_LEN {
+                        *last_bytes_read = read_buf.len();
+                        Ok(None)
+                    } else {
+                        Err(DecoderError::HeadTooLong {
+                            len: read_buf.len(),
+                        })
+                    };
                 };
 
                 let mut control_line = read_buf.split_to(control_line_len + "\r\n".len());
@@ -118,10 +124,6 @@ pub(super) fn decode(
                 } else if let Some(info) = control_line.strip_prefix(b"INFO ") {
                     let info = serde_json::from_slice(info).map_err(DecoderError::InvalidInfo)?;
                     Ok(Some(ServerOp::Info { info }))
-                } else if read_buf.len() > MAX_HEAD_LEN {
-                    Err(DecoderError::HeadTooLong {
-                        len: read_buf.len(),
-                    })
                 } else {
                     Err(DecoderError::InvalidCommand)
                 };

watermelon-proto/src/proto/decoder/stream.rs

@@ -47,13 +47,13 @@ impl Default for StreamDecoder {
 #[cfg(test)]
 mod tests {
     use bytes::{BufMut as _, Bytes};
-    use claims::assert_ok_eq;
+    use claims::{assert_matches, assert_ok_eq};
 
     use crate::{
         error::ServerError,
         headers::HeaderMap,
         message::{MessageBase, ServerMessage},
-        proto::server::ServerOp,
+        proto::{error::DecoderError, server::ServerOp},
         Subject,
     };
 
@@ -121,4 +121,14 @@ mod tests {
         );
         assert_ok_eq!(decoder.decode(), None);
     }
+
+    #[test]
+    fn head_too_long() {
+        let mut decoder = StreamDecoder::new();
+        decoder.read_buf().put_bytes(0, 20000);
+        assert_matches!(
+            decoder.decode(),
+            Err(DecoderError::HeadTooLong { len: 20000 })
+        );
+    }
 }