feat(http2): Implement Client-side CONNECT support over HTTP/2 (#2523)

Closes #2508

tests/client.rs

@@ -2261,14 +2261,16 @@ mod conn {
     use std::thread;
     use std::time::Duration;
 
+    use bytes::Buf;
     use futures_channel::oneshot;
     use futures_util::future::{self, poll_fn, FutureExt, TryFutureExt};
     use futures_util::StreamExt;
+    use hyper::upgrade::OnUpgrade;
     use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _, ReadBuf};
     use tokio::net::{TcpListener as TkTcpListener, TcpStream};
 
     use hyper::client::conn;
-    use hyper::{self, Body, Method, Request};
+    use hyper::{self, Body, Method, Request, Response, StatusCode};
 
     use super::{concat, s, support, tcp_connect, FutureHyperExt};
 
@@ -2984,6 +2986,125 @@ mod conn {
             .expect("client should be open");
     }
 
+    #[tokio::test]
+    async fn h2_connect() {
+        let _ = pretty_env_logger::try_init();
+
+        let listener = TkTcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
+            .await
+            .unwrap();
+        let addr = listener.local_addr().unwrap();
+
+        // Spawn an HTTP2 server that asks for bread and responds with baguette.
+        tokio::spawn(async move {
+            let sock = listener.accept().await.unwrap().0;
+            let mut h2 = h2::server::handshake(sock).await.unwrap();
+
+            let (req, mut respond) = h2.accept().await.unwrap().unwrap();
+            tokio::spawn(async move {
+                poll_fn(|cx| h2.poll_closed(cx)).await.unwrap();
+            });
+            assert_eq!(req.method(), Method::CONNECT);
+
+            let mut body = req.into_body();
+
+            let mut send_stream = respond.send_response(Response::default(), false).unwrap();
+
+            send_stream.send_data("Bread?".into(), true).unwrap();
+
+            let bytes = body.data().await.unwrap().unwrap();
+            assert_eq!(&bytes[..], b"Baguette!");
+            let _ = body.flow_control().release_capacity(bytes.len());
+
+            assert!(body.data().await.is_none());
+        });
+
+        let io = tcp_connect(&addr).await.expect("tcp connect");
+        let (mut client, conn) = conn::Builder::new()
+            .http2_only(true)
+            .handshake::<_, Body>(io)
+            .await
+            .expect("http handshake");
+
+        tokio::spawn(async move {
+            conn.await.expect("client conn shouldn't error");
+        });
+
+        let req = Request::connect("localhost")
+            .body(hyper::Body::empty())
+            .unwrap();
+        let res = client.send_request(req).await.expect("send_request");
+        assert_eq!(res.status(), StatusCode::OK);
+
+        let mut upgraded = hyper::upgrade::on(res).await.unwrap();
+
+        let mut vec = vec![];
+        upgraded.read_to_end(&mut vec).await.unwrap();
+        assert_eq!(s(&vec), "Bread?");
+
+        upgraded.write_all(b"Baguette!").await.unwrap();
+
+        upgraded.shutdown().await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn h2_connect_rejected() {
+        let _ = pretty_env_logger::try_init();
+
+        let listener = TkTcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
+            .await
+            .unwrap();
+        let addr = listener.local_addr().unwrap();
+        let (done_tx, done_rx) = oneshot::channel();
+
+        tokio::spawn(async move {
+            let sock = listener.accept().await.unwrap().0;
+            let mut h2 = h2::server::handshake(sock).await.unwrap();
+
+            let (req, mut respond) = h2.accept().await.unwrap().unwrap();
+            tokio::spawn(async move {
+                poll_fn(|cx| h2.poll_closed(cx)).await.unwrap();
+            });
+            assert_eq!(req.method(), Method::CONNECT);
+
+            let res = Response::builder().status(400).body(()).unwrap();
+            let mut send_stream = respond.send_response(res, false).unwrap();
+            send_stream
+                .send_data("No bread for you!".into(), true)
+                .unwrap();
+            done_rx.await.unwrap();
+        });
+
+        let io = tcp_connect(&addr).await.expect("tcp connect");
+        let (mut client, conn) = conn::Builder::new()
+            .http2_only(true)
+            .handshake::<_, Body>(io)
+            .await
+            .expect("http handshake");
+
+        tokio::spawn(async move {
+            conn.await.expect("client conn shouldn't error");
+        });
+
+        let req = Request::connect("localhost")
+            .body(hyper::Body::empty())
+            .unwrap();
+        let res = client.send_request(req).await.expect("send_request");
+        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
+        assert!(res.extensions().get::<OnUpgrade>().is_none());
+
+        let mut body = String::new();
+        hyper::body::aggregate(res.into_body())
+            .await
+            .unwrap()
+            .reader()
+            .read_to_string(&mut body)
+            .unwrap();
+        assert_eq!(body, "No bread for you!");
+
+        done_tx.send(()).unwrap();
+    }
+
     async fn drain_til_eof<T: AsyncRead + Unpin>(mut sock: T) -> io::Result<()> {
         let mut buf = [0u8; 1024];
         loop {

tests/server.rs

@@ -13,10 +13,13 @@ use std::task::{Context, Poll};
 use std::thread;
 use std::time::Duration;
 
+use bytes::Bytes;
 use futures_channel::oneshot;
 use futures_util::future::{self, Either, FutureExt, TryFutureExt};
 #[cfg(feature = "stream")]
 use futures_util::stream::StreamExt as _;
+use h2::client::SendRequest;
+use h2::{RecvStream, SendStream};
 use http::header::{HeaderName, HeaderValue};
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
 use tokio::net::{TcpListener, TcpStream as TkTcpStream};
@@ -1482,6 +1485,339 @@ async fn http_connect_new() {
     assert_eq!(s(&vec), "bar=foo");
 }
 
+#[tokio::test]
+async fn h2_connect() {
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    let _ = pretty_env_logger::try_init();
+    let listener = tcp_bind(&"127.0.0.1:0".parse().unwrap()).unwrap();
+    let addr = listener.local_addr().unwrap();
+    let conn = connect_async(addr).await;
+
+    let (h2, connection) = h2::client::handshake(conn).await.unwrap();
+    tokio::spawn(async move {
+        connection.await.unwrap();
+    });
+    let mut h2 = h2.ready().await.unwrap();
+
+    async fn connect_and_recv_bread(
+        h2: &mut SendRequest<Bytes>,
+    ) -> (RecvStream, SendStream<Bytes>) {
+        let request = Request::connect("localhost").body(()).unwrap();
+        let (response, send_stream) = h2.send_request(request, false).unwrap();
+        let response = response.await.unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+
+        let mut body = response.into_body();
+        let bytes = body.data().await.unwrap().unwrap();
+        assert_eq!(&bytes[..], b"Bread?");
+        let _ = body.flow_control().release_capacity(bytes.len());
+
+        (body, send_stream)
+    }
+
+    tokio::spawn(async move {
+        let (mut recv_stream, mut send_stream) = connect_and_recv_bread(&mut h2).await;
+
+        send_stream.send_data("Baguette!".into(), true).unwrap();
+
+        assert!(recv_stream.data().await.unwrap().unwrap().is_empty());
+    });
+
+    let svc = service_fn(move |req: Request<Body>| {
+        let on_upgrade = hyper::upgrade::on(req);
+
+        tokio::spawn(async move {
+            let mut upgraded = on_upgrade.await.expect("on_upgrade");
+            upgraded.write_all(b"Bread?").await.unwrap();
+
+            let mut vec = vec![];
+            upgraded.read_to_end(&mut vec).await.unwrap();
+            assert_eq!(s(&vec), "Baguette!");
+
+            upgraded.shutdown().await.unwrap();
+        });
+
+        future::ok::<_, hyper::Error>(
+            Response::builder()
+                .status(200)
+                .body(hyper::Body::empty())
+                .unwrap(),
+        )
+    });
+
+    let (socket, _) = listener.accept().await.unwrap();
+    Http::new()
+        .http2_only(true)
+        .serve_connection(socket, svc)
+        .with_upgrades()
+        .await
+        .unwrap();
+}
+
+#[tokio::test]
+async fn h2_connect_multiplex() {
+    use futures_util::stream::FuturesUnordered;
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    let _ = pretty_env_logger::try_init();
+    let listener = tcp_bind(&"127.0.0.1:0".parse().unwrap()).unwrap();
+    let addr = listener.local_addr().unwrap();
+    let conn = connect_async(addr).await;
+
+    let (h2, connection) = h2::client::handshake(conn).await.unwrap();
+    tokio::spawn(async move {
+        connection.await.unwrap();
+    });
+    let mut h2 = h2.ready().await.unwrap();
+
+    tokio::spawn(async move {
+        let mut streams = vec![];
+        for i in 0..80 {
+            let request = Request::connect(format!("localhost_{}", i % 4))
+                .body(())
+                .unwrap();
+            let (response, send_stream) = h2.send_request(request, false).unwrap();
+            streams.push((i, response, send_stream));
+        }
+
+        let futures = streams
+            .into_iter()
+            .map(|(i, response, mut send_stream)| async move {
+                if i % 4 == 0 {
+                    return;
+                }
+
+                let response = response.await.unwrap();
+                assert_eq!(response.status(), StatusCode::OK);
+
+                if i % 4 == 1 {
+                    return;
+                }
+
+                let mut body = response.into_body();
+                let bytes = body.data().await.unwrap().unwrap();
+                assert_eq!(&bytes[..], b"Bread?");
+                let _ = body.flow_control().release_capacity(bytes.len());
+
+                if i % 4 == 2 {
+                    return;
+                }
+
+                send_stream.send_data("Baguette!".into(), true).unwrap();
+
+                assert!(body.data().await.unwrap().unwrap().is_empty());
+            })
+            .collect::<FuturesUnordered<_>>();
+
+        futures.for_each(future::ready).await;
+    });
+
+    let svc = service_fn(move |req: Request<Body>| {
+        let authority = req.uri().authority().unwrap().to_string();
+        let on_upgrade = hyper::upgrade::on(req);
+
+        tokio::spawn(async move {
+            let upgrade_res = on_upgrade.await;
+            if authority == "localhost_0" {
+                assert!(upgrade_res.expect_err("upgrade cancelled").is_canceled());
+                return;
+            }
+            let mut upgraded = upgrade_res.expect("upgrade successful");
+
+            upgraded.write_all(b"Bread?").await.unwrap();
+
+            let mut vec = vec![];
+            let read_res = upgraded.read_to_end(&mut vec).await;
+
+            if authority == "localhost_1" || authority == "localhost_2" {
+                let err = read_res.expect_err("read failed");
+                assert_eq!(err.kind(), io::ErrorKind::Other);
+                assert_eq!(
+                    err.get_ref()
+                        .unwrap()
+                        .downcast_ref::<h2::Error>()
+                        .unwrap()
+                        .reason(),
+                    Some(h2::Reason::CANCEL),
+                );
+                return;
+            }
+
+            read_res.unwrap();
+            assert_eq!(s(&vec), "Baguette!");
+
+            upgraded.shutdown().await.unwrap();
+        });
+
+        future::ok::<_, hyper::Error>(
+            Response::builder()
+                .status(200)
+                .body(hyper::Body::empty())
+                .unwrap(),
+        )
+    });
+
+    let (socket, _) = listener.accept().await.unwrap();
+    Http::new()
+        .http2_only(true)
+        .serve_connection(socket, svc)
+        .with_upgrades()
+        .await
+        .unwrap();
+}
+
+#[tokio::test]
+async fn h2_connect_large_body() {
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    let _ = pretty_env_logger::try_init();
+    let listener = tcp_bind(&"127.0.0.1:0".parse().unwrap()).unwrap();
+    let addr = listener.local_addr().unwrap();
+    let conn = connect_async(addr).await;
+
+    let (h2, connection) = h2::client::handshake(conn).await.unwrap();
+    tokio::spawn(async move {
+        connection.await.unwrap();
+    });
+    let mut h2 = h2.ready().await.unwrap();
+
+    const NO_BREAD: &str = "All work and no bread makes nox a dull boy.\n";
+
+    async fn connect_and_recv_bread(
+        h2: &mut SendRequest<Bytes>,
+    ) -> (RecvStream, SendStream<Bytes>) {
+        let request = Request::connect("localhost").body(()).unwrap();
+        let (response, send_stream) = h2.send_request(request, false).unwrap();
+        let response = response.await.unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+
+        let mut body = response.into_body();
+        let bytes = body.data().await.unwrap().unwrap();
+        assert_eq!(&bytes[..], b"Bread?");
+        let _ = body.flow_control().release_capacity(bytes.len());
+
+        (body, send_stream)
+    }
+
+    tokio::spawn(async move {
+        let (mut recv_stream, mut send_stream) = connect_and_recv_bread(&mut h2).await;
+
+        let large_body = Bytes::from(NO_BREAD.repeat(9000));
+
+        send_stream.send_data(large_body.clone(), false).unwrap();
+        send_stream.send_data(large_body, true).unwrap();
+
+        assert!(recv_stream.data().await.unwrap().unwrap().is_empty());
+    });
+
+    let svc = service_fn(move |req: Request<Body>| {
+        let on_upgrade = hyper::upgrade::on(req);
+
+        tokio::spawn(async move {
+            let mut upgraded = on_upgrade.await.expect("on_upgrade");
+            upgraded.write_all(b"Bread?").await.unwrap();
+
+            let mut vec = vec![];
+            if upgraded.read_to_end(&mut vec).await.is_err() {
+                return;
+            }
+            assert_eq!(vec.len(), NO_BREAD.len() * 9000 * 2);
+
+            upgraded.shutdown().await.unwrap();
+        });
+
+        future::ok::<_, hyper::Error>(
+            Response::builder()
+                .status(200)
+                .body(hyper::Body::empty())
+                .unwrap(),
+        )
+    });
+
+    let (socket, _) = listener.accept().await.unwrap();
+    Http::new()
+        .http2_only(true)
+        .serve_connection(socket, svc)
+        .with_upgrades()
+        .await
+        .unwrap();
+}
+
+#[tokio::test]
+async fn h2_connect_empty_frames() {
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    let _ = pretty_env_logger::try_init();
+    let listener = tcp_bind(&"127.0.0.1:0".parse().unwrap()).unwrap();
+    let addr = listener.local_addr().unwrap();
+    let conn = connect_async(addr).await;
+
+    let (h2, connection) = h2::client::handshake(conn).await.unwrap();
+    tokio::spawn(async move {
+        connection.await.unwrap();
+    });
+    let mut h2 = h2.ready().await.unwrap();
+
+    async fn connect_and_recv_bread(
+        h2: &mut SendRequest<Bytes>,
+    ) -> (RecvStream, SendStream<Bytes>) {
+        let request = Request::connect("localhost").body(()).unwrap();
+        let (response, send_stream) = h2.send_request(request, false).unwrap();
+        let response = response.await.unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+
+        let mut body = response.into_body();
+        let bytes = body.data().await.unwrap().unwrap();
+        assert_eq!(&bytes[..], b"Bread?");
+        let _ = body.flow_control().release_capacity(bytes.len());
+
+        (body, send_stream)
+    }
+
+    tokio::spawn(async move {
+        let (mut recv_stream, mut send_stream) = connect_and_recv_bread(&mut h2).await;
+
+        send_stream.send_data("".into(), false).unwrap();
+        send_stream.send_data("".into(), false).unwrap();
+        send_stream.send_data("".into(), false).unwrap();
+        send_stream.send_data("Baguette!".into(), false).unwrap();
+        send_stream.send_data("".into(), true).unwrap();
+
+        assert!(recv_stream.data().await.unwrap().unwrap().is_empty());
+    });
+
+    let svc = service_fn(move |req: Request<Body>| {
+        let on_upgrade = hyper::upgrade::on(req);
+
+        tokio::spawn(async move {
+            let mut upgraded = on_upgrade.await.expect("on_upgrade");
+            upgraded.write_all(b"Bread?").await.unwrap();
+
+            let mut vec = vec![];
+            upgraded.read_to_end(&mut vec).await.unwrap();
+            assert_eq!(s(&vec), "Baguette!");
+
+            upgraded.shutdown().await.unwrap();
+        });
+
+        future::ok::<_, hyper::Error>(
+            Response::builder()
+                .status(200)
+                .body(hyper::Body::empty())
+                .unwrap(),
+        )
+    });
+
+    let (socket, _) = listener.accept().await.unwrap();
+    Http::new()
+        .http2_only(true)
+        .serve_connection(socket, svc)
+        .with_upgrades()
+        .await
+        .unwrap();
+}
+
 #[tokio::test]
 async fn parse_errors_send_4xx_response() {
     let listener = tcp_bind(&"127.0.0.1:0".parse().unwrap()).unwrap();