From 4af656039e1f8728bdd2a7b96f95c1e6a9485f27 Mon Sep 17 00:00:00 2001 From: ssrlive <30760636+ssrlive@users.noreply.github.com> Date: Thu, 23 Mar 2023 16:31:33 +0800 Subject: [PATCH 1/3] refine code --- src/http.rs | 15 +++++++-------- src/lib.rs | 4 +++- src/socks5.rs | 17 ++++++++--------- src/tun2proxy.rs | 40 +++++++++++++++++----------------------- 4 files changed, 35 insertions(+), 41 deletions(-) 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), } } } From 6fe44f805583e82e6b317016dd3cfd58d5b3bdbc Mon Sep 17 00:00:00 2001 From: ssrlive <30760636+ssrlive@users.noreply.github.com> Date: Thu, 23 Mar 2023 18:31:09 +0800 Subject: [PATCH 2/3] IpProtocol values --- src/http.rs | 5 +++-- src/socks5.rs | 5 +++-- src/tun2proxy.rs | 11 ++++++----- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/http.rs b/src/http.rs index 9ed5679..ca31b50 100644 --- a/src/http.rs +++ b/src/http.rs @@ -4,6 +4,7 @@ use crate::tun2proxy::{ OutgoingDataEvent, OutgoingDirection, TcpProxy, }; use base64::Engine; +use smoltcp::wire::IpProtocol; use std::collections::VecDeque; use std::net::SocketAddr; use std::rc::Rc; @@ -170,7 +171,7 @@ pub struct HttpManager { impl ConnectionManager for HttpManager { fn handles_connection(&self, connection: &Connection) -> bool { - connection.proto == smoltcp::wire::IpProtocol::Tcp.into() + connection.proto == IpProtocol::Tcp.into() } fn new_connection( @@ -178,7 +179,7 @@ impl ConnectionManager for HttpManager { connection: &Connection, manager: Rc, ) -> Option> { - if connection.proto != smoltcp::wire::IpProtocol::Tcp.into() { + if connection.proto != IpProtocol::Tcp.into() { return None; } Some(Box::new(HttpConnection::new(connection, manager))) diff --git a/src/socks5.rs b/src/socks5.rs index 640b006..cb109fd 100644 --- a/src/socks5.rs +++ b/src/socks5.rs @@ -3,6 +3,7 @@ use crate::tun2proxy::{ Connection, ConnectionManager, Credentials, DestinationHost, IncomingDataEvent, IncomingDirection, OutgoingDataEvent, OutgoingDirection, TcpProxy, }; +use smoltcp::wire::IpProtocol; use std::collections::VecDeque; use std::net::{IpAddr, SocketAddr}; use std::rc::Rc; @@ -297,7 +298,7 @@ pub struct Socks5Manager { impl ConnectionManager for Socks5Manager { fn handles_connection(&self, connection: &Connection) -> bool { - connection.proto == smoltcp::wire::IpProtocol::Tcp.into() + connection.proto == IpProtocol::Tcp.into() } fn new_connection( @@ -305,7 +306,7 @@ impl ConnectionManager for Socks5Manager { connection: &Connection, manager: Rc, ) -> Option> { - if connection.proto != smoltcp::wire::IpProtocol::Tcp.into() { + if connection.proto != IpProtocol::Tcp.into() { return None; } Some(Box::new(SocksConnection::new(connection, manager))) diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index 48c85de..bd0a985 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -10,7 +10,8 @@ use smoltcp::phy::{Device, Medium, RxToken, TunTapInterface, TxToken}; use smoltcp::socket::tcp; use smoltcp::time::Instant; use smoltcp::wire::{ - IpAddress, IpCidr, Ipv4Address, Ipv4Packet, Ipv6Address, Ipv6Packet, TcpPacket, UdpPacket, + IpAddress, IpCidr, IpProtocol, Ipv4Address, Ipv4Packet, Ipv6Address, Ipv6Packet, TcpPacket, + UdpPacket, }; use std::collections::HashMap; use std::convert::From; @@ -115,7 +116,7 @@ fn get_transport_info( transport_offset: usize, packet: &[u8], ) -> Option<((u16, u16), bool, usize, usize)> { - if proto == smoltcp::wire::IpProtocol::Udp.into() { + if proto == IpProtocol::Udp.into() { match UdpPacket::new_checked(packet) { Ok(result) => Some(( (result.src_port(), result.dst_port()), @@ -125,7 +126,7 @@ fn get_transport_info( )), Err(_) => None, } - } else if proto == smoltcp::wire::IpProtocol::Tcp.into() { + } else if proto == IpProtocol::Tcp.into() { match TcpPacket::new_checked(packet) { Ok(result) => Some(( (result.src_port(), result.dst_port()), @@ -380,7 +381,7 @@ impl<'a> TunToProxy<'a> { if let Some((connection, first_packet, _payload_offset, _payload_size)) = connection_tuple(frame) { - if connection.proto == smoltcp::wire::IpProtocol::Tcp.into() { + if connection.proto == IpProtocol::Tcp.into() { let cm = self.get_connection_manager(&connection); if cm.is_none() { return; @@ -447,7 +448,7 @@ impl<'a> TunToProxy<'a> { // The connection handler builds up the connection or encapsulates the data. // Therefore, we now expect it to write data to the server. self.write_to_server(&connection); - } else if connection.proto == smoltcp::wire::IpProtocol::Udp.into() { + } else if connection.proto == IpProtocol::Udp.into() { // UDP is not yet supported /*if _payload_offset > frame.len() || _payload_offset + _payload_offset > frame.len() { return; From d4eed896bb315ffe12cd61e3c13182dab2e15baf Mon Sep 17 00:00:00 2001 From: ssrlive <30760636+ssrlive@users.noreply.github.com> Date: Thu, 23 Mar 2023 20:00:59 +0800 Subject: [PATCH 3/3] SocketAddr issues --- src/http.rs | 2 +- src/socks5.rs | 2 +- src/tun2proxy.rs | 23 +++++++++++------------ 3 files changed, 13 insertions(+), 14 deletions(-) diff --git a/src/http.rs b/src/http.rs index ca31b50..8f6e6b4 100644 --- a/src/http.rs +++ b/src/http.rs @@ -165,7 +165,7 @@ impl TcpProxy for HttpConnection { } pub struct HttpManager { - server: std::net::SocketAddr, + server: SocketAddr, credentials: Option, } diff --git a/src/socks5.rs b/src/socks5.rs index cb109fd..a9d2416 100644 --- a/src/socks5.rs +++ b/src/socks5.rs @@ -292,7 +292,7 @@ impl TcpProxy for SocksConnection { } pub struct Socks5Manager { - server: std::net::SocketAddr, + server: SocketAddr, credentials: Option, } diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index bd0a985..677c473 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -14,7 +14,7 @@ use smoltcp::wire::{ UdpPacket, }; use std::collections::HashMap; -use std::convert::From; +use std::convert::{From, TryFrom}; use std::fmt::{Display, Formatter}; use std::io::{Read, Write}; use std::net::Shutdown::Both; @@ -43,17 +43,18 @@ pub(crate) struct Destination { pub(crate) port: u16, } -impl From for SocketAddr { - fn from(value: Destination) -> Self { - SocketAddr::new( +impl TryFrom for SocketAddr { + type Error = Error; + fn try_from(value: Destination) -> Result { + Ok(SocketAddr::new( match value.host { DestinationHost::Address(addr) => addr, - DestinationHost::Hostname(_) => { - panic!("Failed to convert hostname destination into socket address") + DestinationHost::Hostname(e) => { + return Err(e.into()); } }, value.port, - ) + )) } } @@ -74,7 +75,7 @@ impl Display for Destination { #[derive(Hash, Clone, Eq, PartialEq)] pub(crate) struct Connection { - pub(crate) src: std::net::SocketAddr, + pub(crate) src: SocketAddr, pub(crate) dst: Destination, pub(crate) proto: u8, } @@ -396,10 +397,8 @@ impl<'a> TunToProxy<'a> { smoltcp::socket::tcp::SocketBuffer::new(vec![0; 4096]), ); socket.set_ack_delay(None); - let dst = connection.dst.clone(); - socket - .listen(>::into(dst)) - .unwrap(); + let dst = SocketAddr::try_from(connection.dst.clone()).unwrap(); + socket.listen(dst).unwrap(); let handle = self.sockets.add(socket); let client = TcpStream::connect(server).unwrap();