From 49c808e1446ba10c250357b5da5bf4d3935d5b4c Mon Sep 17 00:00:00 2001
From: Sergio Benitez <sb@sergio.bz>
Date: Mon, 12 Feb 2024 15:47:18 -0800
Subject: [PATCH] add 'Connection:poll_recv_data()'

---
 h3/src/connection.rs | 30 +++++++++++++++++++-----------
 h3/src/server.rs     |  8 ++++++++
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/h3/src/connection.rs b/h3/src/connection.rs
index bdd2ab77..f05ef0bc 100644
--- a/h3/src/connection.rs
+++ b/h3/src/connection.rs
@@ -623,16 +623,21 @@ where
     S: quic::RecvStream,
 {
     /// Receive some of the request body.
-    pub async fn recv_data(&mut self) -> Result<Option<impl Buf>, Error> {
+    pub fn poll_recv_data(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<Option<impl Buf>, Error>> {
         if !self.stream.has_data() {
-            let frame = future::poll_fn(|cx| self.stream.poll_next(cx))
-                .await
+            let frame = self
+                .stream
+                .poll_next(cx)
                 .map_err(|e| self.maybe_conn_err(e))?;
-            match frame {
+
+            match ready!(frame) {
                 Some(Frame::Data { .. }) => (),
                 Some(Frame::Headers(encoded)) => {
                     self.trailers = Some(encoded);
-                    return Ok(None);
+                    return Poll::Ready(Ok(None));
                 }
 
                 //= https://www.rfc-editor.org/rfc/rfc9114#section-4.1
@@ -657,15 +662,18 @@ where
                 //# The MAX_PUSH_ID frame is always sent on the control stream.  Receipt
                 //# of a MAX_PUSH_ID frame on any other stream MUST be treated as a
                 //# connection error of type H3_FRAME_UNEXPECTED.
-                Some(_) => return Err(Code::H3_FRAME_UNEXPECTED.into()),
-                None => return Ok(None),
+                Some(_) => return Poll::Ready(Err(Code::H3_FRAME_UNEXPECTED.into())),
+                None => return Poll::Ready(Ok(None)),
             }
         }
 
-        let data = future::poll_fn(|cx| self.stream.poll_data(cx))
-            .await
-            .map_err(|e| self.maybe_conn_err(e))?;
-        Ok(data)
+        self.stream
+            .poll_data(cx)
+            .map_err(|e| self.maybe_conn_err(e))
+    }
+    /// Receive some of the request body.
+    pub async fn recv_data(&mut self) -> Result<Option<impl Buf>, Error> {
+        future::poll_fn(|cx| self.poll_recv_data(cx)).await
     }
 
     /// Receive trailers
diff --git a/h3/src/server.rs b/h3/src/server.rs
index b3ecf093..0963430c 100644
--- a/h3/src/server.rs
+++ b/h3/src/server.rs
@@ -674,6 +674,14 @@ where
         self.inner.recv_data().await
     }
 
+    /// Poll for data sent from the client
+    pub fn poll_recv_data(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<Option<impl Buf>, Error>> {
+        self.inner.poll_recv_data(cx)
+    }
+
     /// Receive an optional set of trailers for the request
     pub async fn recv_trailers(&mut self) -> Result<Option<HeaderMap>, Error> {
         self.inner.recv_trailers().await