From 14d2d7a3fb663174cf5c7a6761ddbb781345628c Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Fri, 27 Dec 2024 08:45:08 -0500
Subject: [PATCH] fix decoding extra empty frame

---
 src/async_impl/decoder.rs | 63 +++++++++++++++++++--------------------
 1 file changed, 31 insertions(+), 32 deletions(-)

diff --git a/src/async_impl/decoder.rs b/src/async_impl/decoder.rs
index 96a27ac45..9a8e509ea 100644
--- 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",