From 66c36c4edb04d8f75ca66b9199546308fe089c0d Mon Sep 17 00:00:00 2001
From: Michael Rodler <mrodler@amazon.de>
Date: Tue, 23 May 2023 15:58:30 +0000
Subject: [PATCH] fix panic on receiving invalid headers frame by making the
 `take_request` function return a Result

Signed-off-by: Michael Rodler <mrodler@amazon.de>
Reviewed-by: Daniele Ahmed <ahmeddan@amazon.de>
---
 src/proto/streams/recv.rs      | 11 ++++++-----
 src/proto/streams/streams.rs   |  2 +-
 src/server.rs                  | 17 ++++++++++++-----
 tests/h2-tests/tests/server.rs | 32 ++++++++++++++++++++++++++++++++
 4 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/src/proto/streams/recv.rs b/src/proto/streams/recv.rs
index ec4db1c79..cd96dce2c 100644
--- a/src/proto/streams/recv.rs
+++ b/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))
+            }
         }
     }
 
diff --git a/src/proto/streams/streams.rs b/src/proto/streams/streams.rs
index dfc5c768b..d64e00970 100644
--- a/src/proto/streams/streams.rs
+++ b/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;
 
diff --git a/src/server.rs b/src/server.rs
index f1f4cf470..148cad517 100644
--- a/src/server.rs
+++ b/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
diff --git a/tests/h2-tests/tests/server.rs b/tests/h2-tests/tests/server.rs
index c8c1c9d1c..2637011ff 100644
--- a/tests/h2-tests/tests/server.rs
+++ b/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;
+}