socks_proxy.rs

  1//! socks proxy
  2
  3use std::net::SocketAddr;
  4
  5use anyhow::{Context as _, Result};
  6use http_client::Url;
  7use tokio::net::TcpStream;
  8use tokio_socks::{
  9    IntoTargetAddr, TargetAddr,
 10    tcp::{Socks4Stream, Socks5Stream},
 11};
 12
 13use crate::proxy::SYSTEM_DNS_RESOLVER;
 14
 15use super::AsyncReadWrite;
 16
 17/// Identification to a Socks V4 Proxy
 18pub(super) struct Socks4Identification<'a> {
 19    user_id: &'a str,
 20}
 21
 22/// Authorization to a Socks V5 Proxy
 23pub(super) struct Socks5Authorization<'a> {
 24    username: &'a str,
 25    password: &'a str,
 26}
 27
 28/// Socks Proxy Protocol Version
 29///
 30/// V4 allows idenfication using a user_id
 31/// V5 allows authorization using a username and password
 32pub(super) enum SocksVersion<'a> {
 33    V4 {
 34        local_dns: bool,
 35        identification: Option<Socks4Identification<'a>>,
 36    },
 37    V5 {
 38        local_dns: bool,
 39        authorization: Option<Socks5Authorization<'a>>,
 40    },
 41}
 42
 43pub(super) fn parse_socks_proxy<'t>(scheme: &str, proxy: &'t Url) -> SocksVersion<'t> {
 44    if scheme.starts_with("socks4") {
 45        let identification = match proxy.username() {
 46            "" => None,
 47            username => Some(Socks4Identification { user_id: username }),
 48        };
 49        SocksVersion::V4 {
 50            local_dns: scheme != "socks4a",
 51            identification,
 52        }
 53    } else {
 54        let authorization = proxy.password().map(|password| Socks5Authorization {
 55            username: proxy.username(),
 56            password,
 57        });
 58        SocksVersion::V5 {
 59            local_dns: scheme != "socks5h",
 60            authorization,
 61        }
 62    }
 63}
 64
 65pub(super) async fn connect_socks_proxy_stream(
 66    stream: TcpStream,
 67    socks_version: SocksVersion<'_>,
 68    rpc_host: (&str, u16),
 69) -> Result<Box<dyn AsyncReadWrite>> {
 70    let rpc_host = rpc_host
 71        .into_target_addr()
 72        .context("Failed to parse target addr")?;
 73
 74    let local_dns = match &socks_version {
 75        SocksVersion::V4 { local_dns, .. } => local_dns,
 76        SocksVersion::V5 { local_dns, .. } => local_dns,
 77    };
 78    let rpc_host = match (rpc_host, local_dns) {
 79        (TargetAddr::Domain(domain, port), true) => {
 80            let ip_addr = SYSTEM_DNS_RESOLVER
 81                .lookup_ip(domain.as_ref())
 82                .await
 83                .with_context(|| format!("Failed to lookup domain {}", domain))?
 84                .into_iter()
 85                .next()
 86                .ok_or_else(|| anyhow::anyhow!("Failed to lookup domain {}", domain))?;
 87            TargetAddr::Ip(SocketAddr::new(ip_addr, port))
 88        }
 89        (rpc_host, _) => rpc_host,
 90    };
 91
 92    match socks_version {
 93        SocksVersion::V4 {
 94            identification: None,
 95            ..
 96        } => {
 97            let socks = Socks4Stream::connect_with_socket(stream, rpc_host)
 98                .await
 99                .context("error connecting to socks")?;
100            Ok(Box::new(socks))
101        }
102        SocksVersion::V4 {
103            identification: Some(Socks4Identification { user_id }),
104            ..
105        } => {
106            let socks = Socks4Stream::connect_with_userid_and_socket(stream, rpc_host, user_id)
107                .await
108                .context("error connecting to socks")?;
109            Ok(Box::new(socks))
110        }
111        SocksVersion::V5 {
112            authorization: None,
113            ..
114        } => {
115            let socks = Socks5Stream::connect_with_socket(stream, rpc_host)
116                .await
117                .context("error connecting to socks")?;
118            Ok(Box::new(socks))
119        }
120        SocksVersion::V5 {
121            authorization: Some(Socks5Authorization { username, password }),
122            ..
123        } => {
124            let socks = Socks5Stream::connect_with_password_and_socket(
125                stream, rpc_host, username, password,
126            )
127            .await
128            .context("error connecting to socks")?;
129            Ok(Box::new(socks))
130        }
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use url::Url;
137
138    use super::*;
139
140    #[test]
141    fn parse_socks4() {
142        let proxy = Url::parse("socks4://proxy.example.com:1080").unwrap();
143        let scheme = proxy.scheme();
144
145        let version = parse_socks_proxy(scheme, &proxy);
146        assert!(matches!(
147            version,
148            SocksVersion::V4 {
149                local_dns: true,
150                identification: None
151            }
152        ))
153    }
154
155    #[test]
156    fn parse_socks4_with_identification() {
157        let proxy = Url::parse("socks4://userid@proxy.example.com:1080").unwrap();
158        let scheme = proxy.scheme();
159
160        let version = parse_socks_proxy(scheme, &proxy);
161        assert!(matches!(
162            version,
163            SocksVersion::V4 {
164                local_dns: true,
165                identification: Some(Socks4Identification { user_id: "userid" })
166            }
167        ))
168    }
169
170    #[test]
171    fn parse_socks4_with_remote_dns() {
172        let proxy = Url::parse("socks4a://proxy.example.com:1080").unwrap();
173        let scheme = proxy.scheme();
174
175        let version = parse_socks_proxy(scheme, &proxy);
176        assert!(matches!(
177            version,
178            SocksVersion::V4 {
179                local_dns: false,
180                identification: None
181            }
182        ))
183    }
184
185    #[test]
186    fn parse_socks5() {
187        let proxy = Url::parse("socks5://proxy.example.com:1080").unwrap();
188        let scheme = proxy.scheme();
189
190        let version = parse_socks_proxy(scheme, &proxy);
191        assert!(matches!(
192            version,
193            SocksVersion::V5 {
194                local_dns: true,
195                authorization: None
196            }
197        ))
198    }
199
200    #[test]
201    fn parse_socks5_with_authorization() {
202        let proxy = Url::parse("socks5://username:password@proxy.example.com:1080").unwrap();
203        let scheme = proxy.scheme();
204
205        let version = parse_socks_proxy(scheme, &proxy);
206        assert!(matches!(
207            version,
208            SocksVersion::V5 {
209                local_dns: true,
210                authorization: Some(Socks5Authorization {
211                    username: "username",
212                    password: "password"
213                })
214            }
215        ))
216    }
217
218    #[test]
219    fn parse_socks5_with_remote_dns() {
220        let proxy = Url::parse("socks5h://proxy.example.com:1080").unwrap();
221        let scheme = proxy.scheme();
222
223        let version = parse_socks_proxy(scheme, &proxy);
224        assert!(matches!(
225            version,
226            SocksVersion::V5 {
227                local_dns: false,
228                authorization: None
229            }
230        ))
231    }
232}