fix panic on receiving invalid headers frame by making the take_request function return a Result

Signed-off-by: Michael Rodler <[email protected]>
Reviewed-by: Daniele Ahmed <[email protected]>

Author: Michael Rodler

src/proto/streams/recv.rs

@@ -251,14 +251,15 @@ impl Recv {
     }
 
     /// Called by the server to get the request
-    ///
-    /// TODO: Should this fn return `Result`?
-    pub fn take_request(&mut self, stream: &mut store::Ptr) -> Request<()> {
+    pub fn take_request(&mut self, stream: &mut store::Ptr) -> Result<Request<()>, proto::Error> {
         use super::peer::PollMessage::*;
 
         match stream.pending_recv.pop_front(&mut self.buffer) {
-            Some(Event::Headers(Server(request))) => request,
-            _ => panic!(),
+            Some(Event::Headers(Server(request))) => Ok(request),
+            _ => {
+                proto_err!(stream: "received invalid request; stream={:?}", stream.id);
+                Err(Error::library_reset(stream.id, Reason::PROTOCOL_ERROR))
+            }
         }
     }
 

src/proto/streams/streams.rs

@@ -1178,7 +1178,7 @@ impl<B> StreamRef<B> {
     /// # Panics
     ///
     /// This function panics if the request isn't present.
-    pub fn take_request(&self) -> Request<()> {
+    pub fn take_request(&self) -> Result<Request<()>, proto::Error> {
         let mut me = self.opaque.inner.lock().unwrap();
         let me = &mut *me;
 

src/server.rs

@@ -425,13 +425,20 @@ where
 
         if let Some(inner) = self.connection.next_incoming() {
             tracing::trace!("received incoming");
-            let (head, _) = inner.take_request().into_parts();
-            let body = RecvStream::new(FlowControl::new(inner.clone_to_opaque()));
+            match inner.take_request() {
+                Ok(req) => {
+                    let (head, _) = req.into_parts();
+                    let body = RecvStream::new(FlowControl::new(inner.clone_to_opaque()));
 
-            let request = Request::from_parts(head, body);
-            let respond = SendResponse { inner };
+                    let request = Request::from_parts(head, body);
+                    let respond = SendResponse { inner };
 
-            return Poll::Ready(Some(Ok((request, respond))));
+                    return Poll::Ready(Some(Ok((request, respond))));
+                }
+                Err(e) => {
+                    return Poll::Ready(Some(Err(e.into())));
+                }
+            }
         }
 
         Poll::Pending

tests/h2-tests/tests/server.rs

@@ -1378,3 +1378,35 @@ async fn reject_non_authority_target_on_connect_request() {
 
     join(client, srv).await;
 }
+
+#[tokio::test]
+async fn reject_response_headers_in_request() {
+    h2_support::trace_init!();
+
+    let (io, mut client) = mock::new();
+
+    let client = async move {
+        let _ = client.assert_server_handshake().await;
+
+        client.send_frame(frames::headers(1).response(128)).await;
+
+        // TODO: is CANCEL the right error code to expect here?
+        client.recv_frame(frames::reset(1).cancel()).await;
+    };
+
+    let srv = async move {
+        let builder = server::Builder::new();
+        let mut srv = builder.handshake::<_, Bytes>(io).await.expect("handshake");
+
+        let res = srv.next().await;
+        tracing::warn!("{:?}", res);
+        assert!(res.is_some());
+        assert!(res.unwrap().is_err());
+
+        poll_fn(move |cx| srv.poll_closed(cx))
+            .await
+            .expect("server");
+    };
+
+    join(client, srv).await;
+}