fix(web): Buffer incomplete messages (#1528)

Previously, if we did not find trailers in the incoming message and the
message len was longer than the actual buffer we returned the buffer.
This would mean we would fail to parse incomplete message buffers. This
fixes `find_trailers` to support three return types, `Trailers(len)`,
`IncompleteBuf`, `Done`. Previously, there was only `Some(len)` and
`None`.

Author: LucioFranco

tonic-web/src/call.rs

@@ -242,33 +242,41 @@ where
         cx: &mut Context<'_>,
     ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
         if self.client && self.direction == Direction::Decode {
-            let buf = ready!(self.as_mut().poll_decode(cx));
-
-            return if let Some(Ok(mut buf)) = buf {
-                // We found some trailers so extract them since we
-                // want to return them via `poll_trailers`.
-                if let Some(len) = find_trailers(&buf[..]) {
-                    // Extract up to len of where the trailers are at
-                    let msg_buf = buf.copy_to_bytes(len);
-                    match decode_trailers_frame(buf) {
-                        Ok(Some(trailers)) => {
-                            self.project().trailers.replace(trailers);
+            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))),
+                                _ => {}
+                            }
+
+                            if msg_buf.has_remaining() {
+                                return Poll::Ready(Some(Ok(msg_buf)));
+                            } else {
+                                return Poll::Ready(None);
+                            }
                         }
-                        Err(e) => return Poll::Ready(Some(Err(e))),
-                        _ => {}
+                        FindTrailers::IncompleteBuf => continue,
+                        FindTrailers::Done => Poll::Ready(Some(Ok(buf.split().freeze()))),
                     }
-
-                    if msg_buf.has_remaining() {
-                        return Poll::Ready(Some(Ok(msg_buf)));
-                    } else {
-                        return Poll::Ready(None);
-                    }
-                }
-
-                Poll::Ready(Some(Ok(buf)))
-            } else {
-                Poll::Ready(buf)
-            };
+                } else {
+                    Poll::Ready(buf)
+                };
+            }
         }
 
         match self.direction {
@@ -413,21 +421,21 @@ 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]) -> Option<usize> {
+fn find_trailers(buf: &[u8]) -> FindTrailers {
     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 None;
+            return FindTrailers::Done;
         }
 
         let header = temp_buf.get_u8();
 
         if header == GRPC_WEB_TRAILERS_BIT {
-            return Some(len);
+            return FindTrailers::Trailer(len);
         }
 
         let msg_len = temp_buf.get_u32();
@@ -437,13 +445,20 @@ fn find_trailers(buf: &[u8]) -> Option<usize> {
         // 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 None;
+            return FindTrailers::IncompleteBuf;
         }
 
         temp_buf = &buf[len as usize..];
     }
 }
 
+#[derive(Debug, PartialEq, Eq)]
+enum FindTrailers {
+    Trailer(usize),
+    IncompleteBuf,
+    Done,
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -494,7 +509,7 @@ mod tests {
 
         let out = find_trailers(&buf[..]);
 
-        assert_eq!(out, Some(0));
+        assert_eq!(out, FindTrailers::Trailer(0));
     }
 
     #[test]
@@ -511,12 +526,48 @@ mod tests {
 
         let out = find_trailers(&buf[..]);
 
-        assert_eq!(out, Some(81));
+        assert_eq!(out, FindTrailers::Trailer(81));
 
         let trailers = decode_trailers_frame(Bytes::copy_from_slice(&buf[81..]))
             .unwrap()
             .unwrap();
         let status = trailers.get("grpc-status").unwrap();
         assert_eq!(status.to_str().unwrap(), "0")
     }
+
+    #[test]
+    fn find_trailers_buffered_incomplete_message() {
+        let buf = vec![
+            0, 0, 0, 9, 238, 10, 233, 19, 18, 230, 19, 10, 9, 10, 1, 120, 26, 4, 84, 69, 88, 84,
+            18, 60, 10, 58, 10, 56, 3, 0, 0, 0, 44, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32,
+            118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 101, 110, 32,
+            118, 105, 97, 32, 119, 114, 105, 116, 101, 32, 100, 101, 108, 101, 103, 97, 116, 105,
+            111, 110, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104,
+            105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116,
+            101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114,
+            101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0,
+            0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114,
+            105, 116, 116, 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100,
+            101, 100, 32, 114, 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0,
+            46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97,
+            115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98,
+            101, 100, 100, 101, 100, 32, 114, 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10,
+            58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117,
+            101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 98, 121, 32, 97, 110,
+            32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, 101, 112, 108, 105, 99, 97, 33, 18,
+            62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 118,
+            97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 98,
+            121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, 101, 112, 108,
+            105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104,
+            105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116,
+            101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114,
+            101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0,
+            0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114,
+            105, 116, 116, 101, 110, 32, 98, 121, 32,
+        ];
+
+        let out = find_trailers(&buf[..]);
+
+        assert_eq!(out, FindTrailers::IncompleteBuf);
+    }
 }