Expose shadowsocks server client

ref #199

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;

src/relay/loadbalancing/server.rs

@@ -459,10 +459,9 @@ impl<S: ServerData + 'static> PingBalancer<S> {
 
         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?;
 

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<ServerClient> {
+    /// Connect to target address via shadowsocks' server
+    pub async fn connect(context: SharedContext, addr: &Address, svr_cfg: &ServerConfig) -> io::Result<ServerClient> {
         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<Result<usize, io::Error>> {
+        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<Result<usize, io::Error>> {
+        Pin::new(&mut self.stream).poll_write(cx, buf)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Result<(), io::Error>> {
+        Pin::new(&mut self.stream).poll_flush(cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Result<(), io::Error>> {
+        Pin::new(&mut self.stream).poll_shutdown(cx)
+    }
+}