Skip to main content

irontide_session/
proxy.rs

1//! Proxy support: SOCKS4, SOCKS5, HTTP CONNECT, and SOCKS5 UDP ASSOCIATE.
2//!
3//! Hand-rolled implementations — no external dependency beyond `tokio`.
4//! Provides [`connect_through_proxy`] for TCP tunneling and
5//! [`socks5_udp_associate`] + [`ProxiedUdpSocket`] for UDP relay.
6
7use std::io;
8use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
9
10use serde::{Deserialize, Serialize};
11use tokio::io::{AsyncReadExt, AsyncWriteExt};
12use tokio::net::{TcpStream, UdpSocket};
13
14// ── ProxyType ────────────────────────────────────────────────────────
15
16/// Supported proxy protocols (matching libtorrent).
17#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
18pub enum ProxyType {
19    /// No proxy (direct connections).
20    #[default]
21    None,
22    /// SOCKS4 proxy (no authentication, no UDP).
23    Socks4,
24    /// SOCKS5 proxy without authentication.
25    Socks5,
26    /// SOCKS5 proxy with username/password authentication.
27    Socks5Password,
28    /// HTTP CONNECT proxy without authentication.
29    Http,
30    /// HTTP CONNECT proxy with username/password authentication.
31    HttpPassword,
32}
33
34// ── ProxyConfig ──────────────────────────────────────────────────────
35
36/// Proxy connection settings.
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct ProxyConfig {
39    /// Proxy protocol to use.
40    pub proxy_type: ProxyType,
41    /// Proxy server hostname or IP address.
42    pub hostname: String,
43    /// Proxy server port.
44    pub port: u16,
45    /// Username for authenticated proxy types.
46    pub username: Option<String>,
47    /// Password for authenticated proxy types.
48    pub password: Option<String>,
49    /// Route peer connections (incl. web seeds) through proxy.
50    #[serde(default = "default_true")]
51    pub proxy_peer_connections: bool,
52    /// Route tracker HTTP connections through proxy.
53    #[serde(default = "default_true")]
54    pub proxy_tracker_connections: bool,
55    /// Resolve hostnames through proxy (SOCKS5/HTTP only).
56    #[serde(default = "default_true")]
57    pub proxy_hostnames: bool,
58    /// Include local endpoint in SOCKS5 UDP ASSOCIATE.
59    #[serde(default)]
60    pub socks5_udp_send_local_ep: bool,
61}
62
63fn default_true() -> bool {
64    true
65}
66
67impl Default for ProxyConfig {
68    fn default() -> Self {
69        Self {
70            proxy_type: ProxyType::None,
71            hostname: String::new(),
72            port: 0,
73            username: None,
74            password: None,
75            proxy_peer_connections: true,
76            proxy_tracker_connections: true,
77            proxy_hostnames: true,
78            socks5_udp_send_local_ep: false,
79        }
80    }
81}
82
83impl ProxyConfig {
84    /// Format as a URL suitable for `reqwest::Proxy::all()`.
85    pub fn to_url(&self) -> String {
86        let scheme = match self.proxy_type {
87            ProxyType::None => return String::new(),
88            ProxyType::Socks4 => "socks4",
89            ProxyType::Socks5 | ProxyType::Socks5Password => "socks5",
90            ProxyType::Http | ProxyType::HttpPassword => "http",
91        };
92
93        match (&self.username, &self.password) {
94            (Some(u), Some(p))
95                if self.proxy_type == ProxyType::Socks5Password
96                    || self.proxy_type == ProxyType::HttpPassword =>
97            {
98                format!("{}://{}:{}@{}:{}", scheme, u, p, self.hostname, self.port)
99            }
100            _ => format!("{}://{}:{}", scheme, self.hostname, self.port),
101        }
102    }
103}
104
105// ── connect_through_proxy ────────────────────────────────────────────
106
107/// Connect to `target` through the configured proxy.
108///
109/// Returns a `TcpStream` with the proxy handshake already completed.
110/// The stream is ready for BitTorrent protocol I/O.
111pub(crate) async fn connect_through_proxy(
112    proxy: &ProxyConfig,
113    target: SocketAddr,
114) -> io::Result<TcpStream> {
115    let proxy_addr = format!("{}:{}", proxy.hostname, proxy.port);
116    let mut stream = TcpStream::connect(&proxy_addr).await?;
117
118    match proxy.proxy_type {
119        ProxyType::Socks4 => socks4_connect(&mut stream, target).await?,
120        ProxyType::Socks5 => socks5_connect(&mut stream, target, None).await?,
121        ProxyType::Socks5Password => {
122            let auth = match (&proxy.username, &proxy.password) {
123                (Some(u), Some(p)) => Some((u.as_str(), p.as_str())),
124                _ => None,
125            };
126            socks5_connect(&mut stream, target, auth).await?;
127        }
128        ProxyType::Http => http_connect(&mut stream, target, None).await?,
129        ProxyType::HttpPassword => {
130            let auth = match (&proxy.username, &proxy.password) {
131                (Some(u), Some(p)) => Some((u.as_str(), p.as_str())),
132                _ => None,
133            };
134            http_connect(&mut stream, target, auth).await?;
135        }
136        ProxyType::None => {
137            return Err(io::Error::new(
138                io::ErrorKind::InvalidInput,
139                "no proxy configured",
140            ));
141        }
142    }
143
144    Ok(stream)
145}
146
147// ── SOCKS4 ───────────────────────────────────────────────────────────
148
149/// SOCKS4 CONNECT handshake (RFC 1928 predecessor).
150///
151/// Request: VN(1)=4, CD(1)=1, DSTPORT(2), DSTIP(4), USERID(var), NULL(1)
152/// Response: VN(1)=0, CD(1), DSTPORT(2), DSTIP(4) — CD=90 means granted
153async fn socks4_connect(stream: &mut TcpStream, target: SocketAddr) -> io::Result<()> {
154    let ip = match target.ip() {
155        IpAddr::V4(v4) => v4,
156        IpAddr::V6(_) => {
157            return Err(io::Error::new(
158                io::ErrorKind::InvalidInput,
159                "SOCKS4 does not support IPv6",
160            ));
161        }
162    };
163
164    // Build request
165    let mut req = Vec::with_capacity(9);
166    req.push(4); // VN
167    req.push(1); // CD = CONNECT
168    req.extend_from_slice(&target.port().to_be_bytes());
169    req.extend_from_slice(&ip.octets());
170    req.push(0); // NULL-terminated empty userid
171    stream.write_all(&req).await?;
172
173    // Read 8-byte response
174    let mut resp = [0u8; 8];
175    stream.read_exact(&mut resp).await?;
176
177    if resp[1] != 90 {
178        return Err(io::Error::new(
179            io::ErrorKind::ConnectionRefused,
180            format!("SOCKS4 request rejected (CD={})", resp[1]),
181        ));
182    }
183
184    Ok(())
185}
186
187// ── SOCKS5 ───────────────────────────────────────────────────────────
188
189/// SOCKS5 CONNECT handshake (RFC 1928 + optional RFC 1929 auth).
190async fn socks5_connect(
191    stream: &mut TcpStream,
192    target: SocketAddr,
193    auth: Option<(&str, &str)>,
194) -> io::Result<()> {
195    // Step 1: Method negotiation
196    socks5_negotiate_method(stream, auth.is_some()).await?;
197
198    // Step 2: Username/password auth if needed
199    if let Some((user, pass)) = auth {
200        socks5_auth(stream, user, pass).await?;
201    }
202
203    // Step 3: Connect request
204    socks5_send_connect(stream, target).await
205}
206
207/// Negotiate authentication method with SOCKS5 proxy.
208async fn socks5_negotiate_method(stream: &mut TcpStream, want_auth: bool) -> io::Result<()> {
209    let methods: &[u8] = if want_auth {
210        &[5, 2, 0, 2] // VER=5, NMETHODS=2, NO_AUTH + USERNAME/PASSWORD
211    } else {
212        &[5, 1, 0] // VER=5, NMETHODS=1, NO_AUTH
213    };
214    stream.write_all(methods).await?;
215
216    let mut resp = [0u8; 2];
217    stream.read_exact(&mut resp).await?;
218
219    if resp[0] != 5 {
220        return Err(io::Error::new(
221            io::ErrorKind::InvalidData,
222            format!("SOCKS5: unexpected version {}", resp[0]),
223        ));
224    }
225
226    match resp[1] {
227        0 => Ok(()),              // NO_AUTH accepted
228        2 if want_auth => Ok(()), // USERNAME/PASSWORD accepted
229        0xFF => Err(io::Error::new(
230            io::ErrorKind::PermissionDenied,
231            "SOCKS5: no acceptable methods",
232        )),
233        m => Err(io::Error::new(
234            io::ErrorKind::InvalidData,
235            format!("SOCKS5: unexpected method {}", m),
236        )),
237    }
238}
239
240/// RFC 1929 username/password subnegotiation.
241///
242/// Request: VER(1)=1, ULEN(1), UNAME(ULEN), PLEN(1), PASSWD(PLEN)
243/// Response: VER(1)=1, STATUS(1) — 0=success
244async fn socks5_auth(stream: &mut TcpStream, user: &str, pass: &str) -> io::Result<()> {
245    if user.len() > 255 || pass.len() > 255 {
246        return Err(io::Error::new(
247            io::ErrorKind::InvalidInput,
248            "SOCKS5: username or password too long (max 255 bytes)",
249        ));
250    }
251
252    let mut req = Vec::with_capacity(3 + user.len() + pass.len());
253    req.push(1); // subneg version
254    req.push(user.len() as u8);
255    req.extend_from_slice(user.as_bytes());
256    req.push(pass.len() as u8);
257    req.extend_from_slice(pass.as_bytes());
258    stream.write_all(&req).await?;
259
260    let mut resp = [0u8; 2];
261    stream.read_exact(&mut resp).await?;
262
263    if resp[1] != 0 {
264        return Err(io::Error::new(
265            io::ErrorKind::PermissionDenied,
266            "SOCKS5: authentication failed",
267        ));
268    }
269
270    Ok(())
271}
272
273/// Send SOCKS5 CONNECT request and read response.
274///
275/// Request: VER(1)=5, CMD(1)=1, RSV(1)=0, ATYP(1), DST.ADDR(var), DST.PORT(2)
276/// Response: VER(1)=5, REP(1), RSV(1)=0, ATYP(1), BND.ADDR(var), BND.PORT(2)
277async fn socks5_send_connect(stream: &mut TcpStream, target: SocketAddr) -> io::Result<()> {
278    let mut req = Vec::with_capacity(22);
279    req.push(5); // VER
280    req.push(1); // CMD = CONNECT
281    req.push(0); // RSV
282    encode_socks5_addr(&mut req, target);
283    stream.write_all(&req).await?;
284
285    read_socks5_response(stream).await
286}
287
288/// Encode a socket address in SOCKS5 format (ATYP + ADDR + PORT).
289fn encode_socks5_addr(buf: &mut Vec<u8>, addr: SocketAddr) {
290    match addr {
291        SocketAddr::V4(v4) => {
292            buf.push(1); // ATYP = IPv4
293            buf.extend_from_slice(&v4.ip().octets());
294            buf.extend_from_slice(&v4.port().to_be_bytes());
295        }
296        SocketAddr::V6(v6) => {
297            buf.push(4); // ATYP = IPv6
298            buf.extend_from_slice(&v6.ip().octets());
299            buf.extend_from_slice(&v6.port().to_be_bytes());
300        }
301    }
302}
303
304/// Read and validate a SOCKS5 response (CONNECT or UDP ASSOCIATE).
305async fn read_socks5_response(stream: &mut TcpStream) -> io::Result<()> {
306    let mut header = [0u8; 4];
307    stream.read_exact(&mut header).await?;
308
309    if header[0] != 5 {
310        return Err(io::Error::new(
311            io::ErrorKind::InvalidData,
312            format!("SOCKS5: unexpected version {}", header[0]),
313        ));
314    }
315
316    if header[1] != 0 {
317        let msg = socks5_error_message(header[1]);
318        return Err(io::Error::new(io::ErrorKind::ConnectionRefused, msg));
319    }
320
321    // Skip bound address
322    match header[3] {
323        1 => {
324            // IPv4: 4 bytes + 2 port
325            let mut skip = [0u8; 6];
326            stream.read_exact(&mut skip).await?;
327        }
328        4 => {
329            // IPv6: 16 bytes + 2 port
330            let mut skip = [0u8; 18];
331            stream.read_exact(&mut skip).await?;
332        }
333        3 => {
334            // Domain: 1 len + domain + 2 port
335            let mut len = [0u8; 1];
336            stream.read_exact(&mut len).await?;
337            let mut skip = vec![0u8; len[0] as usize + 2];
338            stream.read_exact(&mut skip).await?;
339        }
340        atyp => {
341            return Err(io::Error::new(
342                io::ErrorKind::InvalidData,
343                format!("SOCKS5: unknown ATYP {}", atyp),
344            ));
345        }
346    }
347
348    Ok(())
349}
350
351/// Read a SOCKS5 response and return the bound address.
352async fn read_socks5_response_addr(stream: &mut TcpStream) -> io::Result<SocketAddr> {
353    let mut header = [0u8; 4];
354    stream.read_exact(&mut header).await?;
355
356    if header[0] != 5 {
357        return Err(io::Error::new(
358            io::ErrorKind::InvalidData,
359            format!("SOCKS5: unexpected version {}", header[0]),
360        ));
361    }
362
363    if header[1] != 0 {
364        let msg = socks5_error_message(header[1]);
365        return Err(io::Error::new(io::ErrorKind::ConnectionRefused, msg));
366    }
367
368    match header[3] {
369        1 => {
370            let mut buf = [0u8; 6];
371            stream.read_exact(&mut buf).await?;
372            let ip = Ipv4Addr::new(buf[0], buf[1], buf[2], buf[3]);
373            let port = u16::from_be_bytes([buf[4], buf[5]]);
374            Ok(SocketAddr::V4(SocketAddrV4::new(ip, port)))
375        }
376        4 => {
377            let mut buf = [0u8; 18];
378            stream.read_exact(&mut buf).await?;
379            let ip = Ipv6Addr::from(<[u8; 16]>::try_from(&buf[..16]).unwrap());
380            let port = u16::from_be_bytes([buf[16], buf[17]]);
381            Ok(SocketAddr::V6(SocketAddrV6::new(ip, port, 0, 0)))
382        }
383        atyp => Err(io::Error::new(
384            io::ErrorKind::InvalidData,
385            format!("SOCKS5: unsupported BND.ATYP {} in UDP ASSOCIATE", atyp),
386        )),
387    }
388}
389
390fn socks5_error_message(code: u8) -> String {
391    match code {
392        1 => "SOCKS5: general failure".into(),
393        2 => "SOCKS5: connection not allowed by ruleset".into(),
394        3 => "SOCKS5: network unreachable".into(),
395        4 => "SOCKS5: host unreachable".into(),
396        5 => "SOCKS5: connection refused".into(),
397        6 => "SOCKS5: TTL expired".into(),
398        7 => "SOCKS5: command not supported".into(),
399        8 => "SOCKS5: address type not supported".into(),
400        _ => format!("SOCKS5: unknown error ({})", code),
401    }
402}
403
404// ── HTTP CONNECT ─────────────────────────────────────────────────────
405
406/// HTTP CONNECT tunnel handshake.
407async fn http_connect(
408    stream: &mut TcpStream,
409    target: SocketAddr,
410    auth: Option<(&str, &str)>,
411) -> io::Result<()> {
412    let host_port = format!("{}:{}", target.ip(), target.port());
413
414    let mut request = format!("CONNECT {} HTTP/1.1\r\nHost: {}\r\n", host_port, host_port,);
415
416    if let Some((user, pass)) = auth {
417        use std::fmt::Write;
418        let credentials = format!("{}:{}", user, pass);
419        let encoded = base64_encode(credentials.as_bytes());
420        let _ = write!(request, "Proxy-Authorization: Basic {}\r\n", encoded);
421    }
422
423    request.push_str("\r\n");
424    stream.write_all(request.as_bytes()).await?;
425
426    // Read until we see \r\n\r\n (end of HTTP headers)
427    let mut response_buf = Vec::with_capacity(256);
428    loop {
429        let mut byte = [0u8; 1];
430        stream.read_exact(&mut byte).await?;
431        response_buf.push(byte[0]);
432
433        if response_buf.len() >= 4 {
434            let len = response_buf.len();
435            if response_buf[len - 4..] == [b'\r', b'\n', b'\r', b'\n'] {
436                break;
437            }
438        }
439
440        if response_buf.len() > 8192 {
441            return Err(io::Error::new(
442                io::ErrorKind::InvalidData,
443                "HTTP CONNECT: response too large",
444            ));
445        }
446    }
447
448    let response_str = String::from_utf8_lossy(&response_buf);
449    let status_line = response_str.lines().next().unwrap_or("");
450
451    // Parse status code from "HTTP/1.x NNN ..."
452    let parts: Vec<&str> = status_line.splitn(3, ' ').collect();
453    if parts.len() < 2 {
454        return Err(io::Error::new(
455            io::ErrorKind::InvalidData,
456            format!("HTTP CONNECT: malformed response: {}", status_line),
457        ));
458    }
459
460    let status_code: u16 = parts[1].parse().map_err(|_| {
461        io::Error::new(
462            io::ErrorKind::InvalidData,
463            format!("HTTP CONNECT: invalid status code: {}", parts[1]),
464        )
465    })?;
466
467    if status_code != 200 {
468        return Err(io::Error::new(
469            io::ErrorKind::ConnectionRefused,
470            format!("HTTP CONNECT: proxy returned status {}", status_code),
471        ));
472    }
473
474    Ok(())
475}
476
477/// Minimal base64 encoder (avoids external dep).
478fn base64_encode(data: &[u8]) -> String {
479    const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
480    let mut result = String::with_capacity(data.len().div_ceil(3) * 4);
481
482    for chunk in data.chunks(3) {
483        let b0 = chunk[0] as u32;
484        let b1 = if chunk.len() > 1 { chunk[1] as u32 } else { 0 };
485        let b2 = if chunk.len() > 2 { chunk[2] as u32 } else { 0 };
486        let triple = (b0 << 16) | (b1 << 8) | b2;
487
488        result.push(CHARS[((triple >> 18) & 0x3F) as usize] as char);
489        result.push(CHARS[((triple >> 12) & 0x3F) as usize] as char);
490
491        if chunk.len() > 1 {
492            result.push(CHARS[((triple >> 6) & 0x3F) as usize] as char);
493        } else {
494            result.push('=');
495        }
496
497        if chunk.len() > 2 {
498            result.push(CHARS[(triple & 0x3F) as usize] as char);
499        } else {
500            result.push('=');
501        }
502    }
503
504    result
505}
506
507// ── SOCKS5 UDP ASSOCIATE ─────────────────────────────────────────────
508
509/// Establish a SOCKS5 UDP relay.
510///
511/// Returns `(relay_addr, control_stream)`. The control stream must be
512/// kept alive for the duration of the relay — dropping it terminates
513/// the association.
514pub(crate) async fn socks5_udp_associate(
515    proxy: &ProxyConfig,
516    local_addr: SocketAddr,
517) -> io::Result<(SocketAddr, TcpStream)> {
518    let proxy_addr = format!("{}:{}", proxy.hostname, proxy.port);
519    let mut stream = TcpStream::connect(&proxy_addr).await?;
520
521    let want_auth = proxy.proxy_type == ProxyType::Socks5Password;
522    socks5_negotiate_method(&mut stream, want_auth).await?;
523
524    if want_auth {
525        let user = proxy.username.as_deref().unwrap_or("");
526        let pass = proxy.password.as_deref().unwrap_or("");
527        socks5_auth(&mut stream, user, pass).await?;
528    }
529
530    // UDP ASSOCIATE request (CMD=3)
531    let mut req = Vec::with_capacity(22);
532    req.push(5); // VER
533    req.push(3); // CMD = UDP ASSOCIATE
534    req.push(0); // RSV
535
536    if proxy.socks5_udp_send_local_ep {
537        encode_socks5_addr(&mut req, local_addr);
538    } else {
539        // Send zeros — let proxy figure it out
540        req.push(1); // ATYP = IPv4
541        req.extend_from_slice(&[0, 0, 0, 0]); // 0.0.0.0
542        req.extend_from_slice(&[0, 0]); // port 0
543    }
544
545    stream.write_all(&req).await?;
546
547    let relay_addr = read_socks5_response_addr(&mut stream).await?;
548
549    Ok((relay_addr, stream))
550}
551
552// ── SOCKS5 UDP header ────────────────────────────────────────────────
553
554/// Encode a SOCKS5 UDP relay header.
555///
556/// Format: RSV(2)=0, FRAG(1)=0, ATYP(1), DST.ADDR(var), DST.PORT(2)
557pub(crate) fn encode_socks5_udp_header(target: SocketAddr) -> Vec<u8> {
558    let mut hdr = Vec::with_capacity(10);
559    hdr.extend_from_slice(&[0, 0]); // RSV
560    hdr.push(0); // FRAG
561    encode_socks5_addr(&mut hdr, target);
562    hdr
563}
564
565/// Decode a SOCKS5 UDP relay header, returning (source_addr, data_offset).
566pub(crate) fn decode_socks5_udp_header(buf: &[u8]) -> io::Result<(SocketAddr, usize)> {
567    if buf.len() < 7 {
568        return Err(io::Error::new(
569            io::ErrorKind::InvalidData,
570            "SOCKS5 UDP header too short",
571        ));
572    }
573
574    // RSV(2) + FRAG(1) = first 3 bytes
575    let atyp = buf[3];
576
577    match atyp {
578        1 => {
579            // IPv4: 4 + 2 = 6 bytes after ATYP
580            if buf.len() < 10 {
581                return Err(io::Error::new(
582                    io::ErrorKind::InvalidData,
583                    "SOCKS5 UDP: IPv4 header truncated",
584                ));
585            }
586            let ip = Ipv4Addr::new(buf[4], buf[5], buf[6], buf[7]);
587            let port = u16::from_be_bytes([buf[8], buf[9]]);
588            Ok((SocketAddr::V4(SocketAddrV4::new(ip, port)), 10))
589        }
590        4 => {
591            // IPv6: 16 + 2 = 18 bytes after ATYP
592            if buf.len() < 22 {
593                return Err(io::Error::new(
594                    io::ErrorKind::InvalidData,
595                    "SOCKS5 UDP: IPv6 header truncated",
596                ));
597            }
598            let ip = Ipv6Addr::from(<[u8; 16]>::try_from(&buf[4..20]).unwrap());
599            let port = u16::from_be_bytes([buf[20], buf[21]]);
600            Ok((SocketAddr::V6(SocketAddrV6::new(ip, port, 0, 0)), 22))
601        }
602        _ => Err(io::Error::new(
603            io::ErrorKind::InvalidData,
604            format!("SOCKS5 UDP: unsupported ATYP {}", atyp),
605        )),
606    }
607}
608
609// ── ProxiedUdpSocket ─────────────────────────────────────────────────
610
611/// A UDP socket that routes through a SOCKS5 UDP relay.
612///
613/// Wraps each datagram with the SOCKS5 UDP header and strips it on receive.
614/// The `_control` stream must remain alive for the relay to function.
615pub(crate) struct ProxiedUdpSocket {
616    socket: UdpSocket,
617    relay_addr: SocketAddr,
618    _control: TcpStream,
619}
620
621impl ProxiedUdpSocket {
622    /// Create a new proxied UDP socket.
623    ///
624    /// `socket` is the local UDP socket, `relay_addr` is the proxy's UDP relay
625    /// endpoint, and `control` is the TCP stream that must stay alive.
626    pub fn new(socket: UdpSocket, relay_addr: SocketAddr, control: TcpStream) -> Self {
627        Self {
628            socket,
629            relay_addr,
630            _control: control,
631        }
632    }
633
634    /// Send data to `target` through the SOCKS5 UDP relay.
635    pub async fn send_to(&self, data: &[u8], target: SocketAddr) -> io::Result<usize> {
636        let header = encode_socks5_udp_header(target);
637        let mut packet = Vec::with_capacity(header.len() + data.len());
638        packet.extend_from_slice(&header);
639        packet.extend_from_slice(data);
640        self.socket.send_to(&packet, self.relay_addr).await
641    }
642
643    /// Receive data from the SOCKS5 UDP relay, stripping the relay header.
644    ///
645    /// Returns `(bytes_read, source_addr)` where `source_addr` is the original
646    /// sender (extracted from the SOCKS5 header, not the relay).
647    pub async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
648        let mut raw = vec![0u8; buf.len() + 22]; // max header overhead
649        let (n, _relay) = self.socket.recv_from(&mut raw).await?;
650
651        let (source, offset) = decode_socks5_udp_header(&raw[..n])?;
652        let data_len = n - offset;
653
654        if data_len > buf.len() {
655            return Err(io::Error::new(
656                io::ErrorKind::InvalidData,
657                "SOCKS5 UDP: data exceeds buffer",
658            ));
659        }
660
661        buf[..data_len].copy_from_slice(&raw[offset..n]);
662        Ok((data_len, source))
663    }
664}
665
666// ── Tests ────────────────────────────────────────────────────────────
667
668#[cfg(test)]
669mod tests {
670    use super::*;
671
672    // ── Test 11: SOCKS5 method negotiation + connect request/response ──
673
674    #[test]
675    fn socks5_encode_method_negotiation() {
676        // No auth
677        let no_auth = [5u8, 1, 0];
678        assert_eq!(no_auth[0], 5); // VER
679        assert_eq!(no_auth[1], 1); // NMETHODS
680        assert_eq!(no_auth[2], 0); // NO_AUTH
681
682        // With auth
683        let with_auth = [5u8, 2, 0, 2];
684        assert_eq!(with_auth[0], 5);
685        assert_eq!(with_auth[1], 2);
686        assert_eq!(with_auth[2], 0); // NO_AUTH
687        assert_eq!(with_auth[3], 2); // USERNAME/PASSWORD
688    }
689
690    #[test]
691    fn socks5_encode_connect_request_ipv4() {
692        let target: SocketAddr = "192.168.1.1:8080".parse().unwrap();
693        let mut buf = Vec::new();
694        buf.push(5); // VER
695        buf.push(1); // CMD = CONNECT
696        buf.push(0); // RSV
697        encode_socks5_addr(&mut buf, target);
698
699        assert_eq!(buf[0], 5); // VER
700        assert_eq!(buf[1], 1); // CMD
701        assert_eq!(buf[2], 0); // RSV
702        assert_eq!(buf[3], 1); // ATYP = IPv4
703        assert_eq!(&buf[4..8], &[192, 168, 1, 1]);
704        assert_eq!(&buf[8..10], &8080u16.to_be_bytes());
705    }
706
707    #[test]
708    fn socks5_encode_connect_request_ipv6() {
709        let target: SocketAddr = "[::1]:9999".parse().unwrap();
710        let mut buf = Vec::new();
711        buf.push(5);
712        buf.push(1);
713        buf.push(0);
714        encode_socks5_addr(&mut buf, target);
715
716        assert_eq!(buf[3], 4); // ATYP = IPv6
717        assert_eq!(buf.len(), 3 + 1 + 16 + 2); // header + atyp + ipv6 + port
718        // Last two bytes = port
719        assert_eq!(&buf[20..22], &9999u16.to_be_bytes());
720    }
721
722    #[test]
723    fn socks5_error_messages() {
724        assert!(socks5_error_message(1).contains("general failure"));
725        assert!(socks5_error_message(5).contains("connection refused"));
726        assert!(socks5_error_message(99).contains("unknown"));
727    }
728
729    // ── Test 12: SOCKS5 password auth subnegotiation ──
730
731    #[test]
732    fn socks5_auth_encode() {
733        // Verify RFC 1929 encoding: VER(1)=1, ULEN(1), UNAME, PLEN(1), PASSWD
734        let user = "alice";
735        let pass = "secret";
736
737        let mut req = Vec::new();
738        req.push(1); // subneg version
739        req.push(user.len() as u8);
740        req.extend_from_slice(user.as_bytes());
741        req.push(pass.len() as u8);
742        req.extend_from_slice(pass.as_bytes());
743
744        assert_eq!(req[0], 1); // VER
745        assert_eq!(req[1], 5); // ULEN
746        assert_eq!(&req[2..7], b"alice");
747        assert_eq!(req[7], 6); // PLEN
748        assert_eq!(&req[8..14], b"secret");
749    }
750
751    // ── Test 13: SOCKS4 handshake encode/decode ──
752
753    #[test]
754    fn socks4_encode_connect_request() {
755        let target: SocketAddr = "1.2.3.4:80".parse().unwrap();
756        let mut req = Vec::new();
757        req.push(4); // VN
758        req.push(1); // CD = CONNECT
759        req.extend_from_slice(&target.port().to_be_bytes());
760        req.extend_from_slice(&[1, 2, 3, 4]);
761        req.push(0); // NULL userid
762
763        assert_eq!(req[0], 4);
764        assert_eq!(req[1], 1);
765        assert_eq!(&req[2..4], &80u16.to_be_bytes());
766        assert_eq!(&req[4..8], &[1, 2, 3, 4]);
767        assert_eq!(req[8], 0);
768    }
769
770    #[test]
771    fn socks4_response_granted() {
772        let resp = [0u8, 90, 0, 0, 0, 0, 0, 0];
773        assert_eq!(resp[1], 90); // granted
774    }
775
776    #[test]
777    fn socks4_response_rejected() {
778        let resp = [0u8, 91, 0, 0, 0, 0, 0, 0];
779        assert_ne!(resp[1], 90); // rejected
780    }
781
782    // ── Test 14: HTTP CONNECT format/parse ──
783
784    #[test]
785    fn http_connect_request_no_auth() {
786        let target: SocketAddr = "93.184.216.34:443".parse().unwrap();
787        let host_port = format!("{}:{}", target.ip(), target.port());
788        let request = format!(
789            "CONNECT {} HTTP/1.1\r\nHost: {}\r\n\r\n",
790            host_port, host_port,
791        );
792
793        assert!(request.starts_with("CONNECT 93.184.216.34:443 HTTP/1.1\r\n"));
794        assert!(request.contains("Host: 93.184.216.34:443\r\n"));
795        assert!(request.ends_with("\r\n\r\n"));
796    }
797
798    #[test]
799    fn http_connect_request_with_auth() {
800        let encoded = base64_encode(b"user:pass");
801        assert_eq!(encoded, "dXNlcjpwYXNz");
802
803        let header = format!("Proxy-Authorization: Basic {}\r\n", encoded);
804        assert!(header.contains("dXNlcjpwYXNz"));
805    }
806
807    #[test]
808    fn http_connect_parse_200_response() {
809        let response = b"HTTP/1.1 200 Connection Established\r\n\r\n";
810        let response_str = String::from_utf8_lossy(response);
811        let status_line = response_str.lines().next().unwrap();
812        let parts: Vec<&str> = status_line.splitn(3, ' ').collect();
813        let status_code: u16 = parts[1].parse().unwrap();
814        assert_eq!(status_code, 200);
815    }
816
817    #[test]
818    fn http_connect_parse_407_response() {
819        let response = b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n";
820        let response_str = String::from_utf8_lossy(response);
821        let status_line = response_str.lines().next().unwrap();
822        let parts: Vec<&str> = status_line.splitn(3, ' ').collect();
823        let status_code: u16 = parts[1].parse().unwrap();
824        assert_eq!(status_code, 407);
825    }
826
827    // ── Test 15: SOCKS5 UDP header encode/decode ──
828
829    #[test]
830    fn socks5_udp_header_ipv4_roundtrip() {
831        let target: SocketAddr = "10.0.0.1:6881".parse().unwrap();
832        let header = encode_socks5_udp_header(target);
833
834        assert_eq!(header[0], 0); // RSV
835        assert_eq!(header[1], 0); // RSV
836        assert_eq!(header[2], 0); // FRAG
837        assert_eq!(header[3], 1); // ATYP = IPv4
838        assert_eq!(&header[4..8], &[10, 0, 0, 1]);
839        assert_eq!(&header[8..10], &6881u16.to_be_bytes());
840
841        let (decoded_addr, offset) = decode_socks5_udp_header(&header).unwrap();
842        assert_eq!(decoded_addr, target);
843        assert_eq!(offset, 10);
844    }
845
846    #[test]
847    fn socks5_udp_header_ipv6_roundtrip() {
848        let target: SocketAddr = "[2001:db8::1]:51413".parse().unwrap();
849        let header = encode_socks5_udp_header(target);
850
851        assert_eq!(header[3], 4); // ATYP = IPv6
852        assert_eq!(header.len(), 22); // 3 + 1 + 16 + 2
853
854        let (decoded_addr, offset) = decode_socks5_udp_header(&header).unwrap();
855        assert_eq!(decoded_addr, target);
856        assert_eq!(offset, 22);
857    }
858
859    #[test]
860    fn socks5_udp_header_too_short() {
861        let short = [0u8; 5];
862        assert!(decode_socks5_udp_header(&short).is_err());
863    }
864
865    // ── Test 16: ProxyConfig::to_url() ──
866
867    #[test]
868    fn proxy_config_to_url_none() {
869        let cfg = ProxyConfig::default();
870        assert_eq!(cfg.to_url(), "");
871    }
872
873    #[test]
874    fn proxy_config_to_url_socks5() {
875        let cfg = ProxyConfig {
876            proxy_type: ProxyType::Socks5,
877            hostname: "proxy.example.com".into(),
878            port: 1080,
879            ..Default::default()
880        };
881        assert_eq!(cfg.to_url(), "socks5://proxy.example.com:1080");
882    }
883
884    #[test]
885    fn proxy_config_to_url_socks5_password() {
886        let cfg = ProxyConfig {
887            proxy_type: ProxyType::Socks5Password,
888            hostname: "proxy.example.com".into(),
889            port: 1080,
890            username: Some("user".into()),
891            password: Some("pass".into()),
892            ..Default::default()
893        };
894        assert_eq!(cfg.to_url(), "socks5://user:pass@proxy.example.com:1080");
895    }
896
897    #[test]
898    fn proxy_config_to_url_http() {
899        let cfg = ProxyConfig {
900            proxy_type: ProxyType::Http,
901            hostname: "httpproxy.local".into(),
902            port: 8080,
903            ..Default::default()
904        };
905        assert_eq!(cfg.to_url(), "http://httpproxy.local:8080");
906    }
907
908    #[test]
909    fn proxy_config_to_url_http_password() {
910        let cfg = ProxyConfig {
911            proxy_type: ProxyType::HttpPassword,
912            hostname: "httpproxy.local".into(),
913            port: 3128,
914            username: Some("admin".into()),
915            password: Some("secret".into()),
916            ..Default::default()
917        };
918        assert_eq!(cfg.to_url(), "http://admin:secret@httpproxy.local:3128");
919    }
920
921    #[test]
922    fn proxy_config_to_url_socks4() {
923        let cfg = ProxyConfig {
924            proxy_type: ProxyType::Socks4,
925            hostname: "s4proxy".into(),
926            port: 1080,
927            ..Default::default()
928        };
929        assert_eq!(cfg.to_url(), "socks4://s4proxy:1080");
930    }
931
932    #[test]
933    fn base64_encode_basic() {
934        assert_eq!(base64_encode(b""), "");
935        assert_eq!(base64_encode(b"f"), "Zg==");
936        assert_eq!(base64_encode(b"fo"), "Zm8=");
937        assert_eq!(base64_encode(b"foo"), "Zm9v");
938        assert_eq!(base64_encode(b"foob"), "Zm9vYg==");
939        assert_eq!(base64_encode(b"fooba"), "Zm9vYmE=");
940        assert_eq!(base64_encode(b"foobar"), "Zm9vYmFy");
941    }
942
943    #[test]
944    fn proxy_config_default_flags() {
945        let cfg = ProxyConfig::default();
946        assert_eq!(cfg.proxy_type, ProxyType::None);
947        assert!(cfg.proxy_peer_connections);
948        assert!(cfg.proxy_tracker_connections);
949        assert!(cfg.proxy_hostnames);
950        assert!(!cfg.socks5_udp_send_local_ep);
951    }
952}