From 431442735d4509f65d490733bc1a968e92162f80 Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Fri, 6 Oct 2017 10:54:49 -0700
Subject: [PATCH] reset streams when receiving invalid psuedo headers

---
 src/frame/headers.rs    |  5 +++++
 src/server.rs           | 12 ++++++------
 tests/server.rs         | 27 +++++++++++++++++++++++++++
 tests/support/frames.rs | 11 +++++++++++
 tests/support/mod.rs    |  1 +
 tests/support/util.rs   |  5 +++++
 6 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/src/frame/headers.rs b/src/frame/headers.rs
index f3dbe1e..ad92fba 100644
--- a/src/frame/headers.rs
+++ b/src/frame/headers.rs
@@ -211,6 +211,11 @@ impl Headers {
         (self.header_block.pseudo, self.header_block.fields)
     }
 
+    #[cfg(feature = "unstable")]
+    pub fn pseudo_mut(&mut self) -> &mut Pseudo {
+        &mut self.header_block.pseudo
+    }
+
     pub fn fields(&self) -> &HeaderMap {
         &self.header_block.fields
     }
diff --git a/src/server.rs b/src/server.rs
index df04226..507eea0 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -513,15 +513,15 @@ impl proto::Peer for Peer {
         let mut parts = uri::Parts::default();
 
         if let Some(scheme) = pseudo.scheme {
-            // TODO: Don't unwrap
-            parts.scheme = Some(uri::Scheme::from_shared(scheme.into_inner()).unwrap());
+            parts.scheme = Some(uri::Scheme::from_shared(scheme.into_inner())
+                .or_else(|_| malformed!())?);
         } else {
             malformed!();
         }
 
         if let Some(authority) = pseudo.authority {
-            // TODO: Don't unwrap
-            parts.authority = Some(uri::Authority::from_shared(authority.into_inner()).unwrap());
+            parts.authority = Some(uri::Authority::from_shared(authority.into_inner())
+                .or_else(|_| malformed!())?);
         }
 
         if let Some(path) = pseudo.path {
@@ -530,8 +530,8 @@ impl proto::Peer for Peer {
                 malformed!();
             }
 
-            // TODO: Don't unwrap
-            parts.path_and_query = Some(uri::PathAndQuery::from_shared(path.into_inner()).unwrap());
+            parts.path_and_query = Some(uri::PathAndQuery::from_shared(path.into_inner())
+                .or_else(|_| malformed!())?);
         }
 
         b.uri(parts);
diff --git a/tests/server.rs b/tests/server.rs
index 5530f04..2b3fbad 100644
--- a/tests/server.rs
+++ b/tests/server.rs
@@ -58,3 +58,30 @@ fn serve_request() {
 #[test]
 #[ignore]
 fn accept_with_pending_connections_after_socket_close() {}
+
+#[test]
+fn sent_invalid_authority() {
+    let _ = ::env_logger::init();
+    let (io, client) = mock::new();
+
+    let bad_auth = util::byte_str("not:a/good authority");
+    let mut bad_headers: frame::Headers = frames::headers(1)
+        .request("GET", "https://example.com/")
+        .eos()
+        .into();
+    bad_headers.pseudo_mut().authority = Some(bad_auth);
+
+    let client = client
+        .assert_server_handshake()
+        .unwrap()
+        .recv_settings()
+        .send_frame(bad_headers)
+        .recv_frame(frames::reset(1).protocol_error())
+        .close();
+
+    let srv = Server::handshake(io)
+        .expect("handshake")
+        .and_then(|srv| srv.into_future().unwrap());
+
+    srv.join(client).wait().expect("wait");
+}
diff --git a/tests/support/frames.rs b/tests/support/frames.rs
index 5a0aa9f..711161b 100644
--- a/tests/support/frames.rs
+++ b/tests/support/frames.rs
@@ -136,6 +136,12 @@ impl Mock<frame::Headers> {
     }
 }
 
+impl From<Mock<frame::Headers>> for frame::Headers {
+    fn from(src: Mock<frame::Headers>) -> Self {
+        src.0
+    }
+}
+
 impl From<Mock<frame::Headers>> for SendFrame {
     fn from(src: Mock<frame::Headers>) -> Self {
         Frame::Headers(src.0)
@@ -234,6 +240,11 @@ impl From<Mock<frame::GoAway>> for SendFrame {
 // ==== Reset helpers
 
 impl Mock<frame::Reset> {
+    pub fn protocol_error(self) -> Self {
+        let id = self.0.stream_id();
+        Mock(frame::Reset::new(id, frame::Reason::ProtocolError))
+    }
+
     pub fn flow_control(self) -> Self {
         let id = self.0.stream_id();
         Mock(frame::Reset::new(id, frame::Reason::FlowControlError))
diff --git a/tests/support/mod.rs b/tests/support/mod.rs
index f05a580..68ab24b 100644
--- a/tests/support/mod.rs
+++ b/tests/support/mod.rs
@@ -12,6 +12,7 @@ pub extern crate futures;
 pub extern crate h2;
 pub extern crate http;
 pub extern crate mock_io;
+pub extern crate string;
 pub extern crate tokio_io;
 
 // Kind of annoying, but we can't use macros from crates that aren't specified
diff --git a/tests/support/util.rs b/tests/support/util.rs
index 58b46a4..2e7cb49 100644
--- a/tests/support/util.rs
+++ b/tests/support/util.rs
@@ -1,8 +1,13 @@
 use h2::client;
 
+use super::string::{String, TryFrom};
 use bytes::Bytes;
 use futures::{Async, Future, Poll};
 
+pub fn byte_str(s: &str) -> String<Bytes> {
+    String::try_from(Bytes::from(s)).unwrap()
+}
+
 pub fn wait_for_capacity(stream: client::Stream<Bytes>, target: usize) -> WaitForCapacity {
     WaitForCapacity {
         stream: Some(stream),