better_socks/
lib.rs

1//! Asynchronous SOCKS5 client for Tokio.
2//!
3//! Fork of [`tokio-socks`](https://github.com/sticnarf/tokio-socks) with UDP
4//! `ASSOCIATE` support via [`udp::Socks5UdpFramed`].
5//!
6//! Username/password authentication (RFC 1929) is sent in cleartext to the
7//! proxy. There is no TLS-to-proxy transport in this crate.
8
9use either::Either;
10use futures_util::{
11    future,
12    stream::{self, Once, Stream},
13};
14use std::{
15    borrow::Cow,
16    fmt,
17    future::Future,
18    io,
19    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs},
20    pin::Pin,
21    task::{Context, Poll},
22    vec,
23};
24use tokio::net::lookup_host;
25
26pub use error::Error;
27
28pub type Result<T> = std::result::Result<T, Error>;
29
30/// A trait for objects which can be converted or resolved to one or more
31/// `SocketAddr` values, which are going to be connected as the the proxy
32/// server.
33///
34/// This trait is similar to `std::net::ToSocketAddrs` but allows asynchronous
35/// name resolution.
36pub trait ToProxyAddrs {
37    type Output: Stream<Item = Result<SocketAddr>> + Unpin;
38
39    fn to_proxy_addrs(&self) -> Self::Output;
40}
41
42macro_rules! trivial_impl_to_proxy_addrs {
43    ($t: ty) => {
44        impl ToProxyAddrs for $t {
45            type Output = Once<future::Ready<Result<SocketAddr>>>;
46
47            fn to_proxy_addrs(&self) -> Self::Output {
48                stream::once(future::ready(Ok(SocketAddr::from(*self))))
49            }
50        }
51    };
52}
53
54trivial_impl_to_proxy_addrs!(SocketAddr);
55trivial_impl_to_proxy_addrs!((IpAddr, u16));
56trivial_impl_to_proxy_addrs!((Ipv4Addr, u16));
57trivial_impl_to_proxy_addrs!((Ipv6Addr, u16));
58trivial_impl_to_proxy_addrs!(SocketAddrV4);
59trivial_impl_to_proxy_addrs!(SocketAddrV6);
60
61impl ToProxyAddrs for &[SocketAddr] {
62    type Output = ProxyAddrsStream;
63
64    fn to_proxy_addrs(&self) -> Self::Output {
65        ProxyAddrsStream::from_addrs(self.to_vec())
66    }
67}
68
69impl ToProxyAddrs for str {
70    type Output = ProxyAddrsStream;
71
72    fn to_proxy_addrs(&self) -> Self::Output {
73        let host = self.to_owned();
74        ProxyAddrsStream::lookup(async move { Ok(lookup_host(host).await?.collect()) })
75    }
76}
77
78impl ToProxyAddrs for (&str, u16) {
79    type Output = ProxyAddrsStream;
80
81    fn to_proxy_addrs(&self) -> Self::Output {
82        let host = self.0.to_owned();
83        let port = self.1;
84        ProxyAddrsStream::lookup(async move { Ok(lookup_host((host, port)).await?.collect()) })
85    }
86}
87
88impl<T: ToProxyAddrs + ?Sized> ToProxyAddrs for &T {
89    type Output = T::Output;
90
91    fn to_proxy_addrs(&self) -> Self::Output {
92        (**self).to_proxy_addrs()
93    }
94}
95
96enum ProxyAddrsInner {
97    Iter(vec::IntoIter<SocketAddr>),
98    Lookup(Pin<Box<dyn Future<Output = io::Result<Vec<SocketAddr>>> + Send>>),
99    Done,
100}
101
102/// Stream of resolved proxy server addresses.
103pub struct ProxyAddrsStream {
104    inner: ProxyAddrsInner,
105}
106
107impl ProxyAddrsStream {
108    fn from_addrs(addrs: Vec<SocketAddr>) -> Self {
109        Self {
110            inner: ProxyAddrsInner::Iter(addrs.into_iter()),
111        }
112    }
113
114    fn lookup<F>(fut: F) -> Self
115    where
116        F: Future<Output = io::Result<Vec<SocketAddr>>> + Send + 'static,
117    {
118        Self {
119            inner: ProxyAddrsInner::Lookup(Box::pin(fut)),
120        }
121    }
122}
123
124impl Stream for ProxyAddrsStream {
125    type Item = Result<SocketAddr>;
126
127    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
128        let this = self.get_mut();
129        loop {
130            match &mut this.inner {
131                ProxyAddrsInner::Iter(iter) => return Poll::Ready(iter.next().map(Ok)),
132                ProxyAddrsInner::Lookup(fut) => match fut.as_mut().poll(cx) {
133                    Poll::Ready(Ok(addrs)) => {
134                        this.inner = ProxyAddrsInner::Iter(addrs.into_iter());
135                    },
136                    Poll::Ready(Err(e)) => {
137                        this.inner = ProxyAddrsInner::Done;
138                        return Poll::Ready(Some(Err(e.into())));
139                    },
140                    Poll::Pending => return Poll::Pending,
141                },
142                ProxyAddrsInner::Done => return Poll::Ready(None),
143            }
144        }
145    }
146}
147
148/// A SOCKS connection target.
149#[derive(Clone, Debug, PartialEq, Eq)]
150pub enum TargetAddr<'a> {
151    /// Connect to an IP address.
152    Ip(SocketAddr),
153
154    /// Connect to a fully-qualified domain name.
155    ///
156    /// The domain name will be passed along to the proxy server and DNS lookup
157    /// will happen there.
158    Domain(Cow<'a, str>, u16),
159}
160
161impl TryFrom<TargetAddr<'_>> for SocketAddr {
162    type Error = Error;
163
164    fn try_from(item: TargetAddr<'_>) -> Result<Self> {
165        match item {
166            TargetAddr::Ip(addr) => Ok(addr),
167            TargetAddr::Domain(_, _) => Err(Error::InvalidTargetAddress(
168                "cannot convert a domain target to SocketAddr without DNS resolution",
169            )),
170        }
171    }
172}
173
174impl TargetAddr<'_> {
175    /// Creates owned `TargetAddr` by cloning. It is usually used to eliminate
176    /// the lifetime bound.
177    pub fn to_owned(&self) -> TargetAddr<'static> {
178        match self {
179            TargetAddr::Ip(addr) => TargetAddr::Ip(*addr),
180            TargetAddr::Domain(domain, port) => TargetAddr::Domain(Cow::Owned(domain.clone().into_owned()), *port),
181        }
182    }
183
184    pub fn port(&self) -> u16 {
185        match self {
186            TargetAddr::Ip(addr) => addr.port(),
187            TargetAddr::Domain(_, port) => *port,
188        }
189    }
190
191    pub fn set_port(&mut self, port: u16) {
192        match self {
193            TargetAddr::Ip(addr) => addr.set_port(port),
194            TargetAddr::Domain(_, p) => *p = port,
195        }
196    }
197}
198
199impl ToSocketAddrs for TargetAddr<'_> {
200    type Iter = Either<std::option::IntoIter<SocketAddr>, std::vec::IntoIter<SocketAddr>>;
201
202    /// Resolves domain names with the local resolver. This bypasses SOCKS
203    /// remote DNS and can leak queries; prefer passing the domain to the proxy.
204    fn to_socket_addrs(&self) -> io::Result<Self::Iter> {
205        Ok(match self {
206            TargetAddr::Ip(addr) => Either::Left(addr.to_socket_addrs()?),
207            TargetAddr::Domain(domain, port) => Either::Right((&**domain, *port).to_socket_addrs()?),
208        })
209    }
210}
211
212/// A trait for objects that can be converted to `TargetAddr`.
213pub trait IntoTargetAddr<'a> {
214    /// Converts the value of self to a `TargetAddr`.
215    fn into_target_addr(self) -> Result<TargetAddr<'a>>;
216}
217
218macro_rules! trivial_impl_into_target_addr {
219    ($t: ty) => {
220        impl<'a> IntoTargetAddr<'a> for $t {
221            fn into_target_addr(self) -> Result<TargetAddr<'a>> {
222                Ok(TargetAddr::Ip(SocketAddr::from(self)))
223            }
224        }
225    };
226}
227
228trivial_impl_into_target_addr!(SocketAddr);
229trivial_impl_into_target_addr!((IpAddr, u16));
230trivial_impl_into_target_addr!((Ipv4Addr, u16));
231trivial_impl_into_target_addr!((Ipv6Addr, u16));
232trivial_impl_into_target_addr!(SocketAddrV4);
233trivial_impl_into_target_addr!(SocketAddrV6);
234
235impl<'a> IntoTargetAddr<'a> for TargetAddr<'a> {
236    fn into_target_addr(self) -> Result<TargetAddr<'a>> {
237        Ok(self)
238    }
239}
240
241impl<'a> IntoTargetAddr<'a> for (&'a str, u16) {
242    fn into_target_addr(self) -> Result<TargetAddr<'a>> {
243        // Try IP address first
244        if let Ok(addr) = self.0.parse::<IpAddr>() {
245            return (addr, self.1).into_target_addr();
246        }
247
248        // Treat as domain name
249        if self.0.len() > 255 {
250            return Err(Error::InvalidTargetAddress("overlong domain"));
251        }
252        // TODO: Should we validate the domain format here?
253
254        Ok(TargetAddr::Domain(self.0.into(), self.1))
255    }
256}
257
258impl<'a> IntoTargetAddr<'a> for &'a str {
259    fn into_target_addr(self) -> Result<TargetAddr<'a>> {
260        // Try IP address first
261        if let Ok(addr) = self.parse::<SocketAddr>() {
262            return addr.into_target_addr();
263        }
264
265        // Unbracketed IPv6 host:port (e.g. `::1:80`) is not supported; use `[::1]:80`.
266        let mut parts_iter = self.rsplitn(2, ':');
267        let port: u16 = parts_iter
268            .next()
269            .and_then(|port_str| port_str.parse().ok())
270            .ok_or(Error::InvalidTargetAddress("invalid address format"))?;
271        let domain = parts_iter
272            .next()
273            .ok_or(Error::InvalidTargetAddress("invalid address format"))?;
274        if domain.len() > 255 {
275            return Err(Error::InvalidTargetAddress("overlong domain"));
276        }
277        Ok(TargetAddr::Domain(domain.into(), port))
278    }
279}
280
281impl IntoTargetAddr<'static> for String {
282    fn into_target_addr(mut self) -> Result<TargetAddr<'static>> {
283        // Try IP address first
284        if let Ok(addr) = self.parse::<SocketAddr>() {
285            return addr.into_target_addr();
286        }
287
288        let mut parts_iter = self.rsplitn(2, ':');
289        let port: u16 = parts_iter
290            .next()
291            .and_then(|port_str| port_str.parse().ok())
292            .ok_or(Error::InvalidTargetAddress("invalid address format"))?;
293        let domain_len = parts_iter
294            .next()
295            .ok_or(Error::InvalidTargetAddress("invalid address format"))?
296            .len();
297        if domain_len > 255 {
298            return Err(Error::InvalidTargetAddress("overlong domain"));
299        }
300        self.truncate(domain_len);
301        Ok(TargetAddr::Domain(self.into(), port))
302    }
303}
304
305impl IntoTargetAddr<'static> for (String, u16) {
306    fn into_target_addr(self) -> Result<TargetAddr<'static>> {
307        let addr = (self.0.as_str(), self.1).into_target_addr()?;
308        if let TargetAddr::Ip(addr) = addr {
309            Ok(TargetAddr::Ip(addr))
310        } else {
311            Ok(TargetAddr::Domain(self.0.into(), self.1))
312        }
313    }
314}
315
316impl<'a, T> IntoTargetAddr<'a> for &'a T
317where T: IntoTargetAddr<'a> + Copy
318{
319    fn into_target_addr(self) -> Result<TargetAddr<'a>> {
320        (*self).into_target_addr()
321    }
322}
323
324impl fmt::Display for TargetAddr<'_> {
325    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
326        match self {
327            Self::Ip(addr) => write!(f, "{addr}"),
328            Self::Domain(domain, port) => write!(f, "{domain}:{port}"),
329        }
330    }
331}
332
333/// Authentication methods
334#[derive(Debug)]
335enum Authentication<'a> {
336    Password { username: &'a str, password: &'a str },
337    None,
338}
339
340mod error;
341pub mod tcp;
342pub mod udp;
343
344#[cfg(test)]
345mod tests {
346    use super::*;
347    use futures_executor::block_on;
348    use futures_util::StreamExt;
349
350    fn to_proxy_addrs<T: ToProxyAddrs>(t: T) -> Result<Vec<SocketAddr>> {
351        Ok(block_on(t.to_proxy_addrs().map(Result::unwrap).collect()))
352    }
353
354    #[test]
355    fn converts_socket_addr_to_proxy_addrs() -> Result<()> {
356        let addr = SocketAddr::from(([1, 1, 1, 1], 443));
357        let res = to_proxy_addrs(addr)?;
358        assert_eq!(&res[..], &[addr]);
359        Ok(())
360    }
361
362    #[test]
363    fn converts_socket_addr_ref_to_proxy_addrs() -> Result<()> {
364        let addr = SocketAddr::from(([1, 1, 1, 1], 443));
365        #[allow(clippy::needless_borrows_for_generic_args)]
366        let res = to_proxy_addrs(&addr)?;
367        assert_eq!(&res[..], &[addr]);
368        Ok(())
369    }
370
371    #[test]
372    fn converts_socket_addrs_to_proxy_addrs() -> Result<()> {
373        let addrs = [
374            SocketAddr::from(([1, 1, 1, 1], 443)),
375            SocketAddr::from(([8, 8, 8, 8], 53)),
376        ];
377        let res = to_proxy_addrs(&addrs[..])?;
378        assert_eq!(&res[..], &addrs);
379        Ok(())
380    }
381
382    fn into_target_addr<'a, T>(t: T) -> Result<TargetAddr<'a>>
383    where T: IntoTargetAddr<'a> {
384        t.into_target_addr()
385    }
386
387    #[test]
388    fn converts_socket_addr_to_target_addr() -> Result<()> {
389        let addr = SocketAddr::from(([1, 1, 1, 1], 443));
390        let res = into_target_addr(addr)?;
391        assert_eq!(TargetAddr::Ip(addr), res);
392        let addr2 = SocketAddr::try_from(res)?;
393        assert_eq!(addr, addr2);
394        Ok(())
395    }
396
397    #[test]
398    fn domain_target_does_not_convert_to_socket_addr() {
399        let res = into_target_addr("www.example.com:80").unwrap();
400        assert!(SocketAddr::try_from(res).is_err());
401    }
402
403    #[test]
404    fn converts_socket_addr_ref_to_target_addr() -> Result<()> {
405        let addr = SocketAddr::from(([1, 1, 1, 1], 443));
406        #[allow(clippy::needless_borrows_for_generic_args)]
407        let res = into_target_addr(&addr)?;
408        assert_eq!(TargetAddr::Ip(addr), res);
409        Ok(())
410    }
411
412    #[test]
413    fn converts_socket_addr_str_to_target_addr() -> Result<()> {
414        let addr = SocketAddr::from(([1, 1, 1, 1], 443));
415        let ip_str = format!("{}", addr);
416        let res = into_target_addr(ip_str.as_str())?;
417        assert_eq!(TargetAddr::Ip(addr), res);
418        Ok(())
419    }
420
421    #[test]
422    fn converts_ip_str_and_port_target_addr() -> Result<()> {
423        let addr = SocketAddr::from(([1, 1, 1, 1], 443));
424        let ip_str = format!("{}", addr.ip());
425        let res = into_target_addr((ip_str.as_str(), addr.port()))?;
426        assert_eq!(TargetAddr::Ip(addr), res);
427        Ok(())
428    }
429
430    #[test]
431    fn converts_domain_to_target_addr() -> Result<()> {
432        let domain = "www.example.com:80";
433        let res = into_target_addr(domain)?;
434        assert_eq!(TargetAddr::Domain(Cow::Borrowed("www.example.com"), 80), res);
435
436        let res = into_target_addr(domain.to_owned())?;
437        assert_eq!(TargetAddr::Domain(Cow::Owned("www.example.com".to_owned()), 80), res);
438        Ok(())
439    }
440
441    #[test]
442    fn converts_domain_and_port_to_target_addr() -> Result<()> {
443        let domain = "www.example.com";
444        let res = into_target_addr((domain, 80))?;
445        assert_eq!(TargetAddr::Domain(Cow::Borrowed("www.example.com"), 80), res);
446        Ok(())
447    }
448
449    #[test]
450    fn overlong_domain_to_target_addr_should_fail() {
451        let domain = format!("www.{:a<1$}.com:80", 'a', 300);
452        assert!(into_target_addr(domain.as_str()).is_err());
453        let domain = format!("www.{:a<1$}.com", 'a', 300);
454        assert!(into_target_addr((domain.as_str(), 80)).is_err());
455    }
456
457    #[test]
458    fn addr_with_invalid_port_to_target_addr_should_fail() {
459        let addr = "[ffff::1]:65536";
460        assert!(into_target_addr(addr).is_err());
461        let addr = "www.example.com:65536";
462        assert!(into_target_addr(addr).is_err());
463    }
464
465    #[test]
466    fn target_addr_to_string() {
467        let domain = "www.example.com:80";
468        let res = into_target_addr(domain).unwrap().to_string();
469        assert_eq!(domain, res.as_str());
470
471        let domain = "0.0.0.0:80";
472        let res = into_target_addr(domain).unwrap().to_string();
473        assert_eq!(domain, res.as_str());
474    }
475
476    #[test]
477    fn target_addr_set_port() {
478        let domain = "www.example.com:80";
479        let mut res = into_target_addr(domain).unwrap();
480        let port = 16u16;
481        res.set_port(16u16);
482        assert_eq!(port, res.port());
483
484        let domain = "0.0.0.0:80";
485        let mut res = into_target_addr(domain).unwrap();
486        res.set_port(16u16);
487        assert_eq!(port, res.port());
488    }
489
490    #[test]
491    fn proxy_addrs_stream_invalid_host_does_not_panic_on_repoll() {
492        let rt = tokio::runtime::Builder::new_current_thread()
493            .enable_all()
494            .build()
495            .unwrap();
496        rt.block_on(async {
497            let mut stream = "256.256.256.256:80".to_proxy_addrs();
498            let first = stream.next().await;
499            assert!(first.unwrap().is_err());
500            assert!(stream.next().await.is_none());
501            assert!(stream.next().await.is_none());
502        });
503    }
504}
better_socks/lib.rs

better_socks/
lib.rs
