smb2_client/
socks.rs

1//! Optional SOCKS5 egress. Every TCP dial in the stack (SMB here, plus RPC/LDAP/KDC/WinRM in the
2//! crates that depend on this one) goes through [`dial`], which routes to a SOCKS5 proxy when one
3//! has been registered with [`set_proxy`] — the pivot support real engagements need. Hand-rolled
4//! (RFC 1928 CONNECT + RFC 1929 user/pass), consistent with the from-scratch stack.
5
6use std::io::{Error, ErrorKind, Result};
7use std::sync::OnceLock;
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9use tokio::net::TcpStream;
10
11/// A SOCKS5 proxy: `host:port` plus optional username/password auth.
12#[derive(Clone, Debug)]
13pub struct Socks5 {
14    pub proxy: String,
15    pub auth: Option<(String, String)>,
16}
17
18impl Socks5 {
19    /// Parse `[user:pass@]host:port`.
20    pub fn parse(s: &str) -> Option<Socks5> {
21        let (auth, hostport) = match s.rsplit_once('@') {
22            Some((creds, hp)) => {
23                let (u, p) = creds.split_once(':')?;
24                (Some((u.to_string(), p.to_string())), hp.to_string())
25            }
26            None => (None, s.to_string()),
27        };
28        if !hostport.contains(':') {
29            return None;
30        }
31        Some(Socks5 {
32            proxy: hostport,
33            auth,
34        })
35    }
36}
37
38static PROXY: OnceLock<Option<Socks5>> = OnceLock::new();
39
40/// Register the process-wide SOCKS5 proxy (call once, at startup). `None` means direct connections.
41pub fn set_proxy(cfg: Option<Socks5>) {
42    let _ = PROXY.set(cfg);
43}
44
45/// The registered proxy, if any.
46pub fn proxy() -> Option<&'static Socks5> {
47    PROXY.get().and_then(|o| o.as_ref())
48}
49
50fn err(msg: &str) -> Error {
51    Error::new(ErrorKind::Other, msg)
52}
53
54/// Split `host` into (host, port), defaulting the port when absent. IPv4/hostnames only.
55fn host_port(host: &str, default_port: u16) -> (String, u16) {
56    if let Some((h, p)) = host.rsplit_once(':') {
57        if let Ok(port) = p.parse::<u16>() {
58            return (h.to_string(), port);
59        }
60    }
61    (host.to_string(), default_port)
62}
63
64/// Dial `host` (with optional `:port`, else `default_port`), routing through the registered SOCKS5
65/// proxy if one is set — otherwise a direct TCP connection. The **hostname is sent to the proxy**
66/// (ATYP=domain) so internal names resolve on the pivot side.
67pub async fn dial(host: &str, default_port: u16) -> Result<TcpStream> {
68    let (h, p) = host_port(host, default_port);
69    match proxy() {
70        Some(cfg) => socks5_connect(cfg, &h, p).await,
71        None => TcpStream::connect((h.as_str(), p)).await,
72    }
73}
74
75async fn socks5_connect(cfg: &Socks5, dst_host: &str, dst_port: u16) -> Result<TcpStream> {
76    let mut s = TcpStream::connect(&cfg.proxy).await?;
77
78    // Greeting: offer no-auth (and user/pass if we have creds).
79    if cfg.auth.is_some() {
80        s.write_all(&[0x05, 0x02, 0x00, 0x02]).await?;
81    } else {
82        s.write_all(&[0x05, 0x01, 0x00]).await?;
83    }
84    let mut sel = [0u8; 2];
85    s.read_exact(&mut sel).await?;
86    if sel[0] != 0x05 {
87        return Err(err("SOCKS: bad version in method reply"));
88    }
89    match sel[1] {
90        0x00 => {}
91        0x02 => {
92            let (u, pw) = cfg.auth.as_ref().ok_or_else(|| err("SOCKS: proxy demands auth but none given"))?;
93            if u.len() > 255 || pw.len() > 255 {
94                return Err(err("SOCKS: credential too long"));
95            }
96            let mut req = vec![0x01, u.len() as u8];
97            req.extend_from_slice(u.as_bytes());
98            req.push(pw.len() as u8);
99            req.extend_from_slice(pw.as_bytes());
100            s.write_all(&req).await?;
101            let mut ar = [0u8; 2];
102            s.read_exact(&mut ar).await?;
103            if ar[1] != 0x00 {
104                return Err(err("SOCKS: username/password auth rejected"));
105            }
106        }
107        0xFF => return Err(err("SOCKS: proxy accepts no offered auth method")),
108        m => return Err(err(&format!("SOCKS: unexpected auth method {m}"))),
109    }
110
111    // CONNECT to the destination as a domain name (proxy-side DNS).
112    if dst_host.len() > 255 {
113        return Err(err("SOCKS: destination host too long"));
114    }
115    let mut req = vec![0x05, 0x01, 0x00, 0x03, dst_host.len() as u8];
116    req.extend_from_slice(dst_host.as_bytes());
117    req.extend_from_slice(&dst_port.to_be_bytes());
118    s.write_all(&req).await?;
119
120    // Reply: VER REP RSV ATYP BND.ADDR BND.PORT — consume the bound address so the stream is clean.
121    let mut head = [0u8; 4];
122    s.read_exact(&mut head).await?;
123    if head[1] != 0x00 {
124        return Err(err(&format!(
125            "SOCKS: CONNECT to {dst_host}:{dst_port} failed (reply code {})",
126            head[1]
127        )));
128    }
129    let addr_len = match head[3] {
130        0x01 => 4,
131        0x04 => 16,
132        0x03 => {
133            let mut l = [0u8; 1];
134            s.read_exact(&mut l).await?;
135            l[0] as usize
136        }
137        a => return Err(err(&format!("SOCKS: bad ATYP {a} in reply"))),
138    };
139    let mut rest = vec![0u8; addr_len + 2];
140    s.read_exact(&mut rest).await?;
141    Ok(s)
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn parse_plain_and_authed() {
150        let a = Socks5::parse("127.0.0.1:1080").unwrap();
151        assert_eq!(a.proxy, "127.0.0.1:1080");
152        assert!(a.auth.is_none());
153        let b = Socks5::parse("bob:s3cret@10.0.0.5:9050").unwrap();
154        assert_eq!(b.proxy, "10.0.0.5:9050");
155        assert_eq!(b.auth, Some(("bob".into(), "s3cret".into())));
156        assert!(Socks5::parse("nohost").is_none());
157    }
158
159    #[test]
160    fn host_port_split() {
161        assert_eq!(host_port("dc.corp:445", 999), ("dc.corp".into(), 445));
162        assert_eq!(host_port("dc.corp", 445), ("dc.corp".into(), 445));
163    }
164}
smb2_client/socks.rs

smb2_client/
socks.rs
