feat: allow to get peer address for any session stream

Author: link2xt

src/net/session.rs

@@ -1,62 +1,114 @@
+use anyhow::Result;
 use fast_socks5::client::Socks5Stream;
 use std::pin::Pin;
+use std::net::SocketAddr;
 use std::time::Duration;
 use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, BufStream, BufWriter};
 use tokio_io_timeout::TimeoutStream;
+use tokio::net::TcpStream;
 
 pub(crate) trait SessionStream:
     AsyncRead + AsyncWrite + Unpin + Send + Sync + std::fmt::Debug
 {
     /// Change the read timeout on the session stream.
     fn set_read_timeout(&mut self, timeout: Option<Duration>);
+
+    /// Returns the remote address that this stream is connected to.
+    ///
+    /// If the connection is proxied, returns `None`.
+    fn peer_addr(&self) -> Result<Option<SocketAddr>>;
 }
 
 impl SessionStream for Box<dyn SessionStream> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.as_mut().set_read_timeout(timeout);
     }
+
+    fn peer_addr(&self) -> Result<Option<SocketAddr>> {
+        let addr = self.as_ref().peer_addr()?;
+        Ok(addr)
+    }
 }
 impl<T: SessionStream> SessionStream for async_native_tls::TlsStream<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_mut().set_read_timeout(timeout);
     }
+
+    fn peer_addr(&self) -> Result<Option<SocketAddr>> {
+        let addr = self.get_ref().peer_addr()?;
+        Ok(addr)
+    }
 }
 impl<T: SessionStream> SessionStream for tokio_rustls::client::TlsStream<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_mut().0.set_read_timeout(timeout);
     }
+    
+    fn peer_addr(&self) -> Result<Option<SocketAddr>> {
+        let addr = self.get_ref().0.peer_addr()?;
+        Ok(addr)
+    }
 }
 impl<T: SessionStream> SessionStream for BufStream<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_mut().set_read_timeout(timeout);
     }
+
+    fn peer_addr(&self) -> Result<Option<SocketAddr>> {
+        let addr = self.get_ref().peer_addr()?;
+        Ok(addr)
+    }
 }
 impl<T: SessionStream> SessionStream for BufWriter<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_mut().set_read_timeout(timeout);
     }
+
+    fn peer_addr(&self) -> Result<Option<SocketAddr>> {
+        let addr = self.get_ref().peer_addr()?;
+        Ok(addr)
+    }
 }
-impl<T: AsyncRead + AsyncWrite + Send + Sync + std::fmt::Debug> SessionStream
-    for Pin<Box<TimeoutStream<T>>>
+impl SessionStream for Pin<Box<TimeoutStream<TcpStream>>>
 {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.as_mut().set_read_timeout_pinned(timeout);
     }
+
+    fn peer_addr(&self) -> Result<Option<SocketAddr>> {
+        let addr = self.get_ref().peer_addr()?;
+        Ok(Some(addr))
+    }
 }
 impl<T: SessionStream> SessionStream for Socks5Stream<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_socket_mut().set_read_timeout(timeout)
     }
+
+    fn peer_addr(&self) -> Result<Option<SocketAddr>> {
+        let addr = self.get_socket_ref().peer_addr()?;
+        Ok(addr)
+    }
 }
 impl<T: SessionStream> SessionStream for shadowsocks::ProxyClientStream<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_mut().set_read_timeout(timeout)
     }
+
+    fn peer_addr(&self) -> Result<Option<SocketAddr>> {
+        let addr = self.get_ref().peer_addr()?;
+        Ok(addr)
+    }
 }
 impl<T: SessionStream> SessionStream for async_imap::DeflateStream<T> {
     fn set_read_timeout(&mut self, timeout: Option<Duration>) {
         self.get_mut().set_read_timeout(timeout)
     }
+
+    fn peer_addr(&self) -> Result<Option<SocketAddr>> {
+        let addr = self.get_ref().peer_addr()?;
+        Ok(addr)
+    }
 }
 
 /// Session stream with a read buffer.