Skip to main content

microsandbox_network/proxy/
socks.rs

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