feat(h2): implement CONNECT support (fixes #2508)

Author: nox

Cargo.toml

@@ -31,7 +31,7 @@ http = "0.2"
 http-body = "0.4"
 httpdate = "1.0"
 httparse = "1.4"
-h2 = { version = "0.3", optional = true }
+h2 = { version = "0.3.3", optional = true }
 itoa = "0.4.1"
 tracing = { version = "0.1", default-features = false, features = ["std"] }
 pin-project = "1.0"

src/body/length.rs

@@ -3,6 +3,17 @@ use std::fmt;
 #[derive(Clone, Copy, PartialEq, Eq)]
 pub(crate) struct DecodedLength(u64);
 
+#[cfg(any(feature = "http1", feature = "http2"))]
+impl From<Option<u64>> for DecodedLength {
+    fn from(len: Option<u64>) -> Self {
+        len.and_then(|len| {
+            // If the length is u64::MAX, oh well, just reported chunked.
+            Self::checked_new(len).ok()
+        })
+        .unwrap_or(DecodedLength::CHUNKED)
+    }
+}
+
 #[cfg(any(feature = "http1", feature = "http2", test))]
 const MAX_LEN: u64 = std::u64::MAX - 2;
 

src/error.rs

@@ -90,7 +90,7 @@ pub(super) enum User {
     /// User tried to send a certain header in an unexpected context.
     ///
     /// For example, sending both `content-length` and `transfer-encoding`.
-    #[cfg(feature = "http1")]
+    #[cfg(any(feature = "http1", feature = "http2"))]
     #[cfg(feature = "server")]
     UnexpectedHeader,
     /// User tried to create a Request with bad version.
@@ -279,7 +279,7 @@ impl Error {
         Error::new(Kind::User(user))
     }
 
-    #[cfg(feature = "http1")]
+    #[cfg(any(feature = "http1", feature = "http2"))]
     #[cfg(feature = "server")]
     pub(super) fn new_user_header() -> Error {
         Error::new_user(User::UnexpectedHeader)
@@ -394,7 +394,7 @@ impl Error {
             Kind::User(User::MakeService) => "error from user's MakeService",
             #[cfg(any(feature = "http1", feature = "http2"))]
             Kind::User(User::Service) => "error from user's Service",
-            #[cfg(feature = "http1")]
+            #[cfg(any(feature = "http1", feature = "http2"))]
             #[cfg(feature = "server")]
             Kind::User(User::UnexpectedHeader) => "user sent unexpected header",
             #[cfg(any(feature = "http1", feature = "http2"))]

src/proto/h2/mod.rs

@@ -1,17 +1,20 @@
-use bytes::Buf;
-use h2::SendStream;
+use bytes::{Buf, Bytes};
+use h2::{RecvStream, SendStream};
 use http::header::{
     HeaderName, CONNECTION, PROXY_AUTHENTICATE, PROXY_AUTHORIZATION, TE, TRAILER,
     TRANSFER_ENCODING, UPGRADE,
 };
 use http::HeaderMap;
 use pin_project::pin_project;
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use std::error::Error as StdError;
-use std::io::IoSlice;
+use std::io::{self, Cursor, IoSlice};
+use std::task::Context;
 
 use crate::body::{DecodedLength, HttpBody};
 use crate::common::{task, Future, Pin, Poll};
 use crate::headers::content_length_parse_all;
+use crate::proto::h2::ping::Recorder;
 
 pub(crate) mod ping;
 
@@ -84,12 +87,7 @@ fn strip_connection_headers(headers: &mut HeaderMap, is_request: bool) {
 }
 
 fn decode_content_length(headers: &HeaderMap) -> DecodedLength {
-    if let Some(len) = content_length_parse_all(headers) {
-        // If the length is u64::MAX, oh well, just reported chunked.
-        DecodedLength::checked_new(len).unwrap_or_else(|_| DecodedLength::CHUNKED)
-    } else {
-        DecodedLength::CHUNKED
-    }
+    content_length_parse_all(headers).into()
 }
 
 // body adapters used by both Client and Server
@@ -172,7 +170,7 @@ where
                             is_eos,
                         );
 
-                        let buf = SendBuf(Some(chunk));
+                        let buf = SendBuf::Buf(chunk);
                         me.body_tx
                             .send_data(buf, is_eos)
                             .map_err(crate::Error::new_body_write)?;
@@ -243,32 +241,152 @@ impl<B: Buf> SendStreamExt for SendStream<SendBuf<B>> {
 
     fn send_eos_frame(&mut self) -> crate::Result<()> {
         trace!("send body eos");
-        self.send_data(SendBuf(None), true)
+        self.send_data(SendBuf::None, true)
             .map_err(crate::Error::new_body_write)
     }
 }
 
-struct SendBuf<B>(Option<B>);
+enum SendBuf<B> {
+    Buf(B),
+    Cursor(Cursor<Box<[u8]>>),
+    None,
+}
 
 impl<B: Buf> Buf for SendBuf<B> {
     #[inline]
     fn remaining(&self) -> usize {
-        self.0.as_ref().map(|b| b.remaining()).unwrap_or(0)
+        match *self {
+            Self::Buf(ref b) => b.remaining(),
+            Self::Cursor(ref c) => c.remaining(),
+            Self::None => 0,
+        }
     }
 
     #[inline]
     fn chunk(&self) -> &[u8] {
-        self.0.as_ref().map(|b| b.chunk()).unwrap_or(&[])
+        match *self {
+            Self::Buf(ref b) => b.chunk(),
+            Self::Cursor(ref c) => c.chunk(),
+            Self::None => &[],
+        }
     }
 
     #[inline]
     fn advance(&mut self, cnt: usize) {
-        if let Some(b) = self.0.as_mut() {
-            b.advance(cnt)
+        match *self {
+            Self::Buf(ref mut b) => b.advance(cnt),
+            Self::Cursor(ref mut c) => c.advance(cnt),
+            Self::None => {},
         }
     }
 
     fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
-        self.0.as_ref().map(|b| b.chunks_vectored(dst)).unwrap_or(0)
+        match *self {
+            Self::Buf(ref b) => b.chunks_vectored(dst),
+            Self::Cursor(ref c) => c.chunks_vectored(dst),
+            Self::None => 0,
+        }
+    }
+}
+
+struct H2Upgraded<B>
+where
+    B: Buf,
+{
+    ping: Recorder,
+    send_stream: SendStream<SendBuf<B>>,
+    recv_stream: RecvStream,
+    buf: Bytes,
+}
+
+impl<B> AsyncRead for H2Upgraded<B>
+where
+    B: Buf,
+{
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        read_buf: &mut ReadBuf<'_>,
+    ) -> Poll<Result<(), io::Error>> {
+        if self.buf.is_empty() {
+            self.buf = loop {
+                match ready!(self.recv_stream.poll_data(cx)) {
+                    None => return Poll::Ready(Ok(())),
+                    Some(Ok(buf)) if buf.is_empty() && !self.recv_stream.is_end_stream() => continue,
+                    Some(Ok(buf)) => {
+                        self.ping.record_data(buf.len());
+                        break buf;
+                    }
+                    Some(Err(e)) => {
+                        return Poll::Ready(Err(h2_to_io_error(e)));
+                    }
+                }
+            };
+        }
+        let cnt = std::cmp::min(self.buf.len(), read_buf.remaining());
+        read_buf.put_slice(&self.buf[..cnt]);
+        self.buf.advance(cnt);
+        let _ = self.recv_stream.flow_control().release_capacity(cnt);
+        Poll::Ready(Ok(()))
+    }
+}
+
+impl<B> AsyncWrite for H2Upgraded<B>
+where
+    B: Buf,
+{
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, io::Error>> {
+        if buf.is_empty() {
+            return Poll::Ready(Ok(0));
+        }
+        // FIXME(nox): PipeToSendStream does some weird stuff, first reserving
+        // one byte and then polling reset if the capacity is 0, should we do
+        // that here too? Should we poll reset somewhere?
+        self.send_stream.reserve_capacity(buf.len());
+        Poll::Ready(match ready!(self.send_stream.poll_capacity(cx)) {
+            None => Ok(0),
+            Some(Ok(cnt)) => self.write(&buf[..cnt], false).map(|()| cnt),
+            Some(Err(e)) => {
+                // FIXME(nox): Should all H2 errors be returned as is with a
+                // ErrorKind::Other, or should some be special-cased, say for
+                // example, CANCEL?
+                Err(h2_to_io_error(e))
+            },
+        })
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_shutdown(
+        mut self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<Result<(), io::Error>> {
+        Poll::Ready(self.write(&[], true))
+    }
+}
+
+impl<B> H2Upgraded<B>
+where
+    B: Buf,
+{
+    fn write(&mut self, buf: &[u8], end_of_stream: bool) -> Result<(), io::Error> {
+        let send_buf = SendBuf::Cursor(Cursor::new(buf.into()));
+        self.send_stream
+            .send_data(send_buf, end_of_stream)
+            .map_err(h2_to_io_error)
+    }
+}
+
+fn h2_to_io_error(e: h2::Error) -> io::Error {
+    if e.is_io() {
+        e.into_io().unwrap()
+    } else {
+        io::Error::new(io::ErrorKind::Other, e)
     }
 }

src/proto/h2/server.rs

@@ -3,19 +3,24 @@ use std::marker::Unpin;
 #[cfg(feature = "runtime")]
 use std::time::Duration;
 
+use bytes::Bytes;
 use h2::server::{Connection, Handshake, SendResponse};
-use h2::Reason;
+use h2::{Reason, RecvStream};
+use http::{Method, Request};
 use pin_project::pin_project;
 use tokio::io::{AsyncRead, AsyncWrite};
 
-use super::{decode_content_length, ping, PipeToSendStream, SendBuf};
+use super::{ping, PipeToSendStream, SendBuf};
 use crate::body::HttpBody;
 use crate::common::exec::ConnStreamExec;
 use crate::common::{date, task, Future, Pin, Poll};
 use crate::headers;
+use crate::proto::h2::ping::Recorder;
+use crate::proto::h2::H2Upgraded;
 use crate::proto::Dispatched;
 use crate::service::HttpService;
 
+use crate::upgrade::{OnUpgrade, Pending, Upgraded};
 use crate::{Body, Response};
 
 // Our defaults are chosen for the "majority" case, which usually are not
@@ -257,9 +262,9 @@ where
 
                 // When the service is ready, accepts an incoming request.
                 match ready!(self.conn.poll_accept(cx)) {
-                    Some(Ok((req, respond))) => {
+                    Some(Ok((req, mut respond))) => {
                         trace!("incoming request");
-                        let content_length = decode_content_length(req.headers());
+                        let content_length = headers::content_length_parse_all(req.headers());
                         let ping = self
                             .ping
                             .as_ref()
@@ -269,8 +274,32 @@ where
                         // Record the headers received
                         ping.record_non_data();
 
-                        let req = req.map(|stream| crate::Body::h2(stream, content_length, ping));
-                        let fut = H2Stream::new(service.call(req), respond);
+                        let is_connect = req.method() == Method::CONNECT;
+                        let (mut parts, stream) = req.into_parts();
+                        let (req, connect_parts) = if !is_connect {
+                            (
+                                Request::from_parts(
+                                    parts,
+                                    crate::Body::h2(stream, content_length.into(), ping),
+                                ),
+                                None,
+                            )
+                        } else {
+                            if content_length.map_or(false, |len| len != 0) {
+                                warn!("h2 connect request with non-zero body not supported");
+                                respond.send_reset(h2::Reason::INTERNAL_ERROR);
+                                return Poll::Ready(Ok(()));
+                            }
+                            let (pending, upgrade) = crate::upgrade::pending();
+                            debug_assert!(parts.extensions.get::<OnUpgrade>().is_none());
+                            parts.extensions.insert(upgrade);
+                            (
+                                Request::from_parts(parts, crate::Body::empty()),
+                                Some((pending, ping, stream)),
+                            )
+                        };
+
+                        let fut = H2Stream::new(service.call(req), connect_parts, respond);
                         exec.execute_h2stream(fut);
                     }
                     Some(Err(e)) => {
@@ -333,18 +362,22 @@ enum H2StreamState<F, B>
 where
     B: HttpBody,
 {
-    Service(#[pin] F),
+    Service(#[pin] F, Option<(Pending, Recorder, RecvStream)>),
     Body(#[pin] PipeToSendStream<B>),
 }
 
 impl<F, B> H2Stream<F, B>
 where
     B: HttpBody,
 {
-    fn new(fut: F, respond: SendResponse<SendBuf<B::Data>>) -> H2Stream<F, B> {
+    fn new(
+        fut: F,
+        connect_parts: Option<(Pending, Recorder, RecvStream)>,
+        respond: SendResponse<SendBuf<B::Data>>,
+    ) -> H2Stream<F, B> {
         H2Stream {
             reply: respond,
-            state: H2StreamState::Service(fut),
+            state: H2StreamState::Service(fut, connect_parts),
         }
     }
 }
@@ -374,7 +407,7 @@ where
         let mut me = self.project();
         loop {
             let next = match me.state.as_mut().project() {
-                H2StreamStateProj::Service(h) => {
+                H2StreamStateProj::Service(h, connect_parts) => {
                     let res = match h.poll(cx) {
                         Poll::Ready(Ok(r)) => r,
                         Poll::Pending => {
@@ -405,6 +438,27 @@ where
                         .entry(::http::header::DATE)
                         .or_insert_with(date::update_and_header_value);
 
+                    if let Some((pending, ping, recv_stream)) = connect_parts.take() {
+                        if res.status().is_success() {
+                            if headers::content_length_parse_all(res.headers()).map_or(false, |len| len != 0) {
+                                warn!("h2 successful response to CONNECT request with body not supported");
+                                me.reply.send_reset(h2::Reason::INTERNAL_ERROR);
+                                return Poll::Ready(Err(crate::Error::new_user_header()));
+                            }
+                            let send_stream = reply!(me, res, false);
+                            pending.fulfill(Upgraded::new(
+                                H2Upgraded {
+                                    ping,
+                                    recv_stream,
+                                    send_stream,
+                                    buf: Bytes::new(),
+                                },
+                                Bytes::new(),
+                            ));
+                            return Poll::Ready(Ok(()));
+                        }
+                    }
+
                     // automatically set Content-Length from body...
                     if let Some(len) = body.size_hint().exact() {
                         headers::set_content_length_if_missing(res.headers_mut(), len);