1use std::net::SocketAddr;
2
3use tokio::io::{AsyncReadExt, AsyncWriteExt};
4
5#[derive(Debug, Clone)]
6pub struct ProbeResult {
7 pub success: bool,
8 pub bytes_sent: u64,
9 pub bytes_received: u64,
10 pub response: Vec<u8>,
11 pub error: Option<String>,
12 pub reply_code: Option<u16>,
13 pub reply_message: Option<String>,
14}
15
16fn ok(response: Vec<u8>, bytes_sent: u64, bytes_received: u64) -> ProbeResult {
17 ProbeResult {
18 success: true,
19 bytes_sent,
20 bytes_received,
21 response,
22 error: None,
23 reply_code: None,
24 reply_message: None,
25 }
26}
27
28fn err(msg: impl Into<String>) -> ProbeResult {
29 ProbeResult {
30 success: false,
31 bytes_sent: 0,
32 bytes_received: 0,
33 response: Vec::new(),
34 error: Some(msg.into()),
35 reply_code: None,
36 reply_message: None,
37 }
38}
39
40fn err_with_reply(msg: impl Into<String>, code: u16) -> ProbeResult {
41 ProbeResult {
42 success: false,
43 bytes_sent: 0,
44 bytes_received: 0,
45 response: Vec::new(),
46 error: Some(msg.into()),
47 reply_code: Some(code),
48 reply_message: None,
49 }
50}
51
52fn build_sock5_target(target: SocketAddr) -> Vec<u8> {
53 let mut buf = vec![0x05, 0x01, 0x00, 0x01];
54 match target.ip() {
55 std::net::IpAddr::V4(ip) => buf.extend_from_slice(&ip.octets()),
56 std::net::IpAddr::V6(ip) => {
57 buf[3] = 0x04;
58 buf.extend_from_slice(&ip.octets());
59 }
60 }
61 buf.extend_from_slice(&target.port().to_be_bytes());
62 buf
63}
64
65fn extract_http_status_line(resp: &[u8]) -> String {
66 let text = String::from_utf8_lossy(resp);
67 text.lines().next().unwrap_or("").to_string()
68}
69
70fn extract_http_body(resp: &[u8]) -> Vec<u8> {
71 let text = String::from_utf8_lossy(resp);
72 if let Some(pos) = text.find("\r\n\r\n") {
73 resp[pos + 4..].to_vec()
74 } else {
75 resp.to_vec()
76 }
77}
78
79pub async fn socks5_tcp_connect(
80 proxy: SocketAddr,
81 target: SocketAddr,
82 payload: &[u8],
83) -> ProbeResult {
84 let mut stream = match tokio::net::TcpStream::connect(proxy).await {
85 Ok(s) => s,
86 Err(e) => return err(format!("connect to proxy: {e}")),
87 };
88
89 if let Err(e) = stream.write_all(&[0x05, 0x01, 0x00]).await {
90 return err(format!("greeting write: {e}"));
91 }
92 let mut buf = [0u8; 2];
93 if let Err(e) = stream.read_exact(&mut buf).await {
94 return err(format!("greeting read: {e}"));
95 }
96 if buf != [0x05, 0x00] {
97 return err(format!("unexpected greeting response: {buf:02x?}"));
98 }
99
100 let req = build_sock5_target(target);
101 let connect_bytes = req.len() as u64;
102 if let Err(e) = stream.write_all(&req).await {
103 return err(format!("connect request write: {e}"));
104 }
105 let mut resp = [0u8; 10];
106 if let Err(e) = stream.read_exact(&mut resp).await {
107 return err(format!("connect response read: {e}"));
108 }
109 if resp[1] != 0x00 {
110 return err_with_reply(
111 format!("SOCKS5 connect failed: reply code {:#04x}", resp[1]),
112 resp[1] as u16,
113 );
114 }
115
116 if let Err(e) = stream.write_all(payload).await {
117 return err(format!("payload write: {e}"));
118 }
119 let _ = stream.shutdown().await;
120
121 let mut response = Vec::new();
122 if let Err(e) = stream.read_to_end(&mut response).await {
123 return err(format!("response read: {e}"));
124 }
125
126 let bytes_sent = connect_bytes + payload.len() as u64;
127 let bytes_received = response.len() as u64;
128 ok(response, bytes_sent, bytes_received)
129}
130
131pub async fn socks5_tcp_connect_auth(
132 proxy: SocketAddr,
133 target: SocketAddr,
134 payload: &[u8],
135 username: &str,
136 password: &str,
137) -> ProbeResult {
138 let mut stream = match tokio::net::TcpStream::connect(proxy).await {
139 Ok(s) => s,
140 Err(e) => return err(format!("connect to proxy: {e}")),
141 };
142
143 if let Err(e) = stream.write_all(&[0x05, 0x01, 0x02]).await {
144 return err(format!("greeting write: {e}"));
145 }
146 let mut buf = [0u8; 2];
147 if let Err(e) = stream.read_exact(&mut buf).await {
148 return err(format!("greeting read: {e}"));
149 }
150 if buf[1] != 0x02 {
151 return err("proxy did not accept username/password auth method".to_string());
152 }
153
154 let mut auth = vec![0x01];
155 auth.push(username.len() as u8);
156 auth.extend_from_slice(username.as_bytes());
157 auth.push(password.len() as u8);
158 auth.extend_from_slice(password.as_bytes());
159 if let Err(e) = stream.write_all(&auth).await {
160 return err(format!("auth write: {e}"));
161 }
162 let mut auth_resp = [0u8; 2];
163 if let Err(e) = stream.read_exact(&mut auth_resp).await {
164 return err(format!("auth read: {e}"));
165 }
166 if auth_resp[1] != 0x00 {
167 return err_with_reply(
168 format!("SOCKS5 auth failed: reply code {:#04x}", auth_resp[1]),
169 auth_resp[1] as u16,
170 );
171 }
172
173 let req = build_sock5_target(target);
174 let connect_bytes = req.len() as u64;
175 if let Err(e) = stream.write_all(&req).await {
176 return err(format!("connect request write: {e}"));
177 }
178 let mut resp = [0u8; 10];
179 if let Err(e) = stream.read_exact(&mut resp).await {
180 return err(format!("connect response read: {e}"));
181 }
182 if resp[1] != 0x00 {
183 return err_with_reply(
184 format!("SOCKS5 connect failed: reply code {:#04x}", resp[1]),
185 resp[1] as u16,
186 );
187 }
188
189 if let Err(e) = stream.write_all(payload).await {
190 return err(format!("payload write: {e}"));
191 }
192 let _ = stream.shutdown().await;
193
194 let mut response = Vec::new();
195 if let Err(e) = stream.read_to_end(&mut response).await {
196 return err(format!("response read: {e}"));
197 }
198
199 let bytes_sent = connect_bytes + payload.len() as u64;
200 let bytes_received = response.len() as u64;
201 ok(response, bytes_sent, bytes_received)
202}
203
204pub async fn socks5_connect_refused(proxy: SocketAddr, target: SocketAddr) -> ProbeResult {
205 let mut stream = match tokio::net::TcpStream::connect(proxy).await {
206 Ok(s) => s,
207 Err(e) => return err(format!("connect to proxy: {e}")),
208 };
209
210 if let Err(e) = stream.write_all(&[0x05, 0x01, 0x00]).await {
211 return err(format!("greeting write: {e}"));
212 }
213 let mut buf = [0u8; 2];
214 if let Err(e) = stream.read_exact(&mut buf).await {
215 return err(format!("greeting read: {e}"));
216 }
217
218 let req = build_sock5_target(target);
219 if let Err(e) = stream.write_all(&req).await {
220 return err(format!("connect request write: {e}"));
221 }
222 let mut resp = [0u8; 10];
223 if let Err(e) = stream.read_exact(&mut resp).await {
224 return err(format!("connect response read: {e}"));
225 }
226
227 let code = resp[1];
228 if code == 0x00 {
229 err("expected SOCKS5 failure but got success".to_string())
230 } else {
231 err_with_reply(
232 format!("SOCKS5 connect refused: reply code {:#04x}", code),
233 code as u16,
234 )
235 }
236}
237
238pub async fn socks5_auth_failure(
239 proxy: SocketAddr,
240 _target: SocketAddr,
241 username: &str,
242 password: &str,
243) -> ProbeResult {
244 let mut stream = match tokio::net::TcpStream::connect(proxy).await {
245 Ok(s) => s,
246 Err(e) => return err(format!("connect to proxy: {e}")),
247 };
248
249 if let Err(e) = stream.write_all(&[0x05, 0x01, 0x02]).await {
250 return err(format!("greeting write: {e}"));
251 }
252 let mut buf = [0u8; 2];
253 if let Err(e) = stream.read_exact(&mut buf).await {
254 return err(format!("greeting read: {e}"));
255 }
256 if buf[1] != 0x02 {
257 return err("proxy did not accept auth method".to_string());
258 }
259
260 let mut auth = vec![0x01];
261 auth.push(username.len() as u8);
262 auth.extend_from_slice(username.as_bytes());
263 auth.push(password.len() as u8);
264 auth.extend_from_slice(password.as_bytes());
265 if let Err(e) = stream.write_all(&auth).await {
266 return err(format!("auth write: {e}"));
267 }
268 let mut auth_resp = [0u8; 2];
269 if let Err(e) = stream.read_exact(&mut auth_resp).await {
270 return err(format!("auth read: {e}"));
271 }
272
273 let code = auth_resp[1];
274 if code == 0x00 {
275 err("expected auth failure but got success".to_string())
276 } else {
277 err_with_reply(
278 format!("SOCKS5 auth rejected: reply code {:#04x}", code),
279 code as u16,
280 )
281 }
282}
283
284pub async fn http_connect(proxy: SocketAddr, target: SocketAddr, payload: &[u8]) -> ProbeResult {
285 let mut stream = match tokio::net::TcpStream::connect(proxy).await {
286 Ok(s) => s,
287 Err(e) => return err(format!("connect to proxy: {e}")),
288 };
289
290 let connect_req = format!("CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n");
291 let connect_bytes = connect_req.len() as u64;
292 if let Err(e) = stream.write_all(connect_req.as_bytes()).await {
293 return err(format!("CONNECT write: {e}"));
294 }
295
296 let mut resp = Vec::new();
297 let mut buf = [0u8; 4096];
298 loop {
299 let n = match stream.read(&mut buf).await {
300 Ok(0) => break,
301 Ok(n) => n,
302 Err(e) => return err(format!("CONNECT response read: {e}")),
303 };
304 resp.extend_from_slice(&buf[..n]);
305 if resp.windows(4).any(|w| w == b"\r\n\r\n") {
306 break;
307 }
308 }
309
310 let status_line = extract_http_status_line(&resp);
311 if !status_line.contains("200") {
312 return err(format!("CONNECT failed: {status_line}"));
313 }
314
315 if let Err(e) = stream.write_all(payload).await {
316 return err(format!("payload write: {e}"));
317 }
318 let _ = stream.shutdown().await;
319
320 let mut response = Vec::new();
321 if let Err(e) = stream.read_to_end(&mut response).await {
322 return err(format!("response read: {e}"));
323 }
324
325 let bytes_sent = connect_bytes + payload.len() as u64;
326 let bytes_received = response.len() as u64;
327 ok(response, bytes_sent, bytes_received)
328}
329
330pub async fn http_connect_refused(proxy: SocketAddr, target: SocketAddr) -> ProbeResult {
331 let mut stream = match tokio::net::TcpStream::connect(proxy).await {
332 Ok(s) => s,
333 Err(e) => return err(format!("connect to proxy: {e}")),
334 };
335
336 let connect_req = format!("CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n");
337 if let Err(e) = stream.write_all(connect_req.as_bytes()).await {
338 return err(format!("CONNECT write: {e}"));
339 }
340
341 let mut resp = Vec::new();
342 let mut buf = [0u8; 4096];
343 loop {
344 let n = match stream.read(&mut buf).await {
345 Ok(0) => break,
346 Ok(n) => n,
347 Err(e) => return err(format!("CONNECT response read: {e}")),
348 };
349 resp.extend_from_slice(&buf[..n]);
350 if resp.windows(4).any(|w| w == b"\r\n\r\n") {
351 break;
352 }
353 }
354
355 let status_line = extract_http_status_line(&resp);
356 if status_line.contains("200") {
357 err("expected CONNECT failure but got 200".to_string())
358 } else {
359 let code = status_line
360 .split_whitespace()
361 .nth(1)
362 .and_then(|s| s.parse::<u16>().ok());
363 ProbeResult {
364 success: false,
365 bytes_sent: 0,
366 bytes_received: resp.len() as u64,
367 response: resp,
368 error: Some(format!("CONNECT refused: {status_line}")),
369 reply_code: code,
370 reply_message: None,
371 }
372 }
373}
374
375pub async fn http_forward_get(proxy: SocketAddr, target: SocketAddr, path: &str) -> ProbeResult {
376 http_forward_inner(proxy, target, "GET", path, &[]).await
377}
378
379pub async fn http_forward_post(
380 proxy: SocketAddr,
381 target: SocketAddr,
382 path: &str,
383 body: &[u8],
384) -> ProbeResult {
385 http_forward_inner(proxy, target, "POST", path, body).await
386}
387
388async fn http_forward_inner(
389 proxy: SocketAddr,
390 target: SocketAddr,
391 method: &str,
392 path: &str,
393 body: &[u8],
394) -> ProbeResult {
395 let mut stream = match tokio::net::TcpStream::connect(proxy).await {
396 Ok(s) => s,
397 Err(e) => return err(format!("connect to proxy: {e}")),
398 };
399
400 let mut request = format!("{method} http://{target}{path} HTTP/1.1\r\nHost: {target}\r\n");
401 if !body.is_empty() {
402 request.push_str(&format!("Content-Length: {}\r\n", body.len()));
403 }
404 request.push_str("Connection: close\r\n\r\n");
405
406 let header_bytes = request.len() as u64;
407 if let Err(e) = stream.write_all(request.as_bytes()).await {
408 return err(format!("request write: {e}"));
409 }
410 if !body.is_empty() {
411 if let Err(e) = stream.write_all(body).await {
412 return err(format!("body write: {e}"));
413 }
414 }
415
416 let mut response = Vec::new();
417 if let Err(e) = stream.read_to_end(&mut response).await {
418 return err(format!("response read: {e}"));
419 }
420
421 let status_line = extract_http_status_line(&response);
422 let body_bytes = extract_http_body(&response);
423 let code = status_line
424 .split_whitespace()
425 .nth(1)
426 .and_then(|s| s.parse::<u16>().ok());
427
428 ProbeResult {
429 success: code.is_some_and(|c| (200..300).contains(&c)),
430 bytes_sent: header_bytes + body.len() as u64,
431 bytes_received: response.len() as u64,
432 response: body_bytes,
433 error: None,
434 reply_code: code,
435 reply_message: Some(status_line),
436 }
437}
438
439pub async fn socks4_connect(proxy: SocketAddr, target: SocketAddr, payload: &[u8]) -> ProbeResult {
440 let mut stream = match tokio::net::TcpStream::connect(proxy).await {
441 Ok(s) => s,
442 Err(e) => return err(format!("connect to proxy: {e}")),
443 };
444
445 let mut req = vec![0x04, 0x01];
446 req.extend_from_slice(&target.port().to_be_bytes());
447 match target.ip() {
448 std::net::IpAddr::V4(ip) => req.extend_from_slice(&ip.octets()),
449 std::net::IpAddr::V6(ip) => {
450 return err(format!("SOCKS4 does not support IPv6: {}", ip));
451 }
452 }
453 req.push(0x00);
454 let connect_bytes = req.len() as u64;
455 if let Err(e) = stream.write_all(&req).await {
456 return err(format!("connect request write: {e}"));
457 }
458
459 let mut resp = [0u8; 8];
460 if let Err(e) = stream.read_exact(&mut resp).await {
461 return err(format!("connect response read: {e}"));
462 }
463 if resp[1] != 0x5A {
464 return err_with_reply(
465 format!("SOCKS4 connect failed: reply code {:#04x}", resp[1]),
466 resp[1] as u16,
467 );
468 }
469
470 if let Err(e) = stream.write_all(payload).await {
471 return err(format!("payload write: {e}"));
472 }
473 let _ = stream.shutdown().await;
474
475 let mut response = Vec::new();
476 if let Err(e) = stream.read_to_end(&mut response).await {
477 return err(format!("response read: {e}"));
478 }
479
480 let bytes_sent = connect_bytes + payload.len() as u64;
481 let bytes_received = response.len() as u64;
482 ok(response, bytes_sent, bytes_received)
483}
484
485pub async fn socks4a_connect(
486 proxy: SocketAddr,
487 target_host: &str,
488 target_port: u16,
489 payload: &[u8],
490) -> ProbeResult {
491 let mut stream = match tokio::net::TcpStream::connect(proxy).await {
492 Ok(s) => s,
493 Err(e) => return err(format!("connect to proxy: {e}")),
494 };
495
496 let mut req = vec![0x04, 0x01];
497 req.extend_from_slice(&target_port.to_be_bytes());
498 req.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]);
499 req.push(0x00);
500 req.extend_from_slice(target_host.as_bytes());
501 req.push(0x00);
502 let connect_bytes = req.len() as u64;
503 if let Err(e) = stream.write_all(&req).await {
504 return err(format!("connect request write: {e}"));
505 }
506
507 let mut resp = [0u8; 8];
508 if let Err(e) = stream.read_exact(&mut resp).await {
509 return err(format!("connect response read: {e}"));
510 }
511 if resp[1] != 0x5A {
512 return err_with_reply(
513 format!("SOCKS4a connect failed: reply code {:#04x}", resp[1]),
514 resp[1] as u16,
515 );
516 }
517
518 if let Err(e) = stream.write_all(payload).await {
519 return err(format!("payload write: {e}"));
520 }
521 let _ = stream.shutdown().await;
522
523 let mut response = Vec::new();
524 if let Err(e) = stream.read_to_end(&mut response).await {
525 return err(format!("response read: {e}"));
526 }
527
528 let bytes_sent = connect_bytes + payload.len() as u64;
529 let bytes_received = response.len() as u64;
530 ok(response, bytes_sent, bytes_received)
531}
532
533#[cfg(test)]
534mod tests {
535 use super::*;
536
537 #[tokio::test]
538 async fn socks5_tcp_connect_echo() {
539 let (addr, jh) = crate::start_echo_server().await;
540 let proxy = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
541 let _ = (addr, proxy);
542 jh.abort();
543 }
544
545 #[tokio::test]
546 async fn http_forward_get_simple() {
547 let proxy = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
548 let target = "127.0.0.1:1".parse::<SocketAddr>().unwrap();
549 let _ = (proxy, target);
550 }
551}