fix(web): Client decoding incomplete buffer bug (#1540)

When the server returns a buffer that contains the message and a few
more bytes we previously did not return the end of the message and thus
tossed the begining of the bytes of the next message causing the
following ones to report incorrect length and triggering eof errors.

tonic-web/src/call.rs

@@ -245,36 +245,41 @@ where
             let mut me = self.as_mut();
 
             loop {
-                let buf = ready!(me.as_mut().poll_decode(cx));
-
-                return if let Some(Ok(incoming_buf)) = buf {
-                    let buf = &mut me.as_mut().project().buf;
-
-                    buf.put(incoming_buf);
-
-                    match find_trailers(&buf[..]) {
-                        FindTrailers::Trailer(len) => {
-                            // Extract up to len of where the trailers are at
-                            let msg_buf = buf.copy_to_bytes(len);
-                            match decode_trailers_frame(buf.split().freeze()) {
-                                Ok(Some(trailers)) => {
-                                    self.project().trailers.replace(trailers);
-                                }
-                                Err(e) => return Poll::Ready(Some(Err(e))),
-                                _ => {}
-                            }
+                let incoming_buf = match ready!(me.as_mut().poll_decode(cx)) {
+                    Some(Ok(incoming_buf)) => incoming_buf,
+                    None => {
+                        // TODO: Consider eofing here?
+                        // Even if the buffer has more data, this will hit the eof branch
+                        // of decode in tonic
+                        return Poll::Ready(None);
+                    }
+                    Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+                };
+
+                let buf = &mut me.as_mut().project().buf;
 
-                            if msg_buf.has_remaining() {
-                                return Poll::Ready(Some(Ok(msg_buf)));
-                            } else {
-                                return Poll::Ready(None);
+                buf.put(incoming_buf);
+
+                return match find_trailers(&buf[..])? {
+                    FindTrailers::Trailer(len) => {
+                        // Extract up to len of where the trailers are at
+                        let msg_buf = buf.copy_to_bytes(len);
+                        match decode_trailers_frame(buf.split().freeze()) {
+                            Ok(Some(trailers)) => {
+                                self.project().trailers.replace(trailers);
                             }
+                            Err(e) => return Poll::Ready(Some(Err(e))),
+                            _ => {}
+                        }
+
+                        if msg_buf.has_remaining() {
+                            Poll::Ready(Some(Ok(msg_buf)))
+                        } else {
+                            Poll::Ready(None)
                         }
-                        FindTrailers::IncompleteBuf => continue,
-                        FindTrailers::Done => Poll::Ready(Some(Ok(buf.split().freeze()))),
                     }
-                } else {
-                    Poll::Ready(buf)
+                    FindTrailers::IncompleteBuf => continue,
+                    FindTrailers::Done(len) => Poll::Ready(Some(Ok(buf.split_to(len).freeze()))),
                 };
             }
         }
@@ -421,21 +426,25 @@ fn make_trailers_frame(trailers: HeaderMap) -> Vec<u8> {
 /// its location in the original buf. If `None` is returned we did
 /// not find a trailers in this buffer either because its incomplete
 /// or the buffer jsut contained grpc message frames.
-fn find_trailers(buf: &[u8]) -> FindTrailers {
+fn find_trailers(buf: &[u8]) -> Result<FindTrailers, Status> {
     let mut len = 0;
     let mut temp_buf = &buf[..];
 
     loop {
         // To check each frame, there must be at least GRPC_HEADER_SIZE
         // amount of bytes available otherwise the buffer is incomplete.
         if temp_buf.is_empty() || temp_buf.len() < GRPC_HEADER_SIZE {
-            return FindTrailers::Done;
+            return Ok(FindTrailers::Done(len));
         }
 
         let header = temp_buf.get_u8();
 
         if header == GRPC_WEB_TRAILERS_BIT {
-            return FindTrailers::Trailer(len);
+            return Ok(FindTrailers::Trailer(len));
+        }
+
+        if !(header == 0 || header == 1) {
+            return Err(Status::internal("Invalid header bit {} expected 0 or 1"));
         }
 
         let msg_len = temp_buf.get_u32();
@@ -445,7 +454,7 @@ fn find_trailers(buf: &[u8]) -> FindTrailers {
         // If the msg len of a non-grpc-web trailer frame is larger than
         // the overall buffer we know within that buffer there are no trailers.
         if len > buf.len() {
-            return FindTrailers::IncompleteBuf;
+            return Ok(FindTrailers::IncompleteBuf);
         }
 
         temp_buf = &buf[len as usize..];
@@ -456,11 +465,13 @@ fn find_trailers(buf: &[u8]) -> FindTrailers {
 enum FindTrailers {
     Trailer(usize),
     IncompleteBuf,
-    Done,
+    Done(usize),
 }
 
 #[cfg(test)]
 mod tests {
+    use tonic::Code;
+
     use super::*;
 
     #[test]
@@ -507,7 +518,7 @@ mod tests {
             128, 0, 0, 0, 15, 103, 114, 112, 99, 45, 115, 116, 97, 116, 117, 115, 58, 48, 13, 10,
         ];
 
-        let out = find_trailers(&buf[..]);
+        let out = find_trailers(&buf[..]).unwrap();
 
         assert_eq!(out, FindTrailers::Trailer(0));
     }
@@ -524,7 +535,7 @@ mod tests {
             15, 103, 114, 112, 99, 45, 115, 116, 97, 116, 117, 115, 58, 48, 13, 10,
         ];
 
-        let out = find_trailers(&buf[..]);
+        let out = find_trailers(&buf[..]).unwrap();
 
         assert_eq!(out, FindTrailers::Trailer(81));
 
@@ -566,8 +577,17 @@ mod tests {
             105, 116, 116, 101, 110, 32, 98, 121, 32,
         ];
 
-        let out = find_trailers(&buf[..]);
+        let out = find_trailers(&buf[..]).unwrap();
 
         assert_eq!(out, FindTrailers::IncompleteBuf);
     }
+
+    #[test]
+    #[ignore]
+    fn find_trailers_buffered_incomplete_buf_bug() {
+        let buf = std::fs::read("tests/incomplete-buf-bug.bin").unwrap();
+        let out = find_trailers(&buf[..]).unwrap_err();
+
+        assert_eq!(out.code(), Code::Internal);
+    }
 }

tonic/src/codec/decode.rs

@@ -253,7 +253,7 @@ impl StreamingInner {
         } else {
             // FIXME: improve buf usage.
             if self.buf.has_remaining() {
-                trace!("unexpected EOF decoding stream");
+                trace!("unexpected EOF decoding stream, state: {:?}", self.state);
                 Err(Status::new(
                     Code::Internal,
                     "Unexpected EOF decoding stream.".to_string(),