fix decoding extra empty frame

seanmonstar · Dec 27, 2024 · 14d2d7a · 14d2d7a
1 parent 409cff3
commit 14d2d7a
Showing 1 changed file with 31 additions and 32 deletions.
diff --git a/src/async_impl/decoder.rs b/src/async_impl/decoder.rs
@@ -378,14 +378,7 @@ impl HttpBody for Decoder {
                     Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
                     None => {
                         // poll inner connection until EOF after gzip stream is finished
-                        let inner_stream = decoder.get_mut().get_mut().get_mut().get_mut();
-                        match futures_core::ready!(Pin::new(inner_stream).poll_next(cx)) {
-                            Some(Ok(_)) => Poll::Ready(Some(Err(crate::error::decode(
-                                "there are extra bytes after body has been decompressed",
-                            )))),
-                            Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
-                            None => Poll::Ready(None),
-                        }
+                        poll_inner_should_be_empty(decoder.get_mut().get_mut().get_mut().get_mut(), cx)
                     }
                 }
             }
@@ -396,14 +389,7 @@ impl HttpBody for Decoder {
                     Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
                     None => {
                         // poll inner connection until EOF after brotli stream is finished
-                        let inner_stream = decoder.get_mut().get_mut().get_mut().get_mut();
-                        match futures_core::ready!(Pin::new(inner_stream).poll_next(cx)) {
-                            Some(Ok(_)) => Poll::Ready(Some(Err(crate::error::decode(
-                                "there are extra bytes after body has been decompressed",
-                            )))),
-                            Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
-                            None => Poll::Ready(None),
-                        }
+                        poll_inner_should_be_empty(decoder.get_mut().get_mut().get_mut().get_mut(), cx)
                     }
                 }
             }
@@ -414,14 +400,7 @@ impl HttpBody for Decoder {
                     Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
                     None => {
                         // poll inner connection until EOF after zstd stream is finished
-                        let inner_stream = decoder.get_mut().get_mut().get_mut().get_mut();
-                        match futures_core::ready!(Pin::new(inner_stream).poll_next(cx)) {
-                            Some(Ok(_)) => Poll::Ready(Some(Err(crate::error::decode(
-                                "there are extra bytes after body has been decompressed",
-                            )))),
-                            Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
-                            None => Poll::Ready(None),
-                        }
+                        poll_inner_should_be_empty(decoder.get_mut().get_mut().get_mut().get_mut(), cx)
                     }
                 }
             }
@@ -432,14 +411,7 @@ impl HttpBody for Decoder {
                     Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
                     None => {
                         // poll inner connection until EOF after deflate stream is finished
-                        let inner_stream = decoder.get_mut().get_mut().get_mut().get_mut();
-                        match futures_core::ready!(Pin::new(inner_stream).poll_next(cx)) {
-                            Some(Ok(_)) => Poll::Ready(Some(Err(crate::error::decode(
-                                "there are extra bytes after body has been decompressed",
-                            )))),
-                            Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
-                            None => Poll::Ready(None),
-                        }
+                        poll_inner_should_be_empty(decoder.get_mut().get_mut().get_mut().get_mut(), cx)
                     }
                 }
             }
@@ -461,6 +433,33 @@ impl HttpBody for Decoder {
     }
 }
 
+#[cfg(any(
+    feature = "gzip",
+    feature = "zstd",
+    feature = "brotli",
+    feature = "deflate"
+))]
+fn poll_inner_should_be_empty(
+        inner: &mut PeekableIoStream,
+        cx: &mut Context,
+    ) -> Poll<Option<Result<Frame<Bytes>, crate::Error>>>
+{
+    // poll inner connection until EOF after deflate stream is finished
+    // loop in case of empty frames
+    let mut inner = Pin::new(inner);
+    loop {
+        match futures_core::ready!(inner.as_mut().poll_next(cx)) {
+            // ignore any empty frames
+            Some(Ok(bytes)) if bytes.is_empty() => continue,
+            Some(Ok(_)) => return Poll::Ready(Some(Err(crate::error::decode(
+                "there are extra bytes after body has been decompressed",
+            )))),
+            Some(Err(err)) => return Poll::Ready(Some(Err(crate::error::decode_io(err)))),
+            None => return Poll::Ready(None),
+        }
+    }
+}
+
 #[cfg(any(
     feature = "gzip",
     feature = "zstd",