Skip to main content

slinger_mitm/
socks5.rs

1//! SOCKS5 server implementation for MITM proxy
2
3use crate::error::{Error, Result};
4use tokio::io::{AsyncReadExt, AsyncWriteExt};
5use tokio::net::TcpStream;
6
7const SOCKS5_VERSION: u8 = 0x05;
8const SOCKS5_NO_AUTHENTICATION: u8 = 0x00;
9const SOCKS5_NO_ACCEPTABLE_METHODS: u8 = 0xFF;
10
11const SOCKS5_CMD_CONNECT: u8 = 0x01;
12
13const SOCKS5_ATYP_IPV4: u8 = 0x01;
14const SOCKS5_ATYP_DOMAIN: u8 = 0x03;
15const SOCKS5_ATYP_IPV6: u8 = 0x04;
16
17const SOCKS5_REP_SUCCESS: u8 = 0x00;
18const SOCKS5_REP_COMMAND_NOT_SUPPORTED: u8 = 0x07;
19const SOCKS5_REP_ADDRESS_TYPE_NOT_SUPPORTED: u8 = 0x08;
20
21/// SOCKS5 target address
22#[derive(Debug, Clone)]
23pub enum TargetAddr {
24  /// IPv4 address
25  Ipv4([u8; 4], u16),
26  /// IPv6 address
27  Ipv6([u8; 16], u16),
28  /// Domain name
29  Domain(String, u16),
30}
31
32impl TargetAddr {
33  /// Get host and port as string
34  pub fn to_host_port(&self) -> String {
35    match self {
36      TargetAddr::Ipv4(ip, port) => {
37        format!("{}.{}.{}.{}:{}", ip[0], ip[1], ip[2], ip[3], port)
38      }
39      TargetAddr::Ipv6(ip, port) => {
40        // Group bytes into 16-bit segments
41        let segments: Vec<u16> = (0..8)
42          .map(|i| u16::from_be_bytes([ip[i * 2], ip[i * 2 + 1]]))
43          .collect();
44        format!(
45          "[{:x}:{:x}:{:x}:{:x}:{:x}:{:x}:{:x}:{:x}]:{}",
46          segments[0],
47          segments[1],
48          segments[2],
49          segments[3],
50          segments[4],
51          segments[5],
52          segments[6],
53          segments[7],
54          port
55        )
56      }
57      TargetAddr::Domain(domain, port) => format!("{}:{}", domain, port),
58    }
59  }
60
61  /// Get host (without port)
62  pub fn host(&self) -> String {
63    match self {
64      TargetAddr::Ipv4(ip, _) => format!("{}.{}.{}.{}", ip[0], ip[1], ip[2], ip[3]),
65      TargetAddr::Ipv6(ip, _) => {
66        // Group bytes into 16-bit segments
67        let segments: Vec<u16> = (0..8)
68          .map(|i| u16::from_be_bytes([ip[i * 2], ip[i * 2 + 1]]))
69          .collect();
70        format!(
71          "{:x}:{:x}:{:x}:{:x}:{:x}:{:x}:{:x}:{:x}",
72          segments[0],
73          segments[1],
74          segments[2],
75          segments[3],
76          segments[4],
77          segments[5],
78          segments[6],
79          segments[7]
80        )
81      }
82      TargetAddr::Domain(domain, _) => domain.clone(),
83    }
84  }
85
86  /// Get port
87  pub fn port(&self) -> u16 {
88    match self {
89      TargetAddr::Ipv4(_, port) | TargetAddr::Ipv6(_, port) | TargetAddr::Domain(_, port) => *port,
90    }
91  }
92}
93
94/// SOCKS5 server for MITM proxy
95pub struct Socks5Server;
96
97impl Socks5Server {
98  /// Handle SOCKS5 client connection
99  /// Returns the target address if successful
100  pub async fn handle_handshake(stream: &mut TcpStream) -> Result<TargetAddr> {
101    // Read version and methods
102    let mut buf = [0u8; 2];
103    stream.read_exact(&mut buf).await?;
104
105    let version = buf[0];
106    let nmethods = buf[1];
107
108    if version != SOCKS5_VERSION {
109      return Err(Error::proxy_error(format!(
110        "Unsupported SOCKS version: {}",
111        version
112      )));
113    }
114
115    Self::handle_handshake_internal(stream, nmethods).await
116  }
117
118  /// Handle SOCKS5 client connection when version byte has already been read
119  /// Returns the target address if successful
120  pub async fn handle_handshake_with_version(stream: &mut TcpStream) -> Result<TargetAddr> {
121    // Version byte (0x05) was already read, now read nmethods
122    let mut buf = [0u8; 1];
123    stream.read_exact(&mut buf).await?;
124    let nmethods = buf[0];
125
126    Self::handle_handshake_internal(stream, nmethods).await
127  }
128
129  /// Internal handshake processing after version check
130  async fn handle_handshake_internal(stream: &mut TcpStream, nmethods: u8) -> Result<TargetAddr> {
131    // Read methods
132    let mut methods = vec![0u8; nmethods as usize];
133    stream.read_exact(&mut methods).await?;
134
135    // For now, we only support no authentication
136    // Check if client supports no auth
137    let selected_method = if methods.contains(&SOCKS5_NO_AUTHENTICATION) {
138      SOCKS5_NO_AUTHENTICATION
139    } else {
140      SOCKS5_NO_ACCEPTABLE_METHODS
141    };
142
143    // Send method selection response
144    let response = [SOCKS5_VERSION, selected_method];
145    stream.write_all(&response).await?;
146
147    if selected_method == SOCKS5_NO_ACCEPTABLE_METHODS {
148      return Err(Error::proxy_error(
149        "No acceptable authentication method".to_string(),
150      ));
151    }
152
153    // Read connection request
154    let mut buf = [0u8; 4];
155    stream.read_exact(&mut buf).await?;
156
157    let version = buf[0];
158    let cmd = buf[1];
159    // buf[2] is reserved
160    let atyp = buf[3];
161
162    if version != SOCKS5_VERSION {
163      return Err(Error::proxy_error(format!(
164        "Invalid SOCKS version in request: {}",
165        version
166      )));
167    }
168
169    // Only support CONNECT command
170    if cmd != SOCKS5_CMD_CONNECT {
171      Self::send_reply(stream, SOCKS5_REP_COMMAND_NOT_SUPPORTED).await?;
172      return Err(Error::proxy_error(format!("Unsupported command: {}", cmd)));
173    }
174
175    // Read target address
176    let target_addr = match atyp {
177      SOCKS5_ATYP_IPV4 => {
178        let mut addr = [0u8; 4];
179        stream.read_exact(&mut addr).await?;
180        let mut port_buf = [0u8; 2];
181        stream.read_exact(&mut port_buf).await?;
182        let port = u16::from_be_bytes(port_buf);
183        TargetAddr::Ipv4(addr, port)
184      }
185      SOCKS5_ATYP_IPV6 => {
186        let mut addr = [0u8; 16];
187        stream.read_exact(&mut addr).await?;
188        let mut port_buf = [0u8; 2];
189        stream.read_exact(&mut port_buf).await?;
190        let port = u16::from_be_bytes(port_buf);
191        TargetAddr::Ipv6(addr, port)
192      }
193      SOCKS5_ATYP_DOMAIN => {
194        let mut len_buf = [0u8; 1];
195        stream.read_exact(&mut len_buf).await?;
196        let len = len_buf[0] as usize;
197        let mut domain = vec![0u8; len];
198        stream.read_exact(&mut domain).await?;
199        let mut port_buf = [0u8; 2];
200        stream.read_exact(&mut port_buf).await?;
201        let port = u16::from_be_bytes(port_buf);
202        let domain_str = String::from_utf8(domain)
203          .map_err(|_| Error::proxy_error("Invalid domain name".to_string()))?;
204        TargetAddr::Domain(domain_str, port)
205      }
206      _ => {
207        Self::send_reply(stream, SOCKS5_REP_ADDRESS_TYPE_NOT_SUPPORTED).await?;
208        return Err(Error::proxy_error(format!(
209          "Unsupported address type: {}",
210          atyp
211        )));
212      }
213    };
214
215    // Send success reply
216    Self::send_reply(stream, SOCKS5_REP_SUCCESS).await?;
217
218    Ok(target_addr)
219  }
220
221  /// Send SOCKS5 reply
222  async fn send_reply(stream: &mut TcpStream, reply_code: u8) -> Result<()> {
223    // Reply format: [VER, REP, RSV, ATYP, BND.ADDR, BND.PORT]
224    // We'll use a simple IPv4 address 0.0.0.0:0 as the bound address
225    let response = [
226      SOCKS5_VERSION,
227      reply_code,
228      0x00, // Reserved
229      SOCKS5_ATYP_IPV4,
230      0x00,
231      0x00,
232      0x00,
233      0x00, // Bind address (0.0.0.0)
234      0x00,
235      0x00, // Bind port (0)
236    ];
237    stream.write_all(&response).await?;
238    Ok(())
239  }
240}