Merge pull request #52 from ssrlive/v6

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());

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::<tcp::Socket>(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);
+        let state = match self.connection_map.get_mut(info) {
-        if state.is_none() {
+            None => return Ok(()),
-            return Ok(());
+            Some(state) => state,
-        }
+        };
-        let state = state.unwrap();
         let mut closed_ends = 0;
         if (state.close_state & SERVER_WRITE_CLOSED) == SERVER_WRITE_CLOSED
             && !state
@@ -452,6 +452,78 @@ impl<'a> TunToProxy<'a> {
                             continue;
                         }
                         let tcp_proxy_handler = tcp_proxy_handler?;
+                        self.create_new_tcp_proxy_connection(
+                            server_addr,
+                            dst,
+                            tcp_proxy_handler,
+                            connection_info.clone(),
+                        )?;
+
+                        log::info!("Connect done {} ({})", connection_info, dst);
+                        done = true;
+                        break;
+                    }
+                    if !done {
+                        log::debug!("No connection manager for {} ({})", connection_info, dst);
+                    }
+                } else if !self.connection_map.contains_key(&connection_info) {
+                    log::debug!("Not found {} ({})", connection_info, dst);
+                    return Ok(());
+                } else {
+                    log::trace!("Subsequent packet {} ({})", connection_info, dst);
+                }
+
+                // 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
+                // Exfiltrate the response packets generated by the socket and inject them
+                // into the tunnel interface.
+                self.expect_smoltcp_send()?;
+
+                // Read from the smoltcp socket and push the data to the connection handler.
+                self.tunsocket_read_and_forward(&connection_info)?;
+
+                // 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_info)?;
+            } 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 response = virtual_dns.receive_query(payload)?;
+                    {
+                        let rx_buffer = udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
+                        let tx_buffer = udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
+                        let mut socket = udp::Socket::new(rx_buffer, tx_buffer);
+                        socket.bind(dst)?;
+                        let meta = UdpMetadata::from(connection_info.src);
+                        socket.send_slice(response.as_slice(), meta)?;
+                        let handle = self.sockets.add(socket);
+                        self.expect_smoltcp_send()?;
+                        self.sockets.remove(handle);
+                    }
+                }
+                // Otherwise, UDP is not yet supported.
+            } else {
+                log::warn!("Unsupported protocol: {} ({})", connection_info, dst);
+            }
+            Ok::<(), Error>(())
+        };
+        if let Err(error) = handler() {
+            log::error!("{}", error);
+        }
+        Ok(())
+    }
+
+    fn create_new_tcp_proxy_connection(
+        &mut self,
+        server_addr: SocketAddr,
+        dst: SocketAddr,
+        tcp_proxy_handler: Box<dyn TcpProxy>,
+        connection_info: ConnectionInfo,
+    ) -> Result<()> {
         let mut socket = tcp::Socket::new(
             tcp::SocketBuffer::new(vec![0; 1024 * 128]),
             tcp::SocketBuffer::new(vec![0; 1024 * 128]),
@@ -477,62 +549,6 @@ impl<'a> TunToProxy<'a> {
         self.connection_map.insert(connection_info.clone(), state);
 
         self.token_to_info.insert(token, connection_info.clone());
-
-                        log::info!("Connect done {} ({})", connection_info, dst);
-                        done = true;
-                        break;
-                    }
-                    if !done {
-                        log::debug!("No connection manager for {} ({})", connection_info, dst);
-                    }
-                } else if !self.connection_map.contains_key(&connection_info) {
-                    log::debug!("Not found {} ({})", connection_info, dst);
-                    return Ok(());
-                } else {
-                    log::trace!("Subsequent packet {} ({})", connection_info, dst);
-                }
-
-                // Inject the packet to advance the smoltcp socket state
-                self.device.inject_packet(frame);
-
-                // Having advanced the socket state, we expect the socket to ACK
-                // Exfiltrate the response packets generated by the socket and inject them
-                // into the tunnel interface.
-                self.expect_smoltcp_send()?;
-
-                // Read from the smoltcp socket and push the data to the connection handler.
-                self.tunsocket_read_and_forward(&connection_info)?;
-
-                // 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_info)?;
-            } else if connection_info.protocol == IpProtocol::Udp {
-                log::trace!("{} ({})", connection_info, dst);
-                let port = connection_info.dst.port();
-                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]);
-                        let tx_buffer = udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
-                        let mut socket = udp::Socket::new(rx_buffer, tx_buffer);
-                        socket.bind(dst)?;
-                        let meta = UdpMetadata::from(connection_info.src);
-                        socket.send_slice(response.as_slice(), meta)?;
-                        let handle = self.sockets.add(socket);
-                        self.expect_smoltcp_send()?;
-                        self.sockets.remove(handle);
-                    }
-                }
-                // Otherwise, UDP is not yet supported.
-            } else {
-                log::warn!("Unsupported protocol: {} ({})", connection_info, dst);
-            }
-            Ok::<(), Error>(())
-        };
-        if let Err(error) = handler() {
-            log::error!("{}", error);
-        }
         Ok(())
     }
 
@@ -680,7 +696,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);
 

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,