feat(client): Add connection capturing API to hyper-util

Author: rcoh

src/client/legacy/client.rs

@@ -18,6 +18,7 @@ use hyper::rt::Timer;
 use hyper::{body::Body, Method, Request, Response, Uri, Version};
 use tracing::{debug, trace, warn};
 
+use super::connect::capture::CaptureConnectionExtension;
 #[cfg(feature = "tokio")]
 use super::connect::HttpConnector;
 use super::connect::{Alpn, Connect, Connected, Connection};
@@ -265,6 +266,10 @@ where
     ) -> Result<Response<hyper::body::Incoming>, Error> {
         let mut pooled = self.connection_for(pool_key).await?;
 
+        req.extensions_mut()
+            .get_mut::<CaptureConnectionExtension>()
+            .map(|conn| conn.set(&pooled.conn_info));
+
         if pooled.is_http1() {
             if req.version() == Version::HTTP_2 {
                 warn!("Connection is HTTP/1, but request requires HTTP/2");

src/client/legacy/connect/mod.rs

@@ -74,6 +74,9 @@ pub mod dns;
 #[cfg(feature = "tokio")]
 mod http;
 
+pub(crate) mod capture;
+pub use capture::{capture_connection, CaptureConnection};
+
 pub use self::sealed::Connect;
 
 /// Describes a type returned by a connector.

tests/legacy_client.rs

@@ -18,7 +18,7 @@ use http_body_util::{Empty, Full, StreamBody};
 use hyper::body::Bytes;
 use hyper::body::Frame;
 use hyper::Request;
-use hyper_util::client::legacy::connect::HttpConnector;
+use hyper_util::client::legacy::connect::{capture_connection, HttpConnector};
 use hyper_util::client::legacy::Client;
 use hyper_util::rt::{TokioExecutor, TokioIo};
 
@@ -876,3 +876,34 @@ fn alpn_h2() {
     );
     drop(client);
 }
+
+#[test]
+fn capture_connection_on_client() {
+    let _ = pretty_env_logger::try_init();
+
+    let rt = runtime();
+    let connector = DebugConnector::new();
+
+    let client = Client::builder(TokioExecutor::new()).build(connector);
+
+    let server = TcpListener::bind("127.0.0.1:0").unwrap();
+    let addr = server.local_addr().unwrap();
+    thread::spawn(move || {
+        let mut sock = server.accept().unwrap().0;
+        //drop(server);
+        sock.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
+        sock.set_write_timeout(Some(Duration::from_secs(5)))
+            .unwrap();
+        let mut buf = [0; 4096];
+        sock.read(&mut buf).expect("read 1");
+        sock.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")
+            .expect("write 1");
+    });
+    let mut req = Request::builder()
+        .uri(&*format!("http://{}/a", addr))
+        .body(Empty::<Bytes>::new())
+        .unwrap();
+    let captured_conn = capture_connection(&mut req);
+    rt.block_on(client.request(req)).expect("200 OK");
+    assert!(captured_conn.connection_metadata().is_some());
+}