discord non-data frames in Body Stream impl

Nugine · Mar 16, 2024 · 4315956 · 4315956
1 parent 71ecdf1
commit 4315956
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/crates/s3s/src/http/body.rs b/crates/s3s/src/http/body.rs
@@ -175,14 +175,15 @@ impl http_body::Body for Body {
 impl Stream for Body {
     type Item = Result<Bytes, StdError>;
 
-    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        match http_body::Body::poll_frame(self, cx) {
-            Poll::Ready(Some(Ok(f))) if f.is_data() => Poll::Ready(Some(Ok(f.into_data().expect("already checked")))),
-            Poll::Ready(Some(Ok(f))) if f.is_trailers() => Poll::Ready(None),
-            Poll::Ready(Some(Ok(_))) => Poll::Ready(Some(Err("unexpected BodyFrame".into()))),
-            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
-            Poll::Ready(None) => Poll::Ready(None),
-            Poll::Pending => Poll::Pending,
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        loop {
+            match std::task::ready!(http_body::Body::poll_frame(self.as_mut(), cx)?) {
+                Some(frame) => match frame.into_data() {
+                    Ok(data) => return Poll::Ready(Some(Ok(data))),
+                    Err(_frame) => continue,
+                },
+                None => return Poll::Ready(None),
+            };
         }
     }
 }