Skip to content

Commit

Permalink
feat(client): Add connection capturing API to hyper-util
Browse files Browse the repository at this point in the history
  • Loading branch information
rcoh committed Apr 3, 2024
1 parent 16daef6 commit 986ba5a
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 1 deletion.
5 changes: 5 additions & 0 deletions src/client/legacy/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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");
Expand Down
3 changes: 3 additions & 0 deletions src/client/legacy/connect/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
33 changes: 32 additions & 1 deletion tests/legacy_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand Down Expand Up @@ -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());
}

0 comments on commit 986ba5a

Please sign in to comment.