Skip to main content

irontide_session/
proxy.rs

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