diff --git a/src/socks.rs b/src/socks.rs index 3c71e2a..ec4bf39 100644 --- a/src/socks.rs +++ b/src/socks.rs @@ -194,6 +194,13 @@ impl SocksProxyImpl { self.state_change() } + fn send_request_socks5(&mut self) -> Result<(), Error> { + protocol::Request::new(protocol::Command::Connect, self.info.dst.clone()) + .write_to_stream(&mut self.server_outbuf)?; + self.state = SocksState::ReceiveResponse; + self.state_change() + } + fn receive_connection_status(&mut self) -> Result<(), Error> { let response = protocol::Response::retrieve_from_stream(&mut self.server_inbuf.clone()); if let Err(e) = &response { @@ -217,14 +224,6 @@ impl SocksProxyImpl { self.state_change() } - fn send_request_socks5(&mut self) -> Result<(), Error> { - // self.server_outbuf.extend(&[self.version as u8, self.command as u8, 0]); - protocol::Request::new(protocol::Command::Connect, self.info.dst.clone()) - .write_to_stream(&mut self.server_outbuf)?; - self.state = SocksState::ReceiveResponse; - self.state_change() - } - fn relay_traffic(&mut self) -> Result<(), Error> { self.client_outbuf.extend(self.server_inbuf.iter()); self.server_outbuf.extend(self.client_inbuf.iter()); diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index 1bb0170..8be7f06 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -12,7 +12,7 @@ use std::{ collections::{HashMap, HashSet}, convert::{From, TryFrom}, io::{Read, Write}, - net::{IpAddr, Ipv4Addr, Ipv6Addr, Shutdown, Shutdown::Both, SocketAddr}, + net::{IpAddr, Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr}, os::unix::io::AsRawFd, rc::Rc, str::FromStr, @@ -278,6 +278,7 @@ impl<'a> TunToProxy<'a> { self.connection_managers.push(manager); } + /// Read data from virtual device (remote server) and inject it into tun interface. fn expect_smoltcp_send(&mut self) -> Result<(), Error> { self.iface.poll(Instant::now(), &mut self.device, &mut self.sockets); @@ -297,7 +298,7 @@ impl<'a> TunToProxy<'a> { fn remove_connection(&mut self, info: &ConnectionInfo) -> Result<(), Error> { if let Some(mut conn) = self.connection_map.remove(info) { - _ = conn.mio_stream.shutdown(Both); + _ = conn.mio_stream.shutdown(Shutdown::Both); if let Some(handle) = conn.smoltcp_handle { let socket = self.sockets.get_mut::(handle); socket.close(); @@ -322,11 +323,10 @@ impl<'a> TunToProxy<'a> { } fn check_change_close_state(&mut self, info: &ConnectionInfo) -> Result<(), Error> { - let state = self.connection_map.get_mut(info); - if state.is_none() { - return Ok(()); - } - let state = state.unwrap(); + let state = match self.connection_map.get_mut(info) { + None => return Ok(()), + Some(state) => state, + }; let mut closed_ends = 0; if (state.close_state & SERVER_WRITE_CLOSED) == SERVER_WRITE_CLOSED && !state @@ -492,7 +492,7 @@ impl<'a> TunToProxy<'a> { log::trace!("Subsequent packet {} ({})", connection_info, dst); } - // Inject the packet to advance the smoltcp socket state + // Inject the packet to advance the remote proxy server smoltcp socket state self.device.inject_packet(frame); // Having advanced the socket state, we expect the socket to ACK @@ -509,8 +509,8 @@ impl<'a> TunToProxy<'a> { } else if connection_info.protocol == IpProtocol::Udp { log::trace!("{} ({})", connection_info, dst); let port = connection_info.dst.port(); + let payload = &frame[payload_offset..payload_offset + payload_size]; if let (Some(virtual_dns), true) = (&mut self.options.virtual_dns, port == 53) { - let payload = &frame[payload_offset..payload_offset + payload_size]; let response = virtual_dns.receive_query(payload)?; { let rx_buffer = udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]); @@ -680,7 +680,7 @@ impl<'a> TunToProxy<'a> { if state.tcp_proxy_handler.reset_connection() { _ = self.poll.registry().deregister(&mut state.mio_stream); // Closes the connection with the proxy - state.mio_stream.shutdown(Both)?; + state.mio_stream.shutdown(Shutdown::Both)?; log::info!("RESET {}", conn_info); diff --git a/src/virtdevice.rs b/src/virtdevice.rs index fc862d9..721466c 100644 --- a/src/virtdevice.rs +++ b/src/virtdevice.rs @@ -3,6 +3,7 @@ use smoltcp::{ time::Instant, }; +/// Virtual device representing the remote proxy server. #[derive(Default)] pub struct VirtualTunDevice { capabilities: DeviceCapabilities,