1use std::io;
11use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
12use std::time::Duration;
13
14use tokio::io::{AsyncReadExt, AsyncWriteExt};
15use tokio::net::{TcpStream, UdpSocket};
16use tokio_util::sync::CancellationToken;
17
18pub(crate) const VER: u8 = 0x05;
20const CMD_UDP_ASSOCIATE: u8 = 0x03;
21const ATYP_IPV4: u8 = 0x01;
22const ATYP_DOMAIN: u8 = 0x03;
23const ATYP_IPV6: u8 = 0x04;
24const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
26
27pub(crate) async fn greet(stream: &mut TcpStream) -> io::Result<()> {
30 stream.write_all(&[VER, 0x01, 0x00]).await?;
31 let mut resp = [0u8; 2];
32 stream.read_exact(&mut resp).await?;
33 if resp[0] != VER {
34 return Err(io::Error::other(format!(
35 "SOCKS5: unsupported version {}",
36 resp[0]
37 )));
38 }
39 if resp[1] != 0x00 {
40 return Err(io::Error::other(format!(
41 "SOCKS5: server chose unsupported auth method 0x{:02X} (expected 0x00 no-auth)",
42 resp[1]
43 )));
44 }
45 Ok(())
46}
47
48pub fn encode_addr(out: &mut Vec<u8>, host: &str, port: u16) {
52 match host.parse::<IpAddr>() {
53 Ok(IpAddr::V4(ip)) => {
54 out.push(ATYP_IPV4);
55 out.extend_from_slice(&ip.octets());
56 }
57 Ok(IpAddr::V6(ip)) => {
58 out.push(ATYP_IPV6);
59 out.extend_from_slice(&ip.octets());
60 }
61 Err(_) => {
62 out.push(ATYP_DOMAIN);
63 let bytes = host.as_bytes();
69 debug_assert!(
70 bytes.len() <= 255,
71 "SOCKS5 domain exceeds 255 bytes ({} bytes): {host}",
72 bytes.len()
73 );
74 let len = bytes.len().min(255);
75 out.push(len as u8);
76 out.extend_from_slice(&bytes[..len]);
77 }
78 }
79 out.extend_from_slice(&port.to_be_bytes());
80}
81
82#[must_use]
86pub fn decode_addr(buf: &[u8]) -> Option<(String, u16, usize)> {
87 let (host, after) = match *buf.first()? {
88 ATYP_IPV4 => {
89 let o: [u8; 4] = buf.get(1..5)?.try_into().ok()?;
90 (Ipv4Addr::from(o).to_string(), 5)
91 }
92 ATYP_IPV6 => {
93 let o: [u8; 16] = buf.get(1..17)?.try_into().ok()?;
94 (Ipv6Addr::from(o).to_string(), 17)
95 }
96 ATYP_DOMAIN => {
97 let len = *buf.get(1)? as usize;
98 let name = std::str::from_utf8(buf.get(2..2 + len)?).ok()?.to_string();
99 (name, 2 + len)
100 }
101 _ => return None,
102 };
103 let port = u16::from_be_bytes([*buf.get(after)?, *buf.get(after + 1)?]);
104 Some((host, port, after + 2))
105}
106
107pub async fn read_addr<R: AsyncReadExt + Unpin>(r: &mut R) -> io::Result<(String, u16)> {
110 let mut atyp = [0u8; 1];
111 r.read_exact(&mut atyp).await?;
112 let host = match atyp[0] {
113 ATYP_IPV4 => {
114 let mut o = [0u8; 4];
115 r.read_exact(&mut o).await?;
116 Ipv4Addr::from(o).to_string()
117 }
118 ATYP_IPV6 => {
119 let mut o = [0u8; 16];
120 r.read_exact(&mut o).await?;
121 Ipv6Addr::from(o).to_string()
122 }
123 ATYP_DOMAIN => {
124 let mut len = [0u8; 1];
125 r.read_exact(&mut len).await?;
126 let mut name = vec![0u8; len[0] as usize];
127 r.read_exact(&mut name).await?;
128 String::from_utf8(name).map_err(|_| {
129 io::Error::new(io::ErrorKind::InvalidData, "SOCKS5: non-UTF8 domain")
130 })?
131 }
132 other => {
133 return Err(io::Error::other(format!(
134 "SOCKS5: unsupported ATYP 0x{other:02X} in response"
135 )));
136 }
137 };
138 let mut port = [0u8; 2];
139 r.read_exact(&mut port).await?;
140 Ok((host, u16::from_be_bytes(port)))
141}
142
143pub(crate) fn rep_error(code: u8) -> io::Error {
145 let reason = match code {
146 0x01 => "general SOCKS server failure",
147 0x02 => "connection not allowed by ruleset",
148 0x03 => "network unreachable",
149 0x04 => "host unreachable",
150 0x05 => "connection refused",
151 0x06 => "TTL expired",
152 0x07 => "command not supported",
153 0x08 => "address type not supported",
154 _ => "unknown error",
155 };
156 io::Error::other(format!("SOCKS5: request failed: {reason} (code {code})"))
157}
158
159fn unspecified_in_family(addr: SocketAddr) -> SocketAddr {
163 if addr.is_ipv6() {
164 (Ipv6Addr::UNSPECIFIED, 0).into()
165 } else {
166 (Ipv4Addr::UNSPECIFIED, 0).into()
167 }
168}
169
170pub async fn connect_via_socks5(proxy: &str, host: &str, port: u16) -> io::Result<TcpStream> {
176 let mut stream = TcpStream::connect(proxy).await?;
177
178 greet(&mut stream).await?;
180
181 let mut req = vec![VER, 0x01, 0x00]; encode_addr(&mut req, host, port);
184 stream.write_all(&req).await?;
185
186 let mut hdr = [0u8; 3];
188 stream.read_exact(&mut hdr).await?;
189 if hdr[0] != VER {
190 return Err(io::Error::other(format!(
191 "SOCKS5: unexpected version in response: {}",
192 hdr[0]
193 )));
194 }
195 if hdr[1] != 0x00 {
196 return Err(rep_error(hdr[1]));
197 }
198 let _bind = read_addr(&mut stream).await?;
199
200 tracing::debug!(
201 proxy = %proxy,
202 target = %format!("{host}:{port}"),
203 "SOCKS5 tunnel established"
204 );
205 Ok(stream)
206}
207
208#[derive(Debug)]
212struct WatcherGuard(tokio::task::JoinHandle<()>);
213
214impl Drop for WatcherGuard {
215 fn drop(&mut self) {
216 self.0.abort();
217 }
218}
219
220#[derive(Debug)]
228pub struct Socks5UdpAssociation {
229 socket: UdpSocket,
230 closed: CancellationToken,
233 _watcher: WatcherGuard,
235}
236
237impl Socks5UdpAssociation {
238 pub async fn associate(proxy: &str) -> io::Result<Self> {
241 tokio::time::timeout(HANDSHAKE_TIMEOUT, Self::handshake(proxy))
242 .await
243 .map_err(|_| {
244 io::Error::new(io::ErrorKind::TimedOut, "SOCKS5 UDP associate timed out")
245 })?
246 }
247
248 async fn handshake(proxy: &str) -> io::Result<Self> {
249 let mut control = TcpStream::connect(proxy).await?;
250 greet(&mut control).await?;
251
252 let mut req = vec![VER, CMD_UDP_ASSOCIATE, 0x00];
255 encode_addr(&mut req, "0.0.0.0", 0);
256 control.write_all(&req).await?;
257
258 let mut hdr = [0u8; 3];
260 control.read_exact(&mut hdr).await?;
261 if hdr[0] != VER {
262 return Err(io::Error::other(format!(
263 "SOCKS5: unexpected version in response: {}",
264 hdr[0]
265 )));
266 }
267 if hdr[1] != 0x00 {
268 return Err(rep_error(hdr[1]));
269 }
270 let (bnd_host, bnd_port) = read_addr(&mut control).await?;
271
272 let relay: SocketAddr = match bnd_host.parse::<IpAddr>() {
275 Ok(ip) if ip.is_unspecified() => {
276 let mut peer = control.peer_addr()?;
277 peer.set_port(bnd_port);
278 peer
279 }
280 Ok(ip) => SocketAddr::new(ip, bnd_port),
281 Err(_) => tokio::net::lookup_host((bnd_host.as_str(), bnd_port))
282 .await?
283 .next()
284 .ok_or_else(|| {
285 io::Error::new(io::ErrorKind::InvalidData, "SOCKS5: BND address unresolved")
286 })?,
287 };
288
289 let socket = UdpSocket::bind(unspecified_in_family(relay)).await?;
292 socket.connect(relay).await?;
293
294 let closed = CancellationToken::new();
299 let watcher_token = closed.clone();
300 let watcher = tokio::spawn(async move {
301 let mut control = control;
302 let mut scratch = [0u8; 1];
303 let _ = control.read(&mut scratch).await;
304 watcher_token.cancel();
305 });
306
307 tracing::debug!(proxy = %proxy, relay = %relay, "SOCKS5 UDP association established");
308 Ok(Self {
309 socket,
310 closed,
311 _watcher: WatcherGuard(watcher),
312 })
313 }
314
315 pub async fn send_to(&self, payload: &[u8], host: &str, port: u16) -> io::Result<usize> {
318 let mut dgram = vec![0x00, 0x00, 0x00]; encode_addr(&mut dgram, host, port);
320 dgram.extend_from_slice(payload);
321 let sent = self.socket.send(&dgram).await?;
322 debug_assert_eq!(
323 sent,
324 dgram.len(),
325 "UDP short-send: {sent} of {} bytes",
326 dgram.len()
327 );
328 Ok(payload.len())
329 }
330
331 pub async fn recv_from(&self, buf: &mut [u8]) -> io::Result<Option<(usize, String, u16)>> {
339 let n = tokio::select! {
340 () = self.closed.cancelled() => {
341 return Err(io::Error::new(
342 io::ErrorKind::ConnectionAborted,
343 "SOCKS5 UDP: control connection closed",
344 ));
345 }
346 r = self.socket.recv(buf) => r?,
347 };
348
349 if n < 4 || buf[2] != 0x00 {
352 return Ok(None);
353 }
354 let Some((host, port, consumed)) = decode_addr(&buf[3..]) else {
355 return Ok(None);
356 };
357 let start = 3 + consumed;
358 if start > n {
359 return Ok(None);
360 }
361 let len = n - start;
362 buf.copy_within(start..n, 0);
363 Ok(Some((len, host, port)))
364 }
365}
366
367#[cfg(test)]
368mod tests {
369 use super::*;
370 use tokio::net::TcpListener;
371
372 #[test]
373 fn addr_codec_round_trips() {
374 for (host, port) in [("example.com", 443u16), ("1.2.3.4", 53), ("::1", 80)] {
375 let mut buf = Vec::new();
376 encode_addr(&mut buf, host, port);
377 let (h, p, used) = decode_addr(&buf).expect("decodes");
378 assert_eq!((h.as_str(), p, used), (host, port, buf.len()));
379 }
380 }
381
382 #[test]
383 fn bind_family_matches_relay() {
384 let v6: SocketAddr = "[::1]:9".parse().unwrap();
386 let v4: SocketAddr = "1.2.3.4:9".parse().unwrap();
387 assert!(
388 unspecified_in_family(v6).is_ipv6(),
389 "IPv6 relay → IPv6 bind"
390 );
391 assert!(
392 unspecified_in_family(v4).is_ipv4(),
393 "IPv4 relay → IPv4 bind"
394 );
395 }
396
397 async fn mock_socks5(listener: TcpListener, expect_host: &'static str, expect_port: u16) {
401 let (mut s, _) = listener.accept().await.unwrap();
402
403 let mut greeting = [0u8; 3];
405 s.read_exact(&mut greeting).await.unwrap();
406 assert_eq!(greeting, [0x05, 0x01, 0x00], "client greeting");
407 s.write_all(&[0x05, 0x00]).await.unwrap(); let mut hdr = [0u8; 4];
411 s.read_exact(&mut hdr).await.unwrap();
412 assert_eq!(hdr, [0x05, 0x01, 0x00, 0x03], "connect request header");
413 let mut len = [0u8; 1];
414 s.read_exact(&mut len).await.unwrap();
415 let mut host = vec![0u8; len[0] as usize];
416 s.read_exact(&mut host).await.unwrap();
417 let mut port = [0u8; 2];
418 s.read_exact(&mut port).await.unwrap();
419 assert_eq!(host, expect_host.as_bytes(), "CONNECT host (by name)");
420 assert_eq!(u16::from_be_bytes(port), expect_port, "CONNECT port");
421
422 s.write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
424 .await
425 .unwrap();
426
427 let mut buf = [0u8; 5];
429 s.read_exact(&mut buf).await.unwrap();
430 s.write_all(&buf).await.unwrap();
431 }
432
433 #[tokio::test]
434 async fn socks5_connects_by_domain_and_tunnels() {
435 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
436 let addr = listener.local_addr().unwrap();
437 let server = tokio::spawn(mock_socks5(listener, "example.com", 443));
438
439 let mut stream = connect_via_socks5(&addr.to_string(), "example.com", 443)
440 .await
441 .expect("SOCKS5 handshake should complete");
442
443 stream.write_all(b"hello").await.unwrap();
445 let mut got = [0u8; 5];
446 stream.read_exact(&mut got).await.unwrap();
447 assert_eq!(&got, b"hello", "payload round-trips through the tunnel");
448 server.await.unwrap();
449 }
450
451 #[tokio::test]
452 async fn socks5_propagates_server_failure() {
453 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
454 let addr = listener.local_addr().unwrap();
455 tokio::spawn(async move {
456 let (mut s, _) = listener.accept().await.unwrap();
457 let mut greeting = [0u8; 3];
458 s.read_exact(&mut greeting).await.unwrap();
459 s.write_all(&[0x05, 0x00]).await.unwrap();
460 let mut hdr = [0u8; 4];
462 s.read_exact(&mut hdr).await.unwrap();
463 let mut len = [0u8; 1];
464 s.read_exact(&mut len).await.unwrap();
465 let mut rest = vec![0u8; len[0] as usize + 2];
466 s.read_exact(&mut rest).await.unwrap();
467 s.write_all(&[0x05, 0x05, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
468 .await
469 .unwrap();
470 });
471
472 let err = connect_via_socks5(&addr.to_string(), "blocked.example", 443)
473 .await
474 .expect_err("server failure must surface as an error");
475 assert!(err.to_string().contains("connection refused"), "{err}");
476 }
477
478 async fn mock_socks5_udp(listener: TcpListener) {
482 let relay = UdpSocket::bind("127.0.0.1:0").await.unwrap();
483 let relay_addr = relay.local_addr().unwrap();
484
485 let (mut ctrl, _) = listener.accept().await.unwrap();
486 let mut g = [0u8; 3];
487 ctrl.read_exact(&mut g).await.unwrap();
488 assert_eq!(g, [0x05, 0x01, 0x00], "greeting");
489 ctrl.write_all(&[0x05, 0x00]).await.unwrap();
490
491 let mut req = [0u8; 4];
493 ctrl.read_exact(&mut req).await.unwrap();
494 assert_eq!(&req[..2], &[0x05, 0x03], "UDP ASSOCIATE cmd");
495 assert_eq!(req[3], 0x01, "wildcard DST is ATYP v4");
496 let mut rest = [0u8; 6];
497 ctrl.read_exact(&mut rest).await.unwrap();
498
499 let mut reply = vec![0x05, 0x00, 0x00];
501 encode_addr(&mut reply, &relay_addr.ip().to_string(), relay_addr.port());
502 ctrl.write_all(&reply).await.unwrap();
503
504 let mut buf = vec![0u8; 2048];
506 let (n, client) = relay.recv_from(&mut buf).await.unwrap();
507 let pkt = &buf[..n];
508 assert_eq!(pkt[2], 0x00, "FRAG 0");
509 let (host, port, consumed) = decode_addr(&pkt[3..]).unwrap();
510 let payload = &pkt[3 + consumed..];
511 let mut echo = vec![0x00, 0x00, 0x00];
512 encode_addr(&mut echo, &host, port);
513 echo.extend_from_slice(payload);
514 relay.send_to(&echo, client).await.unwrap();
515
516 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
518 drop(ctrl);
519 }
520
521 #[tokio::test]
522 async fn udp_associate_round_trips_through_relay() {
523 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
524 let proxy = listener.local_addr().unwrap();
525 let server = tokio::spawn(mock_socks5_udp(listener));
526
527 let assoc = Socks5UdpAssociation::associate(&proxy.to_string())
528 .await
529 .expect("association established");
530 assoc.send_to(b"ping", "example.com", 53).await.unwrap();
531
532 let mut buf = [0u8; 64];
533 let (n, host, port) = assoc
534 .recv_from(&mut buf)
535 .await
536 .unwrap()
537 .expect("a deliverable datagram");
538 assert_eq!(&buf[..n], b"ping", "payload round-trips through the relay");
539 assert_eq!(
540 (host.as_str(), port),
541 ("example.com", 53),
542 "target echoed back"
543 );
544 server.await.unwrap();
545 }
546
547 #[tokio::test]
548 async fn udp_associate_surfaces_rep_error() {
549 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
550 let proxy = listener.local_addr().unwrap();
551 tokio::spawn(async move {
552 let (mut ctrl, _) = listener.accept().await.unwrap();
553 let mut g = [0u8; 3];
554 ctrl.read_exact(&mut g).await.unwrap();
555 ctrl.write_all(&[0x05, 0x00]).await.unwrap();
556 let mut req = [0u8; 4];
557 ctrl.read_exact(&mut req).await.unwrap();
558 let mut rest = [0u8; 6];
559 ctrl.read_exact(&mut rest).await.unwrap();
560 ctrl.write_all(&[0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
562 .await
563 .unwrap();
564 });
565
566 let err = Socks5UdpAssociation::associate(&proxy.to_string())
567 .await
568 .expect_err("REP failure surfaces");
569 assert!(err.to_string().contains("command not supported"), "{err}");
570 }
571
572 #[tokio::test]
574 async fn fragmented_datagram_is_skipped_not_fatal() {
575 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
576 let proxy = listener.local_addr().unwrap();
577 tokio::spawn(async move {
578 let relay = UdpSocket::bind("127.0.0.1:0").await.unwrap();
579 let relay_addr = relay.local_addr().unwrap();
580 let (mut ctrl, _) = listener.accept().await.unwrap();
581 let mut g = [0u8; 3];
582 ctrl.read_exact(&mut g).await.unwrap();
583 ctrl.write_all(&[0x05, 0x00]).await.unwrap();
584 let mut req = [0u8; 4];
585 ctrl.read_exact(&mut req).await.unwrap();
586 let mut rest = [0u8; 6];
587 ctrl.read_exact(&mut rest).await.unwrap();
588 let mut reply = vec![0x05, 0x00, 0x00];
589 encode_addr(&mut reply, &relay_addr.ip().to_string(), relay_addr.port());
590 ctrl.write_all(&reply).await.unwrap();
591 let (_n, client) = relay.recv_from(&mut [0u8; 64]).await.unwrap();
592 let mut frag = vec![0x00, 0x00, 0x01];
594 encode_addr(&mut frag, "1.2.3.4", 9);
595 frag.extend_from_slice(b"x");
596 relay.send_to(&frag, client).await.unwrap();
597 tokio::time::sleep(std::time::Duration::from_millis(80)).await;
598 drop(ctrl);
599 });
600
601 let assoc = Socks5UdpAssociation::associate(&proxy.to_string())
602 .await
603 .unwrap();
604 assoc.send_to(b"hi", "1.2.3.4", 9).await.unwrap();
605 let mut buf = [0u8; 64];
606 let got = assoc.recv_from(&mut buf).await.unwrap();
607 assert!(
608 got.is_none(),
609 "fragmented datagram is skipped, flow survives"
610 );
611 }
612}