diff --git a/src/lib.rs b/src/lib.rs index e92c8c1a8472..391ed936b12a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -97,7 +97,7 @@ pub use self::{ pub mod acl; pub mod config; -mod context; +pub mod context; pub mod crypto; pub mod plugin; pub mod relay; diff --git a/src/relay/loadbalancing/server.rs b/src/relay/loadbalancing/server.rs index 499ca17f5712..0692e9fb3f15 100644 --- a/src/relay/loadbalancing/server.rs +++ b/src/relay/loadbalancing/server.rs @@ -459,10 +459,9 @@ impl PingBalancer { let addr = Address::DomainNameAddress("dl.google.com".to_owned(), 80); - let TcpServerClient { mut stream } = - TcpServerClient::connect(stat.clone_context(), &addr, stat.server_config()).await?; + let mut stream = TcpServerClient::connect(stat.clone_context(), &addr, stat.server_config()).await?; stream.write_all(GET_BODY).await?; - stream.flush().await?; + let mut buf = [0u8; 1]; stream.read_exact(&mut buf).await?; diff --git a/src/relay/tcprelay/client.rs b/src/relay/tcprelay/client.rs index 77bff96137a2..a7f979e647f3 100644 --- a/src/relay/tcprelay/client.rs +++ b/src/relay/tcprelay/client.rs @@ -127,17 +127,35 @@ impl AsyncWrite for Socks5Client { } } -pub(crate) struct ServerClient { - pub stream: ProxyStream, +/// Shadowsocks' TCP client +pub struct ServerClient { + stream: ProxyStream, } impl ServerClient { - pub(crate) async fn connect( - context: SharedContext, - addr: &Address, - svr_cfg: &ServerConfig, - ) -> io::Result { + /// Connect to target address via shadowsocks' server + pub async fn connect(context: SharedContext, addr: &Address, svr_cfg: &ServerConfig) -> io::Result { let stream = ProxyStream::connect_proxied(context, svr_cfg, addr).await?; Ok(ServerClient { stream }) } } + +impl AsyncRead for ServerClient { + fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context, buf: &mut [u8]) -> Poll> { + Pin::new(&mut self.stream).poll_read(cx, buf) + } +} + +impl AsyncWrite for ServerClient { + fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context, buf: &[u8]) -> Poll> { + Pin::new(&mut self.stream).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context) -> Poll> { + Pin::new(&mut self.stream).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context) -> Poll> { + Pin::new(&mut self.stream).poll_shutdown(cx) + } +}