From 46ca342aba1a5ba285645a289916a5479f4102f1 Mon Sep 17 00:00:00 2001 From: ssrlive <30760636+ssrlive@users.noreply.github.com> Date: Fri, 11 Aug 2023 15:33:04 +0800 Subject: [PATCH] connection_managers renamed to connection_manager --- src/http.rs | 4 ---- src/lib.rs | 2 +- src/socks.rs | 8 ------- src/tun2proxy.rs | 62 +++++++++++++++--------------------------------- 4 files changed, 20 insertions(+), 56 deletions(-) diff --git a/src/http.rs b/src/http.rs index 05f72d4..880e3da 100644 --- a/src/http.rs +++ b/src/http.rs @@ -391,10 +391,6 @@ pub(crate) struct HttpManager { } impl ConnectionManager for HttpManager { - fn handles_connection(&self, info: &ConnectionInfo) -> bool { - info.protocol == IpProtocol::Tcp - } - fn new_tcp_proxy(&self, info: &ConnectionInfo, _: bool) -> Result, Error> { if info.protocol != IpProtocol::Tcp { return Err("Invalid protocol".into()); diff --git a/src/lib.rs b/src/lib.rs index 01ea7ba..bdbf7a0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -130,7 +130,7 @@ pub fn tun_to_proxy<'a>( ProxyType::Socks5 => Rc::new(SocksProxyManager::new(server, Version::V5, credentials)) as Rc, ProxyType::Http => Rc::new(HttpManager::new(server, credentials)) as Rc, }; - ttp.add_connection_manager(mgr); + ttp.set_connection_manager(Some(mgr)); Ok(ttp) } diff --git a/src/socks.rs b/src/socks.rs index 02b893d..7847025 100644 --- a/src/socks.rs +++ b/src/socks.rs @@ -5,7 +5,6 @@ use crate::{ OutgoingDirection, TcpProxy, }, }; -use smoltcp::wire::IpProtocol; use socks5_impl::protocol::{self, handshake, password_method, Address, AuthMethod, StreamOperation, UserKey, Version}; use std::{collections::VecDeque, convert::TryFrom, net::SocketAddr}; @@ -336,14 +335,7 @@ pub(crate) struct SocksProxyManager { } impl ConnectionManager for SocksProxyManager { - fn handles_connection(&self, info: &ConnectionInfo) -> bool { - info.protocol == IpProtocol::Tcp - } - fn new_tcp_proxy(&self, info: &ConnectionInfo, udp_associate: bool) -> Result> { - if info.protocol != IpProtocol::Tcp { - return Err("Invalid protocol".into()); - } use socks5_impl::protocol::Command::{Connect, UdpAssociate}; let command = if udp_associate { UdpAssociate } else { Connect }; Ok(Box::new(SocksProxyImpl::new( diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index 4bdf1a8..83b57c3 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -192,7 +192,6 @@ 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>; fn close_connection(&self, info: &ConnectionInfo); fn get_server_addr(&self) -> SocketAddr; @@ -207,7 +206,7 @@ pub struct TunToProxy<'a> { poll: Poll, iface: Interface, connection_map: HashMap, - connection_managers: Vec>, + connection_manager: Option>, next_token: usize, sockets: SocketSet<'a>, device: VirtualTunDevice, @@ -255,7 +254,7 @@ impl<'a> TunToProxy<'a> { iface, connection_map: HashMap::default(), next_token: usize::from(EXIT_TOKEN) + 1, - connection_managers: Vec::default(), + connection_manager: None, sockets: SocketSet::new([]), device, options, @@ -272,8 +271,8 @@ impl<'a> TunToProxy<'a> { token } - pub(crate) fn add_connection_manager(&mut self, manager: Rc) { - self.connection_managers.push(manager); + pub(crate) fn set_connection_manager(&mut self, manager: Option>) { + self.connection_manager = manager; } /// Read data from virtual device (remote server) and inject it into tun interface. @@ -320,13 +319,8 @@ impl<'a> TunToProxy<'a> { Ok(()) } - fn get_connection_manager(&self, info: &ConnectionInfo) -> Option> { - for manager in self.connection_managers.iter() { - if manager.handles_connection(info) { - return Some(manager.clone()); - } - } - None + fn get_connection_manager(&self) -> Option> { + self.connection_manager.clone() } /// Scan connection state machine and check if any connection should be closed. @@ -451,29 +445,17 @@ impl<'a> TunToProxy<'a> { } } }; - if connection_info.protocol == IpProtocol::Tcp { - let server_addr = self - .get_connection_manager(&connection_info) - .ok_or("get_connection_manager")? - .get_server_addr(); - if first_packet { - let mut done = false; - for manager in self.connection_managers.iter_mut() { - let tcp_proxy_handler = manager.new_tcp_proxy(&connection_info, false); - if tcp_proxy_handler.is_err() { - continue; - } - let tcp_proxy_handler = tcp_proxy_handler?; - let state = self.create_new_tcp_connection_state(server_addr, dst, tcp_proxy_handler)?; - self.connection_map.insert(connection_info.clone(), state); - log::info!("Connect done {} ({})", connection_info, dst); - done = true; - break; - } - if !done { - log::debug!("No connection manager for {} ({})", connection_info, dst); - } + let manager = self.get_connection_manager().ok_or("get connection manager")?; + let server_addr = manager.get_server_addr(); + + if connection_info.protocol == IpProtocol::Tcp { + if first_packet { + let tcp_proxy_handler = manager.new_tcp_proxy(&connection_info, false)?; + let state = self.create_new_tcp_connection_state(server_addr, dst, tcp_proxy_handler)?; + self.connection_map.insert(connection_info.clone(), state); + + log::info!("Connect done {} ({})", connection_info, dst); } else if !self.connection_map.contains_key(&connection_info) { log::debug!("Not found {} ({})", connection_info, dst); return Ok(()); @@ -514,12 +496,6 @@ impl<'a> TunToProxy<'a> { } } else { // Another UDP packet - let manager = self.get_connection_manager(&connection_info); - if manager.is_none() { - return Ok(()); - } - let manager = manager.ok_or("")?; - let server_addr = manager.get_server_addr(); let tcp_proxy_handler = manager.new_tcp_proxy(&connection_info, true)?; let state = self.create_new_tcp_connection_state(server_addr, dst, tcp_proxy_handler)?; self.connection_map.insert(connection_info.clone(), state); @@ -664,19 +640,19 @@ impl<'a> TunToProxy<'a> { } fn mio_socket_event(&mut self, event: &Event) -> Result<(), Error> { - let e = "connection not found"; let conn_info = match self.find_info_by_token(event.token()) { Some(conn_info) => conn_info.clone(), None => { // We may have closed the connection in an earlier iteration over the poll events, // e.g. because an event through the tunnel interface indicated that the connection // should be closed. - log::trace!("{e}"); + log::trace!("Connection info not found"); return Ok(()); } }; - let server = self.get_connection_manager(&conn_info).ok_or(e)?.get_server_addr(); + let e = "connection manager not found"; + let server = self.get_connection_manager().ok_or(e)?.get_server_addr(); let mut block = || -> Result<(), Error> { if event.is_readable() || event.is_read_closed() {