socks_proxy.rs

  1//! socks proxy
  2
  3use anyhow::{Context, Result};
  4use tokio::net::TcpStream;
  5use tokio_socks::tcp::{Socks4Stream, Socks5Stream};
  6use url::Url;
  7
  8use super::AsyncReadWrite;
  9
 10/// Identification to a Socks V4 Proxy
 11pub(super) struct Socks4Identification<'a> {
 12    user_id: &'a str,
 13}
 14
 15/// Authorization to a Socks V5 Proxy
 16pub(super) struct Socks5Authorization<'a> {
 17    username: &'a str,
 18    password: &'a str,
 19}
 20
 21/// Socks Proxy Protocol Version
 22///
 23/// V4 allows idenfication using a user_id
 24/// V5 allows authorization using a username and password
 25pub(super) enum SocksVersion<'a> {
 26    V4(Option<Socks4Identification<'a>>),
 27    V5(Option<Socks5Authorization<'a>>),
 28}
 29
 30pub(super) fn parse_socks_proxy<'t>(scheme: &str, proxy: &'t Url) -> SocksVersion<'t> {
 31    if scheme.starts_with("socks4") {
 32        let identification = match proxy.username() {
 33            "" => None,
 34            username => Some(Socks4Identification { user_id: username }),
 35        };
 36        SocksVersion::V4(identification)
 37    } else {
 38        let authorization = proxy.password().map(|password| Socks5Authorization {
 39            username: proxy.username(),
 40            password,
 41        });
 42        SocksVersion::V5(authorization)
 43    }
 44}
 45
 46pub(super) async fn connect_socks_proxy_stream(
 47    stream: TcpStream,
 48    socks_version: SocksVersion<'_>,
 49    rpc_host: (&str, u16),
 50) -> Result<Box<dyn AsyncReadWrite>> {
 51    match socks_version {
 52        SocksVersion::V4(None) => {
 53            let socks = Socks4Stream::connect_with_socket(stream, rpc_host)
 54                .await
 55                .context("error connecting to socks")?;
 56            Ok(Box::new(socks))
 57        }
 58        SocksVersion::V4(Some(Socks4Identification { user_id })) => {
 59            let socks = Socks4Stream::connect_with_userid_and_socket(stream, rpc_host, user_id)
 60                .await
 61                .context("error connecting to socks")?;
 62            Ok(Box::new(socks))
 63        }
 64        SocksVersion::V5(None) => {
 65            let socks = Socks5Stream::connect_with_socket(stream, rpc_host)
 66                .await
 67                .context("error connecting to socks")?;
 68            Ok(Box::new(socks))
 69        }
 70        SocksVersion::V5(Some(Socks5Authorization { username, password })) => {
 71            let socks = Socks5Stream::connect_with_password_and_socket(
 72                stream, rpc_host, username, password,
 73            )
 74            .await
 75            .context("error connecting to socks")?;
 76            Ok(Box::new(socks))
 77        }
 78    }
 79}
 80
 81#[cfg(test)]
 82mod tests {
 83    use url::Url;
 84
 85    use super::*;
 86
 87    #[test]
 88    fn parse_socks4() {
 89        let proxy = Url::parse("socks4://proxy.example.com:1080").unwrap();
 90        let scheme = proxy.scheme();
 91
 92        let version = parse_socks_proxy(scheme, &proxy);
 93        assert!(matches!(version, SocksVersion::V4(None)))
 94    }
 95
 96    #[test]
 97    fn parse_socks4_with_identification() {
 98        let proxy = Url::parse("socks4://userid@proxy.example.com:1080").unwrap();
 99        let scheme = proxy.scheme();
100
101        let version = parse_socks_proxy(scheme, &proxy);
102        assert!(matches!(
103            version,
104            SocksVersion::V4(Some(Socks4Identification { user_id: "userid" }))
105        ))
106    }
107
108    #[test]
109    fn parse_socks5() {
110        let proxy = Url::parse("socks5://proxy.example.com:1080").unwrap();
111        let scheme = proxy.scheme();
112
113        let version = parse_socks_proxy(scheme, &proxy);
114        assert!(matches!(version, SocksVersion::V5(None)))
115    }
116
117    #[test]
118    fn parse_socks5_with_authorization() {
119        let proxy = Url::parse("socks5://username:password@proxy.example.com:1080").unwrap();
120        let scheme = proxy.scheme();
121
122        let version = parse_socks_proxy(scheme, &proxy);
123        assert!(matches!(
124            version,
125            SocksVersion::V5(Some(Socks5Authorization {
126                username: "username",
127                password: "password"
128            }))
129        ))
130    }
131}