diff --git a/src/http.rs b/src/http.rs index c3aec4d..9ed5679 100644 --- a/src/http.rs +++ b/src/http.rs @@ -6,6 +6,7 @@ use crate::tun2proxy::{ use base64::Engine; use std::collections::VecDeque; use std::net::SocketAddr; +use std::rc::Rc; #[derive(Eq, PartialEq, Debug)] #[allow(dead_code)] @@ -27,7 +28,7 @@ pub struct HttpConnection { } impl HttpConnection { - fn new(connection: &Connection, manager: std::rc::Rc) -> Self { + fn new(connection: &Connection, manager: Rc) -> Self { let mut server_outbuf: VecDeque = VecDeque::new(); { let credentials = manager.get_credentials(); @@ -175,14 +176,12 @@ impl ConnectionManager for HttpManager { fn new_connection( &self, connection: &Connection, - manager: std::rc::Rc, - ) -> Option> { + manager: Rc, + ) -> Option> { if connection.proto != smoltcp::wire::IpProtocol::Tcp.into() { return None; } - Some(std::boxed::Box::new(HttpConnection::new( - connection, manager, - ))) + Some(Box::new(HttpConnection::new(connection, manager))) } fn close_connection(&self, _: &Connection) {} @@ -197,8 +196,8 @@ impl ConnectionManager for HttpManager { } impl HttpManager { - pub fn new(server: SocketAddr, credentials: Option) -> std::rc::Rc { - std::rc::Rc::new(Self { + pub fn new(server: SocketAddr, credentials: Option) -> Rc { + Rc::new(Self { server, credentials, }) diff --git a/src/lib.rs b/src/lib.rs index 248e8b8..e3d6023 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -85,5 +85,7 @@ pub fn main_entry(tun: &str, proxy: Proxy) { ttp.add_connection_manager(HttpManager::new(proxy.addr, proxy.credentials)); } } - ttp.run(); + if let Err(e) = ttp.run() { + log::error!("{e}"); + } } diff --git a/src/socks5.rs b/src/socks5.rs index 9f4e6c4..640b006 100644 --- a/src/socks5.rs +++ b/src/socks5.rs @@ -5,6 +5,7 @@ use crate::tun2proxy::{ }; use std::collections::VecDeque; use std::net::{IpAddr, SocketAddr}; +use std::rc::Rc; #[derive(Eq, PartialEq, Debug)] #[allow(dead_code)] @@ -60,11 +61,11 @@ pub(crate) struct SocksConnection { client_outbuf: VecDeque, server_outbuf: VecDeque, data_buf: VecDeque, - manager: std::rc::Rc, + manager: Rc, } impl SocksConnection { - pub fn new(connection: &Connection, manager: std::rc::Rc) -> Self { + pub fn new(connection: &Connection, manager: Rc) -> Self { let mut result = Self { connection: connection.clone(), state: SocksState::ServerHello, @@ -302,14 +303,12 @@ impl ConnectionManager for Socks5Manager { fn new_connection( &self, connection: &Connection, - manager: std::rc::Rc, - ) -> Option> { + manager: Rc, + ) -> Option> { if connection.proto != smoltcp::wire::IpProtocol::Tcp.into() { return None; } - Some(std::boxed::Box::new(SocksConnection::new( - connection, manager, - ))) + Some(Box::new(SocksConnection::new(connection, manager))) } fn close_connection(&self, _: &Connection) {} @@ -324,8 +323,8 @@ impl ConnectionManager for Socks5Manager { } impl Socks5Manager { - pub fn new(server: SocketAddr, credentials: Option) -> std::rc::Rc { - std::rc::Rc::new(Self { + pub fn new(server: SocketAddr, credentials: Option) -> Rc { + Rc::new(Self { server, credentials, }) diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index 28a3ed3..48c85de 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -19,6 +19,7 @@ use std::io::{Read, Write}; use std::net::Shutdown::Both; use std::net::{IpAddr, Shutdown, SocketAddr}; use std::os::unix::io::AsRawFd; +use std::rc::Rc; #[derive(Hash, Clone, Eq, PartialEq)] pub enum DestinationHost { @@ -197,7 +198,7 @@ struct ConnectionState { smoltcp_handle: SocketHandle, mio_stream: TcpStream, token: Token, - handler: std::boxed::Box, + handler: Box, } #[derive(Default, Clone, Debug)] @@ -227,21 +228,22 @@ pub(crate) trait ConnectionManager { fn new_connection( &self, connection: &Connection, - manager: std::rc::Rc, - ) -> Option>; + manager: Rc, + ) -> Option>; fn close_connection(&self, connection: &Connection); fn get_server(&self) -> SocketAddr; fn get_credentials(&self) -> &Option; } +const TCP_TOKEN: Token = Token(0); +const UDP_TOKEN: Token = Token(1); + pub(crate) struct TunToProxy<'a> { tun: TunTapInterface, poll: Poll, - tun_token: Token, - udp_token: Token, iface: Interface, connections: HashMap, - connection_managers: Vec>, + connection_managers: Vec>, next_token: usize, token_to_connection: HashMap, sockets: SocketSet<'a>, @@ -250,13 +252,12 @@ pub(crate) struct TunToProxy<'a> { impl<'a> TunToProxy<'a> { pub(crate) fn new(interface: &str) -> Self { - let tun_token = Token(0); let tun = TunTapInterface::new(interface, Medium::Ip).unwrap(); let poll = Poll::new().unwrap(); poll.registry() .register( &mut SourceFd(&tun.as_raw_fd()), - tun_token, + TCP_TOKEN, Interest::READABLE, ) .unwrap(); @@ -285,8 +286,6 @@ impl<'a> TunToProxy<'a> { Self { tun, poll, - tun_token, - udp_token: Token(1), iface, connections: Default::default(), next_token: 2, @@ -297,7 +296,7 @@ impl<'a> TunToProxy<'a> { } } - pub(crate) fn add_connection_manager(&mut self, manager: std::rc::Rc) { + pub(crate) fn add_connection_manager(&mut self, manager: Rc) { self.connection_managers.push(manager); } @@ -328,10 +327,7 @@ impl<'a> TunToProxy<'a> { info!("CLOSE {}", connection); } - fn get_connection_manager( - &self, - connection: &Connection, - ) -> Option> { + fn get_connection_manager(&self, connection: &Connection) -> Option> { for manager in self.connection_managers.iter() { if manager.handles_connection(connection) { return Some(manager.clone()); @@ -570,18 +566,16 @@ impl<'a> TunToProxy<'a> { fn udp_event(&mut self, _event: &Event) {} - pub(crate) fn run(&mut self) { + pub(crate) fn run(&mut self) -> Result<(), Error> { let mut events = Events::with_capacity(1024); loop { - self.poll.poll(&mut events, None).unwrap(); + self.poll.poll(&mut events, None)?; for event in events.iter() { - if event.token() == self.tun_token { - self.tun_event(event); - } else if event.token() == self.udp_token { - self.udp_event(event); - } else { - self.mio_socket_event(event); + match event.token() { + TCP_TOKEN => self.tun_event(event), + UDP_TOKEN => self.udp_event(event), + _ => self.mio_socket_event(event), } } }