1use 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#[derive(Debug, Clone)]
23pub enum TargetAddr {
24 Ipv4([u8; 4], u16),
26 Ipv6([u8; 16], u16),
28 Domain(String, u16),
30}
31
32impl TargetAddr {
33 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 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 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 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 pub fn port(&self) -> u16 {
88 match self {
89 TargetAddr::Ipv4(_, port) | TargetAddr::Ipv6(_, port) | TargetAddr::Domain(_, port) => *port,
90 }
91 }
92}
93
94pub struct Socks5Server;
96
97impl Socks5Server {
98 pub async fn handle_handshake(stream: &mut TcpStream) -> Result<TargetAddr> {
101 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 pub async fn handle_handshake_with_version(stream: &mut TcpStream) -> Result<TargetAddr> {
121 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 async fn handle_handshake_internal(stream: &mut TcpStream, nmethods: u8) -> Result<TargetAddr> {
131 let mut methods = vec![0u8; nmethods as usize];
133 stream.read_exact(&mut methods).await?;
134
135 let selected_method = if methods.contains(&SOCKS5_NO_AUTHENTICATION) {
138 SOCKS5_NO_AUTHENTICATION
139 } else {
140 SOCKS5_NO_ACCEPTABLE_METHODS
141 };
142
143 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 let mut buf = [0u8; 4];
155 stream.read_exact(&mut buf).await?;
156
157 let version = buf[0];
158 let cmd = buf[1];
159 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 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 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 Self::send_reply(stream, SOCKS5_REP_SUCCESS).await?;
217
218 Ok(target_addr)
219 }
220
221 async fn send_reply(stream: &mut TcpStream, reply_code: u8) -> Result<()> {
223 let response = [
226 SOCKS5_VERSION,
227 reply_code,
228 0x00, SOCKS5_ATYP_IPV4,
230 0x00,
231 0x00,
232 0x00,
233 0x00, 0x00,
235 0x00, ];
237 stream.write_all(&response).await?;
238 Ok(())
239 }
240}