From 489d5fec001baee6d5fc29fab403cb00aa8d7988 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Thu, 10 Aug 2023 14:17:07 +0800
Subject: [PATCH] ConnectionInfo

---
 src/http.rs      | 16 ++++++++++------
 src/socks.rs     |  4 ++++
 src/tun2proxy.rs | 22 +++++++++++-----------
 3 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/src/http.rs b/src/http.rs
index d3ca650..05f72d4 100644
--- a/src/http.rs
+++ b/src/http.rs
@@ -8,7 +8,7 @@ use crate::{
 use base64::Engine;
 use httparse::Response;
 use smoltcp::wire::IpProtocol;
-use socks5_impl::protocol::{Address, UserKey};
+use socks5_impl::protocol::UserKey;
 use std::{
     cell::RefCell,
     collections::{hash_map::RandomState, HashMap, VecDeque},
@@ -52,7 +52,7 @@ pub struct HttpConnection {
     digest_state: Rc<RefCell<Option<DigestState>>>,
     before: bool,
     credentials: Option<UserKey>,
-    destination: Address,
+    info: ConnectionInfo,
 }
 
 static PROXY_AUTHENTICATE: &str = "Proxy-Authenticate";
@@ -80,7 +80,7 @@ impl HttpConnection {
             digest_state,
             before: false,
             credentials,
-            destination: info.dst.clone(),
+            info: info.clone(),
         };
 
         res.send_tunnel_request()?;
@@ -89,9 +89,9 @@ impl HttpConnection {
 
     fn send_tunnel_request(&mut self) -> Result<(), Error> {
         self.server_outbuf.extend(b"CONNECT ");
-        self.server_outbuf.extend(self.destination.to_string().as_bytes());
+        self.server_outbuf.extend(self.info.dst.to_string().as_bytes());
         self.server_outbuf.extend(b" HTTP/1.1\r\nHost: ");
-        self.server_outbuf.extend(self.destination.to_string().as_bytes());
+        self.server_outbuf.extend(self.info.dst.to_string().as_bytes());
         self.server_outbuf.extend(b"\r\n");
 
         self.send_auth_data(if self.digest_state.borrow().is_none() {
@@ -111,7 +111,7 @@ impl HttpConnection {
 
         match scheme {
             AuthenticationScheme::Digest => {
-                let uri = self.destination.to_string();
+                let uri = self.info.dst.to_string();
 
                 let context = digest_auth::AuthContext::new_with_method(
                     &credentials.username,
@@ -318,6 +318,10 @@ impl HttpConnection {
 }
 
 impl TcpProxy for HttpConnection {
+    fn get_connection_info(&self) -> &ConnectionInfo {
+        &self.info
+    }
+
     fn push_data(&mut self, event: IncomingDataEvent<'_>) -> Result<(), Error> {
         let direction = event.direction;
         let buffer = event.buffer;
diff --git a/src/socks.rs b/src/socks.rs
index 4d01125..02b893d 100644
--- a/src/socks.rs
+++ b/src/socks.rs
@@ -263,6 +263,10 @@ impl SocksProxyImpl {
 }
 
 impl TcpProxy for SocksProxyImpl {
+    fn get_connection_info(&self) -> &ConnectionInfo {
+        &self.info
+    }
+
     fn push_data(&mut self, event: IncomingDataEvent<'_>) -> Result<(), Error> {
         let direction = event.direction;
         let buffer = event.buffer;
diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs
index 0cabacf..33743b5 100644
--- a/src/tun2proxy.rs
+++ b/src/tun2proxy.rs
@@ -18,7 +18,7 @@ use std::{
     str::FromStr,
 };
 
-#[derive(Hash, Clone, Eq, PartialEq, Debug)]
+#[derive(Hash, Clone, Eq, PartialEq, PartialOrd, Ord, Debug)]
 pub(crate) struct ConnectionInfo {
     pub(crate) src: SocketAddr,
     pub(crate) dst: Address,
@@ -36,7 +36,6 @@ impl Default for ConnectionInfo {
 }
 
 impl ConnectionInfo {
-    #[allow(dead_code)]
     pub fn new(src: SocketAddr, dst: Address, protocol: IpProtocol) -> Self {
         Self { src, dst, protocol }
     }
@@ -133,11 +132,11 @@ fn connection_tuple(frame: &[u8]) -> Result<(ConnectionInfo, bool, usize, usize)
 
         let (ports, first_packet, payload_offset, payload_size) =
             get_transport_info(protocol, header_len, &frame[header_len..])?;
-        let info = ConnectionInfo {
-            src: SocketAddr::new(src_addr, ports.0),
-            dst: SocketAddr::new(dst_addr, ports.1).into(),
+        let info = ConnectionInfo::new(
+            SocketAddr::new(src_addr, ports.0),
+            SocketAddr::new(dst_addr, ports.1).into(),
             protocol,
-        };
+        );
         return Ok((info, first_packet, payload_offset, payload_size));
     }
 
@@ -154,11 +153,11 @@ fn connection_tuple(frame: &[u8]) -> Result<(ConnectionInfo, bool, usize, usize)
 
         let (ports, first_packet, payload_offset, payload_size) =
             get_transport_info(protocol, header_len, &frame[header_len..])?;
-        let info = ConnectionInfo {
-            src: SocketAddr::new(src_addr, ports.0),
-            dst: SocketAddr::new(dst_addr, ports.1).into(),
+        let info = ConnectionInfo::new(
+            SocketAddr::new(src_addr, ports.0),
+            SocketAddr::new(dst_addr, ports.1).into(),
             protocol,
-        };
+        );
         return Ok((info, first_packet, payload_offset, payload_size));
     }
     Err("Neither IPv6 nor IPv4 packet".into())
@@ -178,6 +177,7 @@ struct TcpConnectState {
 }
 
 pub(crate) trait TcpProxy {
+    fn get_connection_info(&self) -> &ConnectionInfo;
     fn push_data(&mut self, event: IncomingDataEvent<'_>) -> Result<(), Error>;
     fn consume_data(&mut self, dir: OutgoingDirection, size: usize);
     fn peek_data(&mut self, dir: OutgoingDirection) -> OutgoingDataEvent;
@@ -193,7 +193,7 @@ pub(crate) trait UdpProxy {
 
 pub(crate) trait ConnectionManager {
     fn handles_connection(&self, info: &ConnectionInfo) -> bool;
-    fn new_tcp_proxy(&self, info: &ConnectionInfo, udp_associate: bool) -> Result<Box<dyn TcpProxy>, Error>;
+    fn new_tcp_proxy(&self, info: &ConnectionInfo, udp_associate: bool) -> Result<Box<dyn TcpProxy>>;
     fn close_connection(&self, info: &ConnectionInfo);
     fn get_server_addr(&self) -> SocketAddr;
     fn get_credentials(&self) -> &Option<UserKey>;