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    let s = match proxy() {
70        Some(cfg) => socks5_connect(cfg, &h, p).await?,
71        None => TcpStream::connect((h.as_str(), p)).await?,
72    };
73    // Disable Nagle: SMB/RPC does many small writes (opens/queries ~90-200B sealed);
74    // Nagle+delayed-ACK adds up to 40ms per call. -300..500ms on secretsdump.
75    let _ = s.set_nodelay(true);
76    Ok(s)
77}
78
79async fn socks5_connect(cfg: &Socks5, dst_host: &str, dst_port: u16) -> Result<TcpStream> {
80    let mut s = TcpStream::connect(&cfg.proxy).await?;
81
82    // Greeting: offer no-auth (and user/pass if we have creds).
83    if cfg.auth.is_some() {
84        s.write_all(&[0x05, 0x02, 0x00, 0x02]).await?;
85    } else {
86        s.write_all(&[0x05, 0x01, 0x00]).await?;
87    }
88    let mut sel = [0u8; 2];
89    s.read_exact(&mut sel).await?;
90    if sel[0] != 0x05 {
91        return Err(err("SOCKS: bad version in method reply"));
92    }
93    match sel[1] {
94        0x00 => {}
95        0x02 => {
96            let (u, pw) = cfg.auth.as_ref().ok_or_else(|| err("SOCKS: proxy demands auth but none given"))?;
97            if u.len() > 255 || pw.len() > 255 {
98                return Err(err("SOCKS: credential too long"));
99            }
100            let mut req = vec![0x01, u.len() as u8];
101            req.extend_from_slice(u.as_bytes());
102            req.push(pw.len() as u8);
103            req.extend_from_slice(pw.as_bytes());
104            s.write_all(&req).await?;
105            let mut ar = [0u8; 2];
106            s.read_exact(&mut ar).await?;
107            if ar[1] != 0x00 {
108                return Err(err("SOCKS: username/password auth rejected"));
109            }
110        }
111        0xFF => return Err(err("SOCKS: proxy accepts no offered auth method")),
112        m => return Err(err(&format!("SOCKS: unexpected auth method {m}"))),
113    }
114
115    // CONNECT to the destination as a domain name (proxy-side DNS).
116    if dst_host.len() > 255 {
117        return Err(err("SOCKS: destination host too long"));
118    }
119    let mut req = vec![0x05, 0x01, 0x00, 0x03, dst_host.len() as u8];
120    req.extend_from_slice(dst_host.as_bytes());
121    req.extend_from_slice(&dst_port.to_be_bytes());
122    s.write_all(&req).await?;
123
124    // Reply: VER REP RSV ATYP BND.ADDR BND.PORT — consume the bound address so the stream is clean.
125    let mut head = [0u8; 4];
126    s.read_exact(&mut head).await?;
127    if head[1] != 0x00 {
128        return Err(err(&format!(
129            "SOCKS: CONNECT to {dst_host}:{dst_port} failed (reply code {})",
130            head[1]
131        )));
132    }
133    let addr_len = match head[3] {
134        0x01 => 4,
135        0x04 => 16,
136        0x03 => {
137            let mut l = [0u8; 1];
138            s.read_exact(&mut l).await?;
139            l[0] as usize
140        }
141        a => return Err(err(&format!("SOCKS: bad ATYP {a} in reply"))),
142    };
143    let mut rest = vec![0u8; addr_len + 2];
144    s.read_exact(&mut rest).await?;
145    Ok(s)
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn parse_plain_and_authed() {
154        let a = Socks5::parse("127.0.0.1:1080").unwrap();
155        assert_eq!(a.proxy, "127.0.0.1:1080");
156        assert!(a.auth.is_none());
157        let b = Socks5::parse("bob:s3cret@10.0.0.5:9050").unwrap();
158        assert_eq!(b.proxy, "10.0.0.5:9050");
159        assert_eq!(b.auth, Some(("bob".into(), "s3cret".into())));
160        assert!(Socks5::parse("nohost").is_none());
161    }
162
163    #[test]
164    fn host_port_split() {
165        assert_eq!(host_port("dc.corp:445", 999), ("dc.corp".into(), 445));
166        assert_eq!(host_port("dc.corp", 445), ("dc.corp".into(), 445));
167    }
168}
smb2_client/socks.rs

smb2_client/
socks.rs
