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, 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::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
97                .auth
98                .as_ref()
99                .ok_or_else(|| err("SOCKS: proxy demands auth but none given"))?;
100            if u.len() > 255 || pw.len() > 255 {
101                return Err(err("SOCKS: credential too long"));
102            }
103            let mut req = vec![0x01, u.len() as u8];
104            req.extend_from_slice(u.as_bytes());
105            req.push(pw.len() as u8);
106            req.extend_from_slice(pw.as_bytes());
107            s.write_all(&req).await?;
108            let mut ar = [0u8; 2];
109            s.read_exact(&mut ar).await?;
110            if ar[1] != 0x00 {
111                return Err(err("SOCKS: username/password auth rejected"));
112            }
113        }
114        0xFF => return Err(err("SOCKS: proxy accepts no offered auth method")),
115        m => return Err(err(&format!("SOCKS: unexpected auth method {m}"))),
116    }
117
118    // CONNECT to the destination as a domain name (proxy-side DNS).
119    if dst_host.len() > 255 {
120        return Err(err("SOCKS: destination host too long"));
121    }
122    let mut req = vec![0x05, 0x01, 0x00, 0x03, dst_host.len() as u8];
123    req.extend_from_slice(dst_host.as_bytes());
124    req.extend_from_slice(&dst_port.to_be_bytes());
125    s.write_all(&req).await?;
126
127    // Reply: VER REP RSV ATYP BND.ADDR BND.PORT — consume the bound address so the stream is clean.
128    let mut head = [0u8; 4];
129    s.read_exact(&mut head).await?;
130    if head[1] != 0x00 {
131        return Err(err(&format!(
132            "SOCKS: CONNECT to {dst_host}:{dst_port} failed (reply code {})",
133            head[1]
134        )));
135    }
136    let addr_len = match head[3] {
137        0x01 => 4,
138        0x04 => 16,
139        0x03 => {
140            let mut l = [0u8; 1];
141            s.read_exact(&mut l).await?;
142            l[0] as usize
143        }
144        a => return Err(err(&format!("SOCKS: bad ATYP {a} in reply"))),
145    };
146    let mut rest = vec![0u8; addr_len + 2];
147    s.read_exact(&mut rest).await?;
148    Ok(s)
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    fn parse_plain_and_authed() {
157        let a = Socks5::parse("127.0.0.1:1080").unwrap();
158        assert_eq!(a.proxy, "127.0.0.1:1080");
159        assert!(a.auth.is_none());
160        let b = Socks5::parse("bob:s3cret@10.0.0.5:9050").unwrap();
161        assert_eq!(b.proxy, "10.0.0.5:9050");
162        assert_eq!(b.auth, Some(("bob".into(), "s3cret".into())));
163        assert!(Socks5::parse("nohost").is_none());
164    }
165
166    #[test]
167    fn host_port_split() {
168        assert_eq!(host_port("dc.corp:445", 999), ("dc.corp".into(), 445));
169        assert_eq!(host_port("dc.corp", 445), ("dc.corp".into(), 445));
170    }
171}
smb2_client/socks.rs

smb2_client/
socks.rs
