Skip to main content

microsandbox_network/proxy/
socks.rs

1//! SOCKS outbound proxy builders, credentials, and transport implementations.
2
3use std::fmt;
4use std::io;
5use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
6
7use microsandbox_types::SecretSource;
8use serde::{Deserialize, Serialize};
9use tokio::io::{AsyncReadExt, AsyncWriteExt};
10use tokio::net::TcpStream;
11use tokio_socks::tcp::Socks4Stream;
12use zeroize::Zeroizing;
13
14use super::types::{
15    OutboundProxy, OutboundProxyBuildError, OutboundProxyBuilder, OutboundProxyConfig,
16    OutboundProxyProtocol, ResolvedOutboundProxy,
17};
18use crate::dns::forwarder::{DnsForwarder, DnsForwarderHandle};
19
20//--------------------------------------------------------------------------------------------------
21// Types
22//--------------------------------------------------------------------------------------------------
23
24/// Environment-backed username/password credentials for a SOCKS5 proxy.
25#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
26pub struct Socks5Credentials {
27    username: String,
28    password: SecretSource,
29}
30
31/// Resolved SOCKS5 credentials held only by the network runtime.
32#[doc(hidden)]
33#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
34pub struct ResolvedSocks5Credentials {
35    username: String,
36    password: Zeroizing<String>,
37}
38
39/// Builds a SOCKS4 outbound proxy.
40#[derive(Debug, Clone)]
41pub struct Socks4ProxyBuilder {
42    address: String,
43    user_id: Option<String>,
44}
45
46/// Builds a SOCKS5 outbound proxy.
47#[derive(Debug, Clone)]
48pub struct Socks5ProxyBuilder {
49    address: String,
50    credentials: Option<Socks5Credentials>,
51}
52
53/// Active SOCKS5 UDP association.
54pub(crate) struct Socks5UdpAssociation {
55    _control: TcpStream,
56    socket: tokio::net::UdpSocket,
57    dns_forwarder: Option<DnsForwarderHandle>,
58}
59
60/// SOCKS5 wire protocol operations shared by TCP and UDP proxying.
61struct Socks5Protocol;
62
63/// Address returned by a SOCKS5 command reply.
64enum Socks5ReplyAddress {
65    Socket(SocketAddr),
66    Domain { name: String, port: u16 },
67}
68
69//--------------------------------------------------------------------------------------------------
70// Methods
71//--------------------------------------------------------------------------------------------------
72
73impl ResolvedOutboundProxy {
74    /// Builds the complete proxy used by the network runtime.
75    #[doc(hidden)]
76    pub fn build(
77        configured: Option<&OutboundProxy>,
78        resolved: Option<ResolvedSocks5Credentials>,
79    ) -> Result<Option<Self>, OutboundProxyBuildError> {
80        let Some(configured) = configured else {
81            if resolved.is_some() {
82                return Err(OutboundProxyBuildError::InvalidSocks5Credentials {
83                    reason: "launch credentials require a configured SOCKS5 proxy",
84                });
85            }
86            return Ok(None);
87        };
88
89        configured.validate()?;
90        match configured {
91            OutboundProxy::Socks4 { address, user_id } => {
92                if resolved.is_some() {
93                    return Err(OutboundProxyBuildError::InvalidSocks5Credentials {
94                        reason: "launch credentials require a configured SOCKS5 proxy",
95                    });
96                }
97                Ok(Some(Self::Socks4 {
98                    address: *address,
99                    user_id: user_id.clone(),
100                }))
101            }
102            OutboundProxy::Socks5 {
103                address,
104                credentials,
105            } => {
106                let credentials = match (credentials, resolved) {
107                    (None, None) => None,
108                    (None, Some(_)) => {
109                        return Err(OutboundProxyBuildError::InvalidSocks5Credentials {
110                            reason: "launch credentials require configured SOCKS5 credentials",
111                        });
112                    }
113                    (Some(_), None) => {
114                        return Err(OutboundProxyBuildError::InvalidSocks5Credentials {
115                            reason: "configured SOCKS5 credentials were not resolved at launch",
116                        });
117                    }
118                    (Some(configured), Some(resolved)) => {
119                        if resolved.username != configured.username {
120                            return Err(OutboundProxyBuildError::InvalidSocks5Credentials {
121                                reason: "launch username does not match the durable configuration",
122                            });
123                        }
124                        resolved.validate()?;
125                        Some(resolved)
126                    }
127                };
128                Ok(Some(Self::Socks5 {
129                    address: *address,
130                    credentials,
131                }))
132            }
133        }
134    }
135
136    /// Connects to `destination` through this outbound proxy.
137    pub(crate) async fn connect(&self, destination: SocketAddr) -> io::Result<TcpStream> {
138        match self {
139            Self::Socks4 { address, user_id } => match user_id {
140                Some(user_id) => {
141                    Socks4Stream::connect_with_userid(*address, destination, user_id).await
142                }
143                None => Socks4Stream::connect(*address, destination).await,
144            }
145            .map(|stream| stream.into_inner())
146            .map_err(io::Error::other),
147            Self::Socks5 {
148                address,
149                credentials,
150            } => {
151                let mut stream = TcpStream::connect(*address).await?;
152                Socks5Protocol::negotiate(&mut stream, credentials.as_ref()).await?;
153                // The bound address is informational for CONNECT. Parse it to
154                // validate the reply, but do not resolve proxy-supplied domains.
155                let _ = Socks5Protocol::command(&mut stream, 0x01, destination).await?;
156                Ok(stream)
157            }
158        }
159    }
160
161    /// Opens a SOCKS5 UDP association for relaying datagrams.
162    pub(crate) async fn associate_udp(
163        &self,
164        dns_forwarder: Option<DnsForwarderHandle>,
165    ) -> io::Result<Socks5UdpAssociation> {
166        let Self::Socks5 {
167            address,
168            credentials,
169        } = self
170        else {
171            return Err(io::Error::new(
172                io::ErrorKind::Unsupported,
173                "SOCKS4 does not support UDP relay",
174            ));
175        };
176
177        let mut control = TcpStream::connect(*address).await?;
178        Socks5Protocol::negotiate(&mut control, credentials.as_ref()).await?;
179
180        let control_local = control.local_addr()?;
181        // The UDP endpoint is not known until the proxy returns its relay.
182        // RFC 1928 requires an all-zero endpoint in that case.
183        let request_address = match control_local.ip() {
184            IpAddr::V4(_) => SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0),
185            IpAddr::V6(_) => SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0),
186        };
187        let relays = match Socks5Protocol::command(&mut control, 0x03, request_address).await? {
188            Socks5ReplyAddress::Socket(relay) => vec![relay],
189            Socks5ReplyAddress::Domain { name, port } => {
190                Socks5Protocol::resolve_domain(dns_forwarder.as_ref(), &name, port).await?
191            }
192        };
193        Socks5UdpAssociation::connect(control, relays, dns_forwarder).await
194    }
195}
196
197impl OutboundProxy {
198    fn validate(&self) -> Result<(), OutboundProxyBuildError> {
199        match self {
200            Self::Socks4 { user_id, .. } => Self::validate_socks4_user_id(user_id.as_deref()),
201            Self::Socks5 { credentials, .. } => credentials
202                .as_ref()
203                .map_or(Ok(()), Socks5Credentials::validate),
204        }
205    }
206
207    fn validate_socks4_user_id(user_id: Option<&str>) -> Result<(), OutboundProxyBuildError> {
208        let Some(user_id) = user_id else {
209            return Ok(());
210        };
211        let reason = if user_id.is_empty() {
212            "must not be empty"
213        } else if user_id.len() > 255 {
214            "must be at most 255 bytes"
215        } else if user_id.contains('\0') {
216            "must not contain a null byte"
217        } else {
218            return Ok(());
219        };
220
221        Err(OutboundProxyBuildError::InvalidSocks4UserId { reason })
222    }
223}
224
225impl ResolvedSocks5Credentials {
226    /// Creates credentials for the private launch contract.
227    #[doc(hidden)]
228    pub fn new(username: impl Into<String>, password: impl Into<String>) -> Self {
229        Self {
230            username: username.into(),
231            password: Zeroizing::new(password.into()),
232        }
233    }
234
235    /// Validates the RFC 1929 one-octet credential lengths.
236    fn validate(&self) -> Result<(), OutboundProxyBuildError> {
237        let reason = if self.username.is_empty() {
238            "username must not be empty"
239        } else if self.username.len() > u8::MAX as usize {
240            "username must be at most 255 bytes"
241        } else if self.password.is_empty() {
242            "password must not be empty"
243        } else if self.password.len() > u8::MAX as usize {
244            "password must be at most 255 bytes"
245        } else {
246            return Ok(());
247        };
248
249        Err(OutboundProxyBuildError::InvalidSocks5Credentials { reason })
250    }
251}
252
253impl Socks5Credentials {
254    pub(crate) fn username(&self) -> &str {
255        &self.username
256    }
257
258    pub(crate) fn password_source(&self) -> &SecretSource {
259        &self.password
260    }
261
262    /// Validates the durable SOCKS5 credential configuration.
263    fn validate(&self) -> Result<(), OutboundProxyBuildError> {
264        if self.username.is_empty() {
265            return Err(OutboundProxyBuildError::InvalidSocks5Credentials {
266                reason: "username must not be empty",
267            });
268        }
269        if self.username.len() > u8::MAX as usize {
270            return Err(OutboundProxyBuildError::InvalidSocks5Credentials {
271                reason: "username must be at most 255 bytes",
272            });
273        }
274        match &self.password {
275            SecretSource::Env { var } if var.is_empty() => {
276                Err(OutboundProxyBuildError::InvalidSocks5Credentials {
277                    reason: "password environment variable must not be empty",
278                })
279            }
280            SecretSource::Env { .. } => Ok(()),
281            SecretSource::Store { .. } => Err(OutboundProxyBuildError::InvalidSocks5Credentials {
282                reason: "store-backed password sources are not supported yet",
283            }),
284        }
285    }
286}
287
288impl Socks5UdpAssociation {
289    /// Connects a UDP socket to the first usable relay address.
290    async fn connect(
291        control: TcpStream,
292        relays: Vec<SocketAddr>,
293        dns_forwarder: Option<DnsForwarderHandle>,
294    ) -> io::Result<Self> {
295        let peer_ip = control.peer_addr()?.ip();
296        let mut last_error = None;
297
298        for relay in relays {
299            let relay = if relay.ip().is_unspecified() {
300                SocketAddr::new(peer_ip, relay.port())
301            } else {
302                relay
303            };
304            let bind_address = match relay.ip() {
305                IpAddr::V4(_) => SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0),
306                IpAddr::V6(_) => SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0),
307            };
308            let socket = match tokio::net::UdpSocket::bind(bind_address).await {
309                Ok(socket) => socket,
310                Err(error) => {
311                    last_error = Some(error);
312                    continue;
313                }
314            };
315            match socket.connect(relay).await {
316                Ok(()) => {
317                    return Ok(Self {
318                        _control: control,
319                        socket,
320                        dns_forwarder,
321                    });
322                }
323                Err(error) => last_error = Some(error),
324            }
325        }
326
327        Err(last_error.unwrap_or_else(|| {
328            io::Error::new(
329                io::ErrorKind::AddrNotAvailable,
330                "SOCKS5 proxy returned no usable UDP relay address",
331            )
332        }))
333    }
334
335    /// Sends one payload to `destination` through the UDP association.
336    pub(crate) async fn send_to(
337        &self,
338        payload: &[u8],
339        destination: SocketAddr,
340    ) -> io::Result<usize> {
341        let mut datagram =
342            Vec::with_capacity(Socks5Protocol::address_len(destination) + 3 + payload.len());
343        datagram.extend_from_slice(&[0x00, 0x00, 0x00]);
344        Socks5Protocol::encode_address(&mut datagram, destination);
345        datagram.extend_from_slice(payload);
346        self.socket.send(&datagram).await.map(|_| payload.len())
347    }
348
349    /// Receives one payload and returns the remote endpoints encoded by the proxy.
350    pub(crate) async fn recv_from(
351        &self,
352        buffer: &mut [u8],
353    ) -> io::Result<(usize, Vec<SocketAddr>)> {
354        let received = self.socket.recv(buffer).await?;
355        let (header_len, source) = Socks5Protocol::decode_udp_header(&buffer[..received])?;
356        let sources = match source {
357            Socks5ReplyAddress::Socket(source) => vec![source],
358            Socks5ReplyAddress::Domain { name, port } => {
359                Socks5Protocol::resolve_domain(self.dns_forwarder.as_ref(), &name, port).await?
360            }
361        };
362        let payload_len = received - header_len;
363        buffer.copy_within(header_len..received, 0);
364        Ok((payload_len, sources))
365    }
366}
367
368impl fmt::Debug for ResolvedSocks5Credentials {
369    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
370        f.debug_struct("ResolvedSocks5Credentials")
371            .field("username", &self.username)
372            .field("password", &"[REDACTED]")
373            .finish()
374    }
375}
376
377impl OutboundProxyBuilder {
378    /// Creates a protocol selector.
379    pub fn new() -> Self {
380        Self
381    }
382
383    /// Starts building a SOCKS4 outbound proxy.
384    pub fn socks4(self, address: impl Into<String>) -> Socks4ProxyBuilder {
385        Socks4ProxyBuilder {
386            address: address.into(),
387            user_id: None,
388        }
389    }
390
391    /// Starts building a SOCKS5 outbound proxy.
392    pub fn socks5(self, address: impl Into<String>) -> Socks5ProxyBuilder {
393        Socks5ProxyBuilder {
394            address: address.into(),
395            credentials: None,
396        }
397    }
398}
399
400impl Socks4ProxyBuilder {
401    /// Sets the optional user ID sent during the SOCKS4 handshake.
402    pub fn user_id(mut self, user_id: impl Into<String>) -> Self {
403        self.user_id = Some(user_id.into());
404        self
405    }
406}
407
408impl Socks5ProxyBuilder {
409    /// Sets username authentication and a host-side password source.
410    pub fn credentials(mut self, username: impl Into<String>, password: SecretSource) -> Self {
411        self.credentials = Some(Socks5Credentials {
412            username: username.into(),
413            password,
414        });
415        self
416    }
417}
418
419//--------------------------------------------------------------------------------------------------
420// Trait Implementations
421//--------------------------------------------------------------------------------------------------
422
423impl OutboundProxyConfig for Socks4ProxyBuilder {
424    fn build(self) -> Result<OutboundProxy, OutboundProxyBuildError> {
425        let address =
426            self.address
427                .parse()
428                .map_err(|source| OutboundProxyBuildError::InvalidAddress {
429                    protocol: OutboundProxyProtocol::Socks4,
430                    address: self.address,
431                    source,
432                })?;
433
434        OutboundProxy::validate_socks4_user_id(self.user_id.as_deref())?;
435        Ok(OutboundProxy::Socks4 {
436            address,
437            user_id: self.user_id,
438        })
439    }
440}
441
442impl OutboundProxyConfig for Socks5ProxyBuilder {
443    fn build(self) -> Result<OutboundProxy, OutboundProxyBuildError> {
444        let address =
445            self.address
446                .parse()
447                .map_err(|source| OutboundProxyBuildError::InvalidAddress {
448                    protocol: OutboundProxyProtocol::Socks5,
449                    address: self.address,
450                    source,
451                })?;
452        if let Some(credentials) = &self.credentials {
453            credentials.validate()?;
454        }
455        Ok(OutboundProxy::Socks5 {
456            address,
457            credentials: self.credentials,
458        })
459    }
460}
461
462impl OutboundProxyConfig for OutboundProxy {
463    fn build(self) -> Result<OutboundProxy, OutboundProxyBuildError> {
464        self.validate()?;
465        Ok(self)
466    }
467}
468
469impl Socks5Protocol {
470    /// Resolves a domain-form SOCKS5 endpoint through the internal DNS path.
471    async fn resolve_domain(
472        dns_forwarder: Option<&DnsForwarderHandle>,
473        name: &str,
474        port: u16,
475    ) -> io::Result<Vec<SocketAddr>> {
476        let dns_forwarder = dns_forwarder.ok_or_else(|| {
477            io::Error::other("DNS forwarder is unavailable for SOCKS5 UDP endpoint resolution")
478        })?;
479        let forwarder = DnsForwarder::wait(dns_forwarder.clone())
480            .await
481            .ok_or_else(|| {
482                io::Error::other("DNS forwarder is unavailable for SOCKS5 UDP endpoint resolution")
483            })?;
484        Ok(forwarder
485            .resolve_proxy_domain(name)
486            .await?
487            .into_iter()
488            .map(|address| SocketAddr::new(address, port))
489            .collect())
490    }
491
492    /// Negotiates a SOCKS5 authentication method and performs username/password
493    /// authentication when selected by the proxy.
494    async fn negotiate(
495        stream: &mut TcpStream,
496        credentials: Option<&ResolvedSocks5Credentials>,
497    ) -> io::Result<()> {
498        if let Some(credentials) = credentials {
499            credentials.validate().map_err(io::Error::other)?;
500        }
501
502        match credentials {
503            Some(_) => stream.write_all(&[0x05, 0x02, 0x00, 0x02]).await?,
504            None => stream.write_all(&[0x05, 0x01, 0x00]).await?,
505        }
506
507        let mut selection = [0u8; 2];
508        stream.read_exact(&mut selection).await?;
509        if selection[0] != 0x05 {
510            return Err(Self::invalid_response("invalid method-selection version"));
511        }
512
513        match selection[1] {
514            0x00 => Ok(()),
515            0x02 => {
516                let credentials = credentials.ok_or_else(|| {
517                    io::Error::new(
518                        io::ErrorKind::PermissionDenied,
519                        "SOCKS5 proxy requires username/password authentication",
520                    )
521                })?;
522                let username = credentials.username.as_bytes();
523                let password = credentials.password.as_bytes();
524                let mut request = Vec::with_capacity(3 + username.len() + password.len());
525                request.extend_from_slice(&[0x01, username.len() as u8]);
526                request.extend_from_slice(username);
527                request.push(password.len() as u8);
528                request.extend_from_slice(password);
529                stream.write_all(&request).await?;
530
531                let mut response = [0u8; 2];
532                stream.read_exact(&mut response).await?;
533                if response[0] != 0x01 {
534                    return Err(Self::invalid_response(
535                        "invalid username/password response version",
536                    ));
537                }
538                if response[1] != 0x00 {
539                    return Err(io::Error::new(
540                        io::ErrorKind::PermissionDenied,
541                        "SOCKS5 username/password authentication failed",
542                    ));
543                }
544                Ok(())
545            }
546            0xff => Err(io::Error::new(
547                io::ErrorKind::PermissionDenied,
548                "SOCKS5 proxy rejected all offered authentication methods",
549            )),
550            method => Err(Self::invalid_response(format!(
551                "SOCKS5 proxy selected unsupported authentication method {method:#04x}"
552            ))),
553        }
554    }
555
556    /// Sends a SOCKS5 command and returns the bound address from its reply.
557    async fn command(
558        stream: &mut TcpStream,
559        command: u8,
560        destination: SocketAddr,
561    ) -> io::Result<Socks5ReplyAddress> {
562        let mut request = Vec::with_capacity(3 + Self::address_len(destination));
563        request.extend_from_slice(&[0x05, command, 0x00]);
564        Self::encode_address(&mut request, destination);
565        stream.write_all(&request).await?;
566
567        let mut header = [0u8; 4];
568        stream.read_exact(&mut header).await?;
569        if header[0] != 0x05 || header[2] != 0x00 {
570            return Err(Self::invalid_response("invalid SOCKS5 command reply"));
571        }
572        if header[1] != 0x00 {
573            return Err(io::Error::other(format!(
574                "SOCKS5 proxy command failed: {}",
575                Self::reply_message(header[1])
576            )));
577        }
578
579        Self::read_address(stream, header[3]).await
580    }
581
582    /// Reads a SOCKS5 address whose address-type byte was already consumed.
583    async fn read_address(
584        stream: &mut TcpStream,
585        address_type: u8,
586    ) -> io::Result<Socks5ReplyAddress> {
587        let ip = match address_type {
588            0x01 => {
589                let mut octets = [0u8; 4];
590                stream.read_exact(&mut octets).await?;
591                IpAddr::V4(Ipv4Addr::from(octets))
592            }
593            0x04 => {
594                let mut octets = [0u8; 16];
595                stream.read_exact(&mut octets).await?;
596                IpAddr::V6(Ipv6Addr::from(octets))
597            }
598            0x03 => {
599                let length = stream.read_u8().await? as usize;
600                let mut domain = vec![0u8; length];
601                stream.read_exact(&mut domain).await?;
602                let domain = String::from_utf8(domain).map_err(|_| {
603                    Self::invalid_response("SOCKS5 reply contains a non-UTF-8 domain")
604                })?;
605                let mut port = [0u8; 2];
606                stream.read_exact(&mut port).await?;
607                return Ok(Socks5ReplyAddress::Domain {
608                    name: domain,
609                    port: u16::from_be_bytes(port),
610                });
611            }
612            _ => return Err(Self::invalid_response("unsupported SOCKS5 address type")),
613        };
614
615        let mut port = [0u8; 2];
616        stream.read_exact(&mut port).await?;
617        Ok(Socks5ReplyAddress::Socket(SocketAddr::new(
618            ip,
619            u16::from_be_bytes(port),
620        )))
621    }
622
623    /// Encodes a socket address in SOCKS5 address form.
624    fn encode_address(output: &mut Vec<u8>, address: SocketAddr) {
625        match address {
626            SocketAddr::V4(address) => {
627                output.push(0x01);
628                output.extend_from_slice(&address.ip().octets());
629                output.extend_from_slice(&address.port().to_be_bytes());
630            }
631            SocketAddr::V6(address) => {
632                output.push(0x04);
633                output.extend_from_slice(&address.ip().octets());
634                output.extend_from_slice(&address.port().to_be_bytes());
635            }
636        }
637    }
638
639    /// Returns the encoded length of a SOCKS5 socket address.
640    fn address_len(address: SocketAddr) -> usize {
641        match address {
642            SocketAddr::V4(_) => 7,
643            SocketAddr::V6(_) => 19,
644        }
645    }
646
647    /// Decodes a SOCKS5 UDP request header and returns its payload offset and endpoint.
648    fn decode_udp_header(datagram: &[u8]) -> io::Result<(usize, Socks5ReplyAddress)> {
649        if datagram.len() < 4 || datagram[..2] != [0x00, 0x00] {
650            return Err(Self::invalid_response("invalid SOCKS5 UDP header"));
651        }
652        if datagram[2] != 0x00 {
653            return Err(Self::invalid_response(
654                "fragmented SOCKS5 UDP datagrams are not supported",
655            ));
656        }
657
658        let (endpoint, port_offset) = match datagram[3] {
659            0x01 if datagram.len() >= 10 => (
660                Socks5ReplyAddress::Socket(SocketAddr::new(
661                    IpAddr::V4(Ipv4Addr::new(
662                        datagram[4],
663                        datagram[5],
664                        datagram[6],
665                        datagram[7],
666                    )),
667                    u16::from_be_bytes([datagram[8], datagram[9]]),
668                )),
669                8,
670            ),
671            0x04 if datagram.len() >= 22 => {
672                let mut octets = [0u8; 16];
673                octets.copy_from_slice(&datagram[4..20]);
674                (
675                    Socks5ReplyAddress::Socket(SocketAddr::new(
676                        IpAddr::V6(Ipv6Addr::from(octets)),
677                        u16::from_be_bytes([datagram[20], datagram[21]]),
678                    )),
679                    20,
680                )
681            }
682            0x03 if datagram.len() >= 7 => {
683                let length = datagram[4] as usize;
684                let port_offset = 5 + length;
685                if length == 0 || datagram.len() < port_offset + 2 {
686                    return Err(Self::invalid_response("invalid SOCKS5 UDP domain address"));
687                }
688                let name = std::str::from_utf8(&datagram[5..port_offset])
689                    .map_err(|_| Self::invalid_response("non-UTF-8 SOCKS5 UDP domain address"))?
690                    .to_owned();
691                (
692                    Socks5ReplyAddress::Domain {
693                        name,
694                        port: u16::from_be_bytes([
695                            datagram[port_offset],
696                            datagram[port_offset + 1],
697                        ]),
698                    },
699                    port_offset,
700                )
701            }
702            _ => return Err(Self::invalid_response("invalid SOCKS5 UDP address")),
703        };
704        Ok((port_offset + 2, endpoint))
705    }
706
707    /// Converts a SOCKS5 reply code into a stable diagnostic.
708    fn reply_message(reply: u8) -> &'static str {
709        match reply {
710            0x01 => "general server failure",
711            0x02 => "connection not allowed by ruleset",
712            0x03 => "network unreachable",
713            0x04 => "host unreachable",
714            0x05 => "connection refused",
715            0x06 => "TTL expired",
716            0x07 => "command not supported",
717            0x08 => "address type not supported",
718            _ => "unknown error",
719        }
720    }
721
722    /// Builds an invalid-data error for malformed SOCKS5 responses.
723    fn invalid_response(message: impl Into<String>) -> io::Error {
724        io::Error::new(io::ErrorKind::InvalidData, message.into())
725    }
726}
727
728//--------------------------------------------------------------------------------------------------
729// Tests
730//--------------------------------------------------------------------------------------------------
731
732#[cfg(test)]
733mod tests {
734    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
735    use std::sync::Arc;
736
737    use hickory_net::proto::op::{Message, MessageType, OpCode};
738    use hickory_net::proto::rr::rdata::{A, AAAA};
739    use hickory_net::proto::rr::{RData, Record, RecordType};
740    use hickory_net::proto::serialize::binary::{BinDecodable, BinEncodable};
741    use microsandbox_types::SecretSource;
742    use tokio::io::{AsyncReadExt, AsyncWriteExt};
743    use tokio::net::{TcpListener, UdpSocket};
744    use tokio::sync::watch;
745
746    use super::{
747        OutboundProxy, OutboundProxyBuildError, OutboundProxyBuilder, OutboundProxyConfig,
748        OutboundProxyProtocol, ResolvedOutboundProxy, ResolvedSocks5Credentials,
749    };
750    use crate::dns::forwarder::DnsForwarder;
751    use crate::netstack::poll::GatewayIps;
752    use crate::netstack::shared::SharedState;
753
754    async fn responding_dns(relay_ipv6: Ipv6Addr, source_ipv4: Ipv4Addr) -> SocketAddr {
755        let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
756        let address = socket.local_addr().unwrap();
757        tokio::spawn(async move {
758            let mut buffer = [0u8; 4096];
759            loop {
760                let Ok((length, source)) = socket.recv_from(&mut buffer).await else {
761                    continue;
762                };
763                let Ok(query) = Message::from_bytes(&buffer[..length]) else {
764                    continue;
765                };
766                let mut response =
767                    Message::new(query.metadata.id, MessageType::Response, OpCode::Query);
768                response.metadata.recursion_desired = query.metadata.recursion_desired;
769                response.metadata.recursion_available = true;
770                if let Some(question) = query.queries.first() {
771                    response.add_query(question.clone());
772                    let domain = question.name().to_string();
773                    let answer = match (domain.trim_end_matches('.'), question.query_type()) {
774                        ("relay.example.com", RecordType::AAAA) => {
775                            Some(RData::AAAA(AAAA::from(relay_ipv6)))
776                        }
777                        ("source.example.com", RecordType::A) => {
778                            Some(RData::A(A::from(source_ipv4)))
779                        }
780                        _ => None,
781                    };
782                    if let Some(answer) = answer {
783                        response.add_answer(Record::from_rdata(
784                            question.name().clone(),
785                            60,
786                            answer,
787                        ));
788                    }
789                }
790                if let Ok(bytes) = response.to_bytes() {
791                    let _ = socket.send_to(&bytes, source).await;
792                }
793            }
794        });
795        address
796    }
797
798    #[test]
799    fn builder_creates_socks4_proxy_with_optional_user_id() {
800        let address = "127.0.0.1:1080".parse().unwrap();
801        let without_user_id = OutboundProxyBuilder::new()
802            .socks4("127.0.0.1:1080")
803            .build()
804            .unwrap();
805        let with_user_id = OutboundProxyBuilder::new()
806            .socks4("127.0.0.1:1080")
807            .user_id("sandbox")
808            .build()
809            .unwrap();
810
811        assert_eq!(
812            without_user_id,
813            OutboundProxy::Socks4 {
814                address,
815                user_id: None,
816            }
817        );
818        assert_eq!(
819            with_user_id,
820            OutboundProxy::Socks4 {
821                address,
822                user_id: Some("sandbox".to_string()),
823            }
824        );
825    }
826
827    #[test]
828    fn builder_creates_socks5_proxy() {
829        let proxy = OutboundProxyBuilder::new()
830            .socks5("127.0.0.1:1080")
831            .build()
832            .unwrap();
833
834        assert_eq!(
835            proxy,
836            OutboundProxy::Socks5 {
837                address: "127.0.0.1:1080".parse().unwrap(),
838                credentials: None,
839            }
840        );
841    }
842
843    #[test]
844    fn builder_creates_socks5_proxy_with_password_source() {
845        let proxy = OutboundProxyBuilder::new()
846            .socks5("127.0.0.1:1080")
847            .credentials(
848                "sandbox",
849                SecretSource::Env {
850                    var: "SOCKS5_PASSWORD".into(),
851                },
852            )
853            .build()
854            .unwrap();
855
856        let debug = format!("{proxy:?}");
857        assert!(debug.contains("sandbox"));
858        assert!(debug.contains("SOCKS5_PASSWORD"));
859
860        let json = serde_json::to_value(&proxy).unwrap();
861        assert_eq!(json["credentials"]["username"], "sandbox");
862        assert_eq!(json["credentials"]["password"]["kind"], "env");
863        assert_eq!(json["credentials"]["password"]["var"], "SOCKS5_PASSWORD");
864        assert!(json["credentials"].get("value").is_none());
865    }
866
867    #[test]
868    fn builder_rejects_invalid_socks5_credentials() {
869        for (username, password_env) in [
870            (String::new(), "SOCKS5_PASSWORD".to_string()),
871            ("username".to_string(), String::new()),
872            ("u".repeat(256), "SOCKS5_PASSWORD".to_string()),
873        ] {
874            assert!(
875                OutboundProxyBuilder::new()
876                    .socks5("127.0.0.1:1080")
877                    .credentials(username, SecretSource::Env { var: password_env })
878                    .build()
879                    .is_err()
880            );
881        }
882
883        assert!(
884            OutboundProxyBuilder::new()
885                .socks5("127.0.0.1:1080")
886                .credentials(
887                    "username",
888                    SecretSource::Store {
889                        reference: "production/socks5-password".into(),
890                    },
891                )
892                .build()
893                .is_err()
894        );
895    }
896
897    #[test]
898    fn resolved_proxy_build_validates_resolved_socks5_password() {
899        let proxy = OutboundProxyBuilder::new()
900            .socks5("127.0.0.1:1080")
901            .credentials(
902                "sandbox",
903                SecretSource::Env {
904                    var: "SOCKS5_PASSWORD".into(),
905                },
906            )
907            .build()
908            .unwrap();
909        assert!(
910            ResolvedOutboundProxy::build(
911                Some(&proxy),
912                Some(ResolvedSocks5Credentials::new("sandbox", "")),
913            )
914            .is_err()
915        );
916        assert!(
917            ResolvedOutboundProxy::build(
918                Some(&proxy),
919                Some(ResolvedSocks5Credentials::new("sandbox", "p".repeat(256),)),
920            )
921            .is_err()
922        );
923        ResolvedOutboundProxy::build(
924            Some(&proxy),
925            Some(ResolvedSocks5Credentials::new("sandbox", "password")),
926        )
927        .unwrap();
928    }
929
930    #[test]
931    fn resolved_proxy_build_requires_matching_resolved_credentials() {
932        let authenticated = OutboundProxyBuilder::new()
933            .socks5("127.0.0.1:1080")
934            .credentials("sandbox", SecretSource::env("SOCKS5_PASSWORD"))
935            .build()
936            .unwrap();
937        let unauthenticated = OutboundProxyBuilder::new()
938            .socks5("127.0.0.1:1080")
939            .build()
940            .unwrap();
941        let resolved = || ResolvedSocks5Credentials::new("sandbox", "password");
942
943        assert!(ResolvedOutboundProxy::build(Some(&authenticated), None).is_err());
944        assert!(ResolvedOutboundProxy::build(Some(&unauthenticated), Some(resolved())).is_err());
945        assert!(ResolvedOutboundProxy::build(None, Some(resolved())).is_err());
946        assert!(
947            ResolvedOutboundProxy::build(
948                Some(&authenticated),
949                Some(ResolvedSocks5Credentials::new("different-user", "password")),
950            )
951            .is_err()
952        );
953    }
954
955    #[test]
956    fn uri_parses_and_formats_for_cli() {
957        let socks4: OutboundProxy = "socks4://127.0.0.1:1080".parse().unwrap();
958        let socks5: OutboundProxy = "socks5://127.0.0.1:1080".parse().unwrap();
959
960        assert_eq!(
961            socks4,
962            OutboundProxy::Socks4 {
963                address: "127.0.0.1:1080".parse().unwrap(),
964                user_id: None,
965            }
966        );
967        assert_eq!(socks4.to_string(), "socks4://127.0.0.1:1080");
968        assert_eq!(
969            socks5,
970            OutboundProxy::Socks5 {
971                address: "127.0.0.1:1080".parse().unwrap(),
972                credentials: None,
973            }
974        );
975        assert_eq!(socks5.to_string(), "socks5://127.0.0.1:1080");
976    }
977
978    #[test]
979    fn uri_rejects_unsupported_forms() {
980        for raw in [
981            "127.0.0.1:1080",
982            "http://127.0.0.1:1080",
983            "socks4://user@127.0.0.1:1080",
984            "socks5://user@127.0.0.1:1080",
985            "socks5://127.0.0.1:1080/path",
986            "socks5://127.0.0.1:1080?option=value",
987            "socks5://127.0.0.1:1080#fragment",
988        ] {
989            assert!(raw.parse::<OutboundProxy>().is_err(), "accepted {raw:?}");
990        }
991    }
992
993    #[test]
994    fn builder_rejects_invalid_socks4_user_ids() {
995        for user_id in [String::new(), "a\0b".to_string(), "a".repeat(256)] {
996            assert!(
997                OutboundProxyBuilder::new()
998                    .socks4("127.0.0.1:1080")
999                    .user_id(user_id)
1000                    .build()
1001                    .is_err()
1002            );
1003        }
1004    }
1005
1006    #[test]
1007    fn outbound_proxy_rejects_invalid_socks4_user_ids() {
1008        for user_id in [String::new(), "a\0b".to_string(), "a".repeat(256)] {
1009            let proxy = OutboundProxy::Socks4 {
1010                address: "127.0.0.1:1080".parse().unwrap(),
1011                user_id: Some(user_id),
1012            };
1013
1014            assert!(proxy.build().is_err());
1015        }
1016    }
1017
1018    #[test]
1019    fn invalid_address_error_uses_typed_protocol() {
1020        let error = OutboundProxyBuilder::new()
1021            .socks5("not-an-address")
1022            .build()
1023            .unwrap_err();
1024
1025        assert!(matches!(
1026            error,
1027            OutboundProxyBuildError::InvalidAddress {
1028                protocol: OutboundProxyProtocol::Socks5,
1029                ..
1030            }
1031        ));
1032    }
1033
1034    #[tokio::test]
1035    async fn connects_through_socks4_proxy_with_user_id() {
1036        let target: SocketAddr = "93.184.216.34:443".parse().unwrap();
1037        let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1038        let proxy_addr = proxy_listener.local_addr().unwrap();
1039        let proxy_task = tokio::spawn(async move {
1040            let (mut client, _) = proxy_listener.accept().await.unwrap();
1041
1042            let mut request = [0u8; 16];
1043            client.read_exact(&mut request).await.unwrap();
1044            assert_eq!(request[0], 0x04, "SOCKS version");
1045            assert_eq!(request[1], 0x01, "CONNECT command");
1046            assert_eq!(u16::from_be_bytes([request[2], request[3]]), 443);
1047            assert_eq!(&request[4..8], &[93, 184, 216, 34]);
1048            assert_eq!(&request[8..], b"sandbox\0");
1049
1050            client
1051                .write_all(&[0x00, 0x5a, 0x01, 0xbb, 93, 184, 216, 34])
1052                .await
1053                .unwrap();
1054
1055            let mut buf = [0u8; 5];
1056            client.read_exact(&mut buf).await.unwrap();
1057            client.write_all(&buf).await.unwrap();
1058        });
1059
1060        let mut stream = ResolvedOutboundProxy::Socks4 {
1061            address: proxy_addr,
1062            user_id: Some("sandbox".to_string()),
1063        }
1064        .connect(target)
1065        .await
1066        .unwrap();
1067        stream.write_all(b"hello").await.unwrap();
1068        let mut echoed = [0u8; 5];
1069        stream.read_exact(&mut echoed).await.unwrap();
1070        assert_eq!(&echoed, b"hello");
1071
1072        proxy_task.await.unwrap();
1073    }
1074
1075    #[tokio::test]
1076    async fn connects_through_socks5_proxy() {
1077        let target: SocketAddr = "93.184.216.34:443".parse().unwrap();
1078        let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1079        let proxy_addr = proxy_listener.local_addr().unwrap();
1080        let proxy_task = tokio::spawn(async move {
1081            let (mut client, _) = proxy_listener.accept().await.unwrap();
1082
1083            let mut greeting = [0u8; 3];
1084            client.read_exact(&mut greeting).await.unwrap();
1085            assert_eq!(greeting, [0x05, 0x01, 0x00]);
1086            client.write_all(&[0x05, 0x00]).await.unwrap();
1087
1088            let mut request = [0u8; 10];
1089            client.read_exact(&mut request).await.unwrap();
1090            assert_eq!(request[0], 0x05, "SOCKS version");
1091            assert_eq!(request[1], 0x01, "CONNECT command");
1092            assert_eq!(request[3], 0x01, "IPv4 address type");
1093            assert_eq!(&request[4..8], &[93, 184, 216, 34]);
1094            assert_eq!(u16::from_be_bytes([request[8], request[9]]), 443);
1095
1096            client
1097                .write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
1098                .await
1099                .unwrap();
1100
1101            let mut buf = [0u8; 5];
1102            client.read_exact(&mut buf).await.unwrap();
1103            client.write_all(&buf).await.unwrap();
1104        });
1105
1106        let mut stream = ResolvedOutboundProxy::Socks5 {
1107            address: proxy_addr,
1108            credentials: None,
1109        }
1110        .connect(target)
1111        .await
1112        .unwrap();
1113        stream.write_all(b"hello").await.unwrap();
1114        let mut echoed = [0u8; 5];
1115        stream.read_exact(&mut echoed).await.unwrap();
1116        assert_eq!(&echoed, b"hello");
1117
1118        proxy_task.await.unwrap();
1119    }
1120
1121    #[tokio::test]
1122    async fn connect_does_not_resolve_domain_from_socks5_reply() {
1123        let target: SocketAddr = "93.184.216.34:443".parse().unwrap();
1124        let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1125        let proxy_addr = proxy_listener.local_addr().unwrap();
1126        let proxy_task = tokio::spawn(async move {
1127            let (mut client, _) = proxy_listener.accept().await.unwrap();
1128
1129            let mut greeting = [0u8; 3];
1130            client.read_exact(&mut greeting).await.unwrap();
1131            client.write_all(&[0x05, 0x00]).await.unwrap();
1132
1133            let mut request = [0u8; 10];
1134            client.read_exact(&mut request).await.unwrap();
1135            let domain = b"does-not-resolve.invalid";
1136            let mut response = vec![0x05, 0x00, 0x00, 0x03, domain.len() as u8];
1137            response.extend_from_slice(domain);
1138            response.extend_from_slice(&443u16.to_be_bytes());
1139            client.write_all(&response).await.unwrap();
1140
1141            let mut buf = [0u8; 5];
1142            client.read_exact(&mut buf).await.unwrap();
1143            client.write_all(&buf).await.unwrap();
1144        });
1145
1146        let mut stream = ResolvedOutboundProxy::Socks5 {
1147            address: proxy_addr,
1148            credentials: None,
1149        }
1150        .connect(target)
1151        .await
1152        .unwrap();
1153        stream.write_all(b"hello").await.unwrap();
1154        let mut echoed = [0u8; 5];
1155        stream.read_exact(&mut echoed).await.unwrap();
1156        assert_eq!(&echoed, b"hello");
1157
1158        proxy_task.await.unwrap();
1159    }
1160
1161    #[tokio::test]
1162    async fn connects_through_authenticated_socks5_proxy() {
1163        let target: SocketAddr = "93.184.216.34:443".parse().unwrap();
1164        let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1165        let proxy_addr = proxy_listener.local_addr().unwrap();
1166        let proxy_task = tokio::spawn(async move {
1167            let (mut client, _) = proxy_listener.accept().await.unwrap();
1168
1169            let mut greeting = [0u8; 4];
1170            client.read_exact(&mut greeting).await.unwrap();
1171            assert_eq!(greeting, [0x05, 0x02, 0x00, 0x02]);
1172            client.write_all(&[0x05, 0x02]).await.unwrap();
1173
1174            let mut auth = [0u8; 18];
1175            client.read_exact(&mut auth).await.unwrap();
1176            assert_eq!(&auth, b"\x01\x07sandbox\x08password");
1177            client.write_all(&[0x01, 0x00]).await.unwrap();
1178
1179            let mut request = [0u8; 10];
1180            client.read_exact(&mut request).await.unwrap();
1181            assert_eq!(request[1], 0x01, "CONNECT command");
1182            client
1183                .write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
1184                .await
1185                .unwrap();
1186        });
1187
1188        let configured = OutboundProxyBuilder::new()
1189            .socks5(proxy_addr.to_string())
1190            .credentials(
1191                "sandbox",
1192                SecretSource::Env {
1193                    var: "SOCKS5_PASSWORD".into(),
1194                },
1195            )
1196            .build()
1197            .unwrap();
1198        let proxy = ResolvedOutboundProxy::build(
1199            Some(&configured),
1200            Some(ResolvedSocks5Credentials::new("sandbox", "password")),
1201        )
1202        .unwrap()
1203        .unwrap();
1204        proxy.connect(target).await.unwrap();
1205
1206        proxy_task.await.unwrap();
1207    }
1208
1209    #[tokio::test]
1210    async fn associates_and_relays_socks5_udp() {
1211        let target: SocketAddr = "93.184.216.34:5353".parse().unwrap();
1212        let relay = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1213        let relay_addr = relay.local_addr().unwrap();
1214        let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1215        let proxy_addr = proxy_listener.local_addr().unwrap();
1216
1217        let proxy_task = tokio::spawn(async move {
1218            let (mut control, _) = proxy_listener.accept().await.unwrap();
1219            let mut greeting = [0u8; 3];
1220            control.read_exact(&mut greeting).await.unwrap();
1221            assert_eq!(greeting, [0x05, 0x01, 0x00]);
1222            control.write_all(&[0x05, 0x00]).await.unwrap();
1223
1224            let mut request = [0u8; 10];
1225            control.read_exact(&mut request).await.unwrap();
1226            assert_eq!(request[1], 0x03, "UDP ASSOCIATE command");
1227
1228            let mut response = vec![0x05, 0x00, 0x00, 0x01];
1229            let SocketAddr::V4(relay_addr) = relay_addr else {
1230                unreachable!()
1231            };
1232            response.extend_from_slice(&relay_addr.ip().octets());
1233            response.extend_from_slice(&relay_addr.port().to_be_bytes());
1234            control.write_all(&response).await.unwrap();
1235
1236            let mut datagram = [0u8; 64];
1237            let (received, client) = relay.recv_from(&mut datagram).await.unwrap();
1238            assert_eq!(&datagram[..10], &[0, 0, 0, 1, 93, 184, 216, 34, 0x14, 0xe9]);
1239            assert_eq!(&datagram[10..received], b"hello");
1240            relay.send_to(&datagram[..received], client).await.unwrap();
1241        });
1242
1243        let configured = OutboundProxyBuilder::new()
1244            .socks5(proxy_addr.to_string())
1245            .build()
1246            .unwrap();
1247        let proxy = ResolvedOutboundProxy::build(Some(&configured), None)
1248            .unwrap()
1249            .unwrap();
1250        let association = proxy.associate_udp(None).await.unwrap();
1251        association.send_to(b"hello", target).await.unwrap();
1252        let mut response = [0u8; 64];
1253        let (received, sources) = association.recv_from(&mut response).await.unwrap();
1254        assert_eq!(sources, vec![target]);
1255        assert_eq!(&response[..received], b"hello");
1256
1257        proxy_task.await.unwrap();
1258    }
1259
1260    #[tokio::test]
1261    async fn udp_association_supports_domain_relay_in_another_address_family() {
1262        let target: SocketAddr = "93.184.216.34:5353".parse().unwrap();
1263        let relay = UdpSocket::bind("[::1]:0").await.unwrap();
1264        let relay_addr = relay.local_addr().unwrap();
1265        let dns_upstream =
1266            responding_dns(Ipv6Addr::LOCALHOST, Ipv4Addr::new(93, 184, 216, 34)).await;
1267        let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1268        let proxy_addr = proxy_listener.local_addr().unwrap();
1269        let proxy_task = tokio::spawn(async move {
1270            let (mut control, _) = proxy_listener.accept().await.unwrap();
1271
1272            let mut greeting = [0u8; 3];
1273            control.read_exact(&mut greeting).await.unwrap();
1274            control.write_all(&[0x05, 0x00]).await.unwrap();
1275
1276            let mut request = [0u8; 10];
1277            control.read_exact(&mut request).await.unwrap();
1278            assert_eq!(request[1], 0x03, "UDP ASSOCIATE command");
1279
1280            let domain = b"relay.example.com";
1281            let mut response = vec![0x05, 0x00, 0x00, 0x03, domain.len() as u8];
1282            response.extend_from_slice(domain);
1283            response.extend_from_slice(&relay_addr.port().to_be_bytes());
1284            control.write_all(&response).await.unwrap();
1285
1286            let mut datagram = [0u8; 64];
1287            let (_, client) = relay.recv_from(&mut datagram).await.unwrap();
1288            let source_domain = b"source.example.com";
1289            let mut response = vec![0x00, 0x00, 0x00, 0x03, source_domain.len() as u8];
1290            response.extend_from_slice(source_domain);
1291            response.extend_from_slice(&target.port().to_be_bytes());
1292            response.extend_from_slice(b"hello");
1293            relay.send_to(&response, client).await.unwrap();
1294        });
1295
1296        let proxy = ResolvedOutboundProxy::Socks5 {
1297            address: proxy_addr,
1298            credentials: None,
1299        };
1300        let forwarder = DnsForwarder::for_proxy_test(
1301            Arc::new(SharedState::new(4)),
1302            GatewayIps {
1303                ipv4: Some("127.0.0.1".parse().unwrap()),
1304                ipv6: None,
1305            },
1306            Some(dns_upstream),
1307        )
1308        .await;
1309        let (_dns_tx, dns_forwarder) = watch::channel(Some(forwarder));
1310        let association = proxy.associate_udp(Some(dns_forwarder)).await.unwrap();
1311        association.send_to(b"hello", target).await.unwrap();
1312        let mut response = [0u8; 64];
1313        let (received, sources) = association.recv_from(&mut response).await.unwrap();
1314        assert_eq!(sources, vec![target]);
1315        assert_eq!(&response[..received], b"hello");
1316
1317        proxy_task.await.unwrap();
1318    }
1319}