diff --git a/talpid-tunnel-config-client/src/lib.rs b/talpid-tunnel-config-client/src/lib.rs index f09cd32074af..3a1b25bb4656 100644 --- a/talpid-tunnel-config-client/src/lib.rs +++ b/talpid-tunnel-config-client/src/lib.rs @@ -5,7 +5,6 @@ use std::net::SocketAddr; #[cfg(not(target_os = "ios"))] use std::net::{IpAddr, Ipv4Addr}; use talpid_types::net::wireguard::{PresharedKey, PublicKey}; -use tokio::io::{AsyncRead, AsyncWrite}; use tonic::transport::Channel; #[cfg(not(target_os = "ios"))] use tonic::transport::Endpoint; @@ -279,6 +278,8 @@ fn xor_assign(dst: &mut [u8; 32], src: &[u8; 32]) { /// value has been speficically lowered, to avoid MTU issues. See the `socket` module. #[cfg(not(target_os = "ios"))] async fn connect_relay_config_client(ip: Ipv4Addr) -> Result { + use hyper_util::rt::tokio::TokioIo; + let endpoint = Endpoint::from_static("tcp://0.0.0.0:0"); let addr = SocketAddr::new(IpAddr::V4(ip), CONFIG_SERVICE_PORT); @@ -286,13 +287,13 @@ async fn connect_relay_config_client(ip: Ipv4Addr) -> Result(hyper_util::rt::tokio::TokioIo::new(sniffer)) + Ok::<_, std::io::Error>(TokioIo::new(sniffer)) })) .await .map_err(Error::GrpcConnectError)?; @@ -300,63 +301,68 @@ async fn connect_relay_config_client(ip: Ipv4Addr) -> Result { - s: S, - rx_bytes: usize, - tx_bytes: usize, - start_time: std::time::Instant, -} - -impl Drop for SocketSniffer { - fn drop(&mut self) { - let duration = self.start_time.elapsed(); - log::debug!( - "Tunnel config client connection ended. RX: {} bytes, TX: {} bytes, duration: {} s", - self.rx_bytes, - self.tx_bytes, - duration.as_secs() - ); +mod socket_sniffer { + pub struct SocketSniffer { + pub s: S, + pub rx_bytes: usize, + pub tx_bytes: usize, + pub start_time: std::time::Instant, } -} + use std::{ + io, + pin::Pin, + task::{Context, Poll}, + }; + + use tokio::io::AsyncWrite; -impl AsyncRead for SocketSniffer { - fn poll_read( - mut self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - buf: &mut tokio::io::ReadBuf<'_>, - ) -> std::task::Poll> { - let bytes = std::task::ready!(std::pin::Pin::new(&mut self.s).poll_read(cx, buf)); - if bytes.is_ok() { - self.rx_bytes += buf.filled().len(); + use tokio::io::{AsyncRead, ReadBuf}; + + impl Drop for SocketSniffer { + fn drop(&mut self) { + let duration = self.start_time.elapsed(); + log::debug!( + "Tunnel config client connection ended. RX: {} bytes, TX: {} bytes, duration: {} s", + self.rx_bytes, + self.tx_bytes, + duration.as_secs() + ); } - std::task::Poll::Ready(bytes) } -} -impl AsyncWrite for SocketSniffer { - fn poll_write( - mut self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - buf: &[u8], - ) -> std::task::Poll> { - let bytes = std::task::ready!(std::pin::Pin::new(&mut self.s).poll_write(cx, buf)); - if bytes.is_ok() { - self.tx_bytes += buf.len(); + impl AsyncRead for SocketSniffer { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let bytes = std::task::ready!(Pin::new(&mut self.s).poll_read(cx, buf)); + if bytes.is_ok() { + self.rx_bytes += buf.filled().len(); + } + Poll::Ready(bytes) } - std::task::Poll::Ready(bytes) } - fn poll_flush( - mut self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - std::pin::Pin::new(&mut self.s).poll_flush(cx) - } + impl AsyncWrite for SocketSniffer { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let bytes = std::task::ready!(Pin::new(&mut self.s).poll_write(cx, buf)); + if bytes.is_ok() { + self.tx_bytes += buf.len(); + } + Poll::Ready(bytes) + } - fn poll_shutdown( - mut self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - std::pin::Pin::new(&mut self.s).poll_shutdown(cx) + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.s).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.s).poll_shutdown(cx) + } } }