Make ConnectInfo work with ListenerExt::tap_io (#3059)

This commit is contained in:
Jonas Platte 2024-12-02 07:36:50 +00:00 committed by GitHub
parent 6d30c579e0
commit 9fab45af46
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 50 additions and 1 deletions

View file

@ -92,6 +92,17 @@ const _: () = {
*stream.remote_addr()
}
}
impl<'a, L, F> Connected<serve::IncomingStream<'a, serve::TapIo<L, F>>> for L::Addr
where
L: serve::Listener,
L::Addr: Clone + Sync + 'static,
F: FnMut(&mut L::Io) + Send + 'static,
{
fn connect_info(stream: serve::IncomingStream<'a, serve::TapIo<L, F>>) -> Self {
stream.remote_addr().clone()
}
}
};
impl Connected<SocketAddr> for SocketAddr {

View file

@ -434,6 +434,7 @@ mod tests {
extract::connect_info::Connected,
handler::{Handler, HandlerWithoutStateExt},
routing::get,
serve::ListenerExt,
Router,
};
@ -452,14 +453,29 @@ mod tests {
let addr = "0.0.0.0:0";
let tcp_nodelay_listener = || async {
TcpListener::bind(addr).await.unwrap().tap_io(|tcp_stream| {
if let Err(err) = tcp_stream.set_nodelay(true) {
eprintln!("failed to set TCP_NODELAY on incoming connection: {err:#}");
}
})
};
// router
serve(TcpListener::bind(addr).await.unwrap(), router.clone());
serve(tcp_nodelay_listener().await, router.clone())
.await
.unwrap();
serve(UnixListener::bind("").unwrap(), router.clone());
serve(
TcpListener::bind(addr).await.unwrap(),
router.clone().into_make_service(),
);
serve(
tcp_nodelay_listener().await,
router.clone().into_make_service(),
);
serve(
UnixListener::bind("").unwrap(),
router.clone().into_make_service(),
@ -471,6 +487,12 @@ mod tests {
.clone()
.into_make_service_with_connect_info::<std::net::SocketAddr>(),
);
serve(
tcp_nodelay_listener().await,
router
.clone()
.into_make_service_with_connect_info::<std::net::SocketAddr>(),
);
serve(
UnixListener::bind("").unwrap(),
router.into_make_service_with_connect_info::<UdsConnectInfo>(),
@ -478,12 +500,17 @@ mod tests {
// method router
serve(TcpListener::bind(addr).await.unwrap(), get(handler));
serve(tcp_nodelay_listener().await, get(handler));
serve(UnixListener::bind("").unwrap(), get(handler));
serve(
TcpListener::bind(addr).await.unwrap(),
get(handler).into_make_service(),
);
serve(
tcp_nodelay_listener().await,
get(handler).into_make_service(),
);
serve(
UnixListener::bind("").unwrap(),
get(handler).into_make_service(),
@ -493,6 +520,10 @@ mod tests {
TcpListener::bind(addr).await.unwrap(),
get(handler).into_make_service_with_connect_info::<std::net::SocketAddr>(),
);
serve(
tcp_nodelay_listener().await,
get(handler).into_make_service_with_connect_info::<std::net::SocketAddr>(),
);
serve(
UnixListener::bind("").unwrap(),
get(handler).into_make_service_with_connect_info::<UdsConnectInfo>(),
@ -503,24 +534,31 @@ mod tests {
TcpListener::bind(addr).await.unwrap(),
handler.into_service(),
);
serve(tcp_nodelay_listener().await, handler.into_service());
serve(UnixListener::bind("").unwrap(), handler.into_service());
serve(
TcpListener::bind(addr).await.unwrap(),
handler.with_state(()),
);
serve(tcp_nodelay_listener().await, handler.with_state(()));
serve(UnixListener::bind("").unwrap(), handler.with_state(()));
serve(
TcpListener::bind(addr).await.unwrap(),
handler.into_make_service(),
);
serve(tcp_nodelay_listener().await, handler.into_make_service());
serve(UnixListener::bind("").unwrap(), handler.into_make_service());
serve(
TcpListener::bind(addr).await.unwrap(),
handler.into_make_service_with_connect_info::<std::net::SocketAddr>(),
);
serve(
tcp_nodelay_listener().await,
handler.into_make_service_with_connect_info::<std::net::SocketAddr>(),
);
serve(
UnixListener::bind("").unwrap(),
handler.into_make_service_with_connect_info::<UdsConnectInfo>(),

View file

@ -102,7 +102,7 @@ impl<L: Listener> ListenerExt for L {}
/// Return type of [`ListenerExt::tap_io`].
///
/// See that method for details.
pub struct TapIo<L: Listener, F> {
pub struct TapIo<L, F> {
listener: L,
tap_fn: F,
}