diff --git a/h3/src/client.rs b/h3/src/client.rs index 83b745e3..f436bf16 100644 --- a/h3/src/client.rs +++ b/h3/src/client.rs @@ -687,6 +687,11 @@ where self.inner.recv_data().await } + /// Poll for data sent from the server. + pub fn poll_data(&mut self, cx: &mut Context<'_>) -> Poll, Error>> { + self.inner.poll_data(cx) + } + /// Receive an optional set of trailers for the response. pub async fn recv_trailers(&mut self) -> Result, Error> { let res = self.inner.recv_trailers().await; diff --git a/h3/src/connection.rs b/h3/src/connection.rs index bdd2ab77..09769137 100644 --- a/h3/src/connection.rs +++ b/h3/src/connection.rs @@ -622,47 +622,57 @@ impl RequestStream where S: quic::RecvStream, { - /// Receive some of the request body. - pub async fn recv_data(&mut self) -> Result, Error> { + pub fn poll_data(&mut self, cx: &mut Context<'_>) -> Poll, Error>> { if !self.stream.has_data() { - let frame = future::poll_fn(|cx| self.stream.poll_next(cx)) - .await - .map_err(|e| self.maybe_conn_err(e))?; - match frame { - Some(Frame::Data { .. }) => (), - Some(Frame::Headers(encoded)) => { - self.trailers = Some(encoded); - return Ok(None); - } + match self.stream.poll_next(cx) { + Poll::Ready(Ok(frame)) => { + match frame { + Some(Frame::Data { .. }) => (), + Some(Frame::Headers(encoded)) => { + self.trailers = Some(encoded); + return Poll::Ready(Ok(None)); + } - //= https://www.rfc-editor.org/rfc/rfc9114#section-4.1 - //# Receipt of an invalid sequence of frames MUST be treated as a - //# connection error of type H3_FRAME_UNEXPECTED. + //= https://www.rfc-editor.org/rfc/rfc9114#section-4.1 + //# Receipt of an invalid sequence of frames MUST be treated as a + //# connection error of type H3_FRAME_UNEXPECTED. - //= https://www.rfc-editor.org/rfc/rfc9114#section-7.2.3 - //# Receiving a - //# CANCEL_PUSH frame on a stream other than the control stream MUST be - //# treated as a connection error of type H3_FRAME_UNEXPECTED. + //= https://www.rfc-editor.org/rfc/rfc9114#section-7.2.3 + //# Receiving a + //# CANCEL_PUSH frame on a stream other than the control stream MUST be + //# treated as a connection error of type H3_FRAME_UNEXPECTED. - //= https://www.rfc-editor.org/rfc/rfc9114#section-7.2.4 - //# If an endpoint receives a SETTINGS frame on a different - //# stream, the endpoint MUST respond with a connection error of type - //# H3_FRAME_UNEXPECTED. - - //= https://www.rfc-editor.org/rfc/rfc9114#section-7.2.6 - //# A client MUST treat a GOAWAY frame on a stream other than - //# the control stream as a connection error of type H3_FRAME_UNEXPECTED. - - //= https://www.rfc-editor.org/rfc/rfc9114#section-7.2.7 - //# 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), + //= https://www.rfc-editor.org/rfc/rfc9114#section-7.2.4 + //# If an endpoint receives a SETTINGS frame on a different + //# stream, the endpoint MUST respond with a connection error of type + //# H3_FRAME_UNEXPECTED. + + //= https://www.rfc-editor.org/rfc/rfc9114#section-7.2.6 + //# A client MUST treat a GOAWAY frame on a stream other than + //# the control stream as a connection error of type H3_FRAME_UNEXPECTED. + + //= https://www.rfc-editor.org/rfc/rfc9114#section-7.2.7 + //# 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 Poll::Ready(Err(Code::H3_FRAME_UNEXPECTED.into())), + None => return Poll::Ready(Ok(None)), + } + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(self.maybe_conn_err(e))), + Poll::Pending => return Poll::Pending, } } + match self.stream.poll_data(cx) { + Poll::Ready(Ok(data)) => Poll::Ready(Ok(data)), + Poll::Ready(Err(e)) => Poll::Ready(Err(self.maybe_conn_err(e))), + Poll::Pending => Poll::Pending, + } + } - let data = future::poll_fn(|cx| self.stream.poll_data(cx)) + /// Receive some of the request body. + pub async fn recv_data(&mut self) -> Result, Error> { + let data = future::poll_fn(|cx| self.poll_data(cx)) .await .map_err(|e| self.maybe_conn_err(e))?; Ok(data) diff --git a/h3/src/server.rs b/h3/src/server.rs index b3ecf093..01c02a5a 100644 --- a/h3/src/server.rs +++ b/h3/src/server.rs @@ -674,6 +674,11 @@ where self.inner.recv_data().await } + /// Poll for data sent from the client + pub fn poll_data(&mut self, cx: &mut Context<'_>) -> Poll, Error>> { + self.inner.poll_data(cx) + } + /// Receive an optional set of trailers for the request pub async fn recv_trailers(&mut self) -> Result, Error> { self.inner.recv_trailers().await