Prepare DNS support

src/socks5.rs

@@ -1,7 +1,7 @@
 use crate::error::Error;
 use crate::tun2proxy::{
-    Connection, ConnectionManager, Credentials, IncomingDataEvent, IncomingDirection,
-    OutgoingDataEvent, OutgoingDirection, TcpProxy,
+    Connection, ConnectionManager, Credentials, DestinationHost, IncomingDataEvent,
+    IncomingDirection, OutgoingDataEvent, OutgoingDirection, TcpProxy,
 };
 use std::collections::VecDeque;
 use std::net::{IpAddr, SocketAddr};
@@ -18,7 +18,6 @@ enum SocksState {
     Established,
 }
 
-#[allow(dead_code)]
 #[repr(u8)]
 #[derive(Copy, Clone)]
 enum SocksAddressType {
@@ -28,7 +27,6 @@ enum SocksAddressType {
 }
 
 #[allow(dead_code)]
-#[repr(u8)]
 enum SocksAuthentication {
     None = 0,
     Password = 2,
@@ -68,7 +66,7 @@ pub(crate) struct SocksConnection {
 impl SocksConnection {
     pub fn new(connection: &Connection, manager: std::rc::Rc<dyn ConnectionManager>) -> Self {
         let mut result = Self {
-            connection: *connection,
+            connection: connection.clone(),
             state: SocksState::ServerHello,
             client_inbuf: Default::default(),
             server_inbuf: Default::default(),
@@ -84,9 +82,11 @@ impl SocksConnection {
     fn send_client_hello(&mut self) {
         let credentials = self.manager.get_credentials();
         if credentials.is_some() {
-            self.server_outbuf.extend(&[5u8, 1, 2]);
+            self.server_outbuf
+                .extend(&[5u8, 1, SocksAuthentication::Password as u8]);
         } else {
-            self.server_outbuf.extend(&[5u8, 1, 0]);
+            self.server_outbuf
+                .extend(&[5u8, 1, SocksAuthentication::None as u8]);
         }
         self.state = SocksState::ServerHello;
     }
@@ -191,16 +191,28 @@ impl SocksConnection {
     }
 
     fn send_request(&mut self) -> Result<(), Error> {
-        let dst_ip = self.connection.dst.ip();
-        let cmd = if dst_ip.is_ipv4() { 1 } else { 4 };
-        self.server_outbuf.extend(&[5u8, 1, 0, cmd]);
-        match dst_ip {
-            IpAddr::V4(ip) => self.server_outbuf.extend(ip.octets().as_ref()),
-            IpAddr::V6(ip) => self.server_outbuf.extend(ip.octets().as_ref()),
-        };
+        self.server_outbuf.extend(&[5u8, 1, 0]);
+        match &self.connection.dst.host {
+            DestinationHost::Address(dst_ip) => {
+                let cmd = if dst_ip.is_ipv4() {
+                    SocksAddressType::Ipv4
+                } else {
+                    SocksAddressType::Ipv6
+                };
+                self.server_outbuf.extend(&[cmd as u8]);
+                match dst_ip {
+                    IpAddr::V4(ip) => self.server_outbuf.extend(ip.octets().as_ref()),
+                    IpAddr::V6(ip) => self.server_outbuf.extend(ip.octets().as_ref()),
+                };
+            }
+            DestinationHost::Hostname(_) => {
+                self.server_outbuf
+                    .extend(&[SocksAddressType::DomainName as u8]);
+            }
+        }
         self.server_outbuf.extend(&[
-            (self.connection.dst.port() >> 8) as u8,
-            (self.connection.dst.port() & 0xff) as u8,
+            (self.connection.dst.port >> 8) as u8,
+            (self.connection.dst.port & 0xff) as u8,
         ]);
         self.state = SocksState::ReceiveResponse;
         self.state_change()

src/tun2proxy.rs

@@ -14,16 +14,67 @@ use smoltcp::wire::{
 };
 use std::collections::HashMap;
 use std::convert::From;
+use std::fmt::{Display, Formatter};
 use std::io::{Read, Write};
 use std::net::Shutdown::Both;
 use std::net::{IpAddr, Shutdown, SocketAddr};
 use std::os::unix::io::AsRawFd;
 
-#[derive(Hash, Clone, Copy, Eq, PartialEq)]
-pub struct Connection {
-    pub src: std::net::SocketAddr,
-    pub dst: std::net::SocketAddr,
-    pub proto: u8,
+#[derive(Hash, Clone, Eq, PartialEq)]
+pub enum DestinationHost {
+    Address(IpAddr),
+    Hostname(String),
+}
+
+impl ToString for DestinationHost {
+    fn to_string(&self) -> String {
+        match self {
+            DestinationHost::Address(addr) => addr.to_string(),
+            DestinationHost::Hostname(name) => name.clone(),
+        }
+    }
+}
+
+#[derive(Hash, Clone, Eq, PartialEq)]
+pub(crate) struct Destination {
+    pub(crate) host: DestinationHost,
+    pub(crate) port: u16,
+}
+
+impl From<Destination> for SocketAddr {
+    fn from(value: Destination) -> Self {
+        SocketAddr::new(
+            match value.host {
+                DestinationHost::Address(addr) => addr,
+                DestinationHost::Hostname(_) => {
+                    panic!("Failed to convert hostname destination into socket address")
+                }
+            },
+            value.port,
+        )
+    }
+}
+
+impl From<SocketAddr> for Destination {
+    fn from(addr: SocketAddr) -> Self {
+        Self {
+            host: DestinationHost::Address(addr.ip()),
+            port: addr.port(),
+        }
+    }
+}
+
+impl Display for Destination {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}", self.host.to_string(), self.port)
+    }
+}
+
+#[derive(Hash, Clone, Eq, PartialEq)]
+pub(crate) struct Connection {
+    pub(crate) src: std::net::SocketAddr,
+    pub(crate) dst: Destination,
+    pub(crate) proto: u8,
 }
 
 impl std::fmt::Display for Connection {
@@ -105,7 +156,7 @@ fn connection_tuple(frame: &[u8]) -> Option<(Connection, bool, usize, usize)> {
         ) {
             let connection = Connection {
                 src: SocketAddr::new(src_addr, ports.0),
-                dst: SocketAddr::new(dst_addr, ports.1),
+                dst: SocketAddr::new(dst_addr, ports.1).into(),
                 proto,
             };
             return Some((connection, first_packet, payload_offset, payload_size));
@@ -130,7 +181,7 @@ fn connection_tuple(frame: &[u8]) -> Option<(Connection, bool, usize, usize)> {
             {
                 let connection = Connection {
                     src: SocketAddr::new(src_addr, ports.0),
-                    dst: SocketAddr::new(dst_addr, ports.1),
+                    dst: SocketAddr::new(dst_addr, ports.1).into(),
                     proto,
                 };
                 Some((connection, first_packet, payload_offset, payload_size))
@@ -348,7 +399,10 @@ impl<'a> TunToProxy<'a> {
                                 smoltcp::socket::tcp::SocketBuffer::new(vec![0; 4096]),
                             );
                             socket.set_ack_delay(None);
-                            socket.listen(connection.dst).unwrap();
+                            let dst = connection.dst.clone();
+                            socket
+                                .listen(<Destination as Into<SocketAddr>>::into(dst))
+                                .unwrap();
                             let handle = self.sockets.add(socket);
 
                             let client = TcpStream::connect(server).unwrap();
@@ -363,7 +417,7 @@ impl<'a> TunToProxy<'a> {
                                 handler,
                             };
 
-                            self.token_to_connection.insert(token, connection);
+                            self.token_to_connection.insert(token, connection.clone());
                             self.poll
                                 .registry()
                                 .register(
@@ -373,7 +427,7 @@ impl<'a> TunToProxy<'a> {
                                 )
                                 .unwrap();
 
-                            self.connections.insert(connection, state);
+                            self.connections.insert(connection.clone(), state);
 
                             info!("CONNECT {}", connection,);
                             break;
@@ -456,7 +510,7 @@ impl<'a> TunToProxy<'a> {
 
     fn mio_socket_event(&mut self, event: &Event) {
         if let Some(conn_ref) = self.token_to_connection.get(&event.token()) {
-            let connection = *conn_ref;
+            let connection = conn_ref.clone();
             if event.is_readable() {
                 {
                     let state = self.connections.get_mut(&connection).unwrap();