socks5-impl 0.9.8

use super::{Address, UserKey};
use crate::{Error, Result};

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
#[cfg_attr(feature = "serde", serde(try_from = "String", into = "String"))]
pub struct ProxyParameters {
    pub proxy_type: ProxyType,
    pub addr: Option<Address>,
    #[cfg_attr(feature = "serde", serde(skip_serializing_if = "Option::is_none"))]
    pub credentials: Option<UserKey>,
}

impl Default for ProxyParameters {
    fn default() -> Self {
        ProxyParameters {
            proxy_type: ProxyType::Socks5,
            addr: Some("127.0.0.1:1080".parse().unwrap()),
            credentials: None,
        }
    }
}

impl std::fmt::Display for ProxyParameters {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.proxy_type == ProxyType::None || self.addr.is_none() {
            return write!(f, "none");
        }

        let auth = match &self.credentials {
            Some(creds) => format!("{creds}"),
            None => "".to_owned(),
        };
        let addr = self.addr.as_ref().unwrap();
        if auth.is_empty() {
            write!(f, "{}://{addr}", self.proxy_type)
        } else {
            write!(f, "{}://{}@{addr}", self.proxy_type, auth)
        }
    }
}

#[cfg(feature = "serde")]
impl From<ProxyParameters> for String {
    fn from(parameters: ProxyParameters) -> Self {
        if parameters.proxy_type == ProxyType::None {
            "none".to_owned()
        } else {
            parameters.to_string()
        }
    }
}

impl ProxyParameters {
    pub fn new(proxy_type: ProxyType, addr: Option<Address>, credentials: Option<UserKey>) -> Self {
        ProxyParameters {
            proxy_type,
            addr,
            credentials,
        }
    }
}

impl TryFrom<String> for ProxyParameters {
    type Error = Error;
    fn try_from(s: String) -> Result<Self> {
        Self::try_from(s.as_str())
    }
}

impl TryFrom<&String> for ProxyParameters {
    type Error = Error;
    fn try_from(s: &String) -> Result<Self> {
        Self::try_from(s.as_str())
    }
}

impl TryFrom<&str> for ProxyParameters {
    type Error = Error;
    fn try_from(s: &str) -> Result<Self> {
        if s == "none" {
            return Ok(ProxyParameters {
                proxy_type: ProxyType::None,
                addr: None,
                credentials: None,
            });
        }

        let e = format!("`{s}` is not a valid proxy URL");
        let url = url::Url::parse(s).map_err(|_| Error::from(&e))?;
        let e = format!("`{s}` does not contain a host");
        let host = url.host_str().ok_or(Error::from(e))?;

        let e = format!("`{s}` does not contain a port");
        let port = url.port_or_known_default().ok_or(Error::from(&e))?;

        let addr = (host, port).into();

        let credentials = if url.username() == "" && url.password().is_none() {
            None
        } else {
            use percent_encoding::percent_decode;
            let username = percent_decode(url.username().as_bytes()).decode_utf8()?;
            let password = percent_decode(url.password().unwrap_or("").as_bytes()).decode_utf8()?;
            Some(UserKey::new(username, password))
        };

        let proxy_type = url.scheme().to_ascii_lowercase().as_str().try_into()?;

        Ok(ProxyParameters {
            proxy_type,
            addr: Some(addr),
            credentials,
        })
    }
}

impl std::str::FromStr for ProxyParameters {
    type Err = Error;
    fn from_str(s: &str) -> Result<Self> {
        Self::try_from(s)
    }
}

impl TryFrom<url::Url> for ProxyParameters {
    type Error = Error;
    fn try_from(url: url::Url) -> Result<Self> {
        let e = format!("`{url}` does not contain a host");
        let host = url.host_str().ok_or(Error::from(e))?;

        let e = format!("`{url}` does not contain a port");
        let port = url.port_or_known_default().ok_or(Error::from(&e))?;

        let addr = (host, port).into();

        let credentials = if url.username() == "" && url.password().is_none() {
            None
        } else {
            use percent_encoding::percent_decode;
            let username = percent_decode(url.username().as_bytes()).decode_utf8()?;
            let password = percent_decode(url.password().unwrap_or("").as_bytes()).decode_utf8()?;
            Some(UserKey::new(username, password))
        };

        let proxy_type = url.scheme().to_ascii_lowercase().as_str().try_into()?;

        Ok(ProxyParameters {
            proxy_type,
            addr: Some(addr),
            credentials,
        })
    }
}

impl TryFrom<ProxyParameters> for url::Url {
    type Error = Error;
    fn try_from(parameters: ProxyParameters) -> Result<Self> {
        let addr = parameters
            .addr
            .ok_or_else(|| Error::from("proxy parameters do not contain an address"))?;
        let mut url = url::Url::parse(&format!("{}://{addr}", parameters.proxy_type))
            .map_err(|e| Error::from(&format!("failed to parse proxy parameters into URL: {e}")))?;
        if let Some(creds) = parameters.credentials {
            url.set_username(&creds.username)
                .map_err(|e| Error::from(&format!("failed to set username in proxy URL: {e:?}")))?;
            url.set_password(Some(&creds.password))
                .map_err(|e| Error::from(&format!("failed to set password in proxy URL: {e:?}")))?;
        }
        Ok(url)
    }
}

#[repr(C)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Default, Hash)]
pub enum ProxyType {
    None = 0,
    Http,
    Socks4,
    #[default]
    Socks5,
    Mixed,
}

impl TryFrom<&str> for ProxyType {
    type Error = Error;
    fn try_from(value: &str) -> Result<Self> {
        match value.to_ascii_lowercase().as_str() {
            "none" => Ok(ProxyType::None),
            "http" => Ok(ProxyType::Http),
            "socks4" => Ok(ProxyType::Socks4),
            "socks5" => Ok(ProxyType::Socks5),
            "mixed" => Ok(ProxyType::Mixed),
            scheme => Err(Error::from(&format!("`{scheme}` is an invalid proxy type"))),
        }
    }
}

impl std::fmt::Display for ProxyType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ProxyType::None => write!(f, "none"),
            ProxyType::Socks4 => write!(f, "socks4"),
            ProxyType::Socks5 => write!(f, "socks5"),
            ProxyType::Http => write!(f, "http"),
            ProxyType::Mixed => write!(f, "mixed"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn default_proxy_parameters() {
        let parameters = ProxyParameters::default();

        assert_eq!(parameters.proxy_type, ProxyType::Socks5);
        assert_eq!(parameters.addr, Some("127.0.0.1:1080".parse().unwrap()));
        assert_eq!(parameters.credentials, None);
        assert_eq!(parameters.to_string(), "socks5://127.0.0.1:1080");
    }

    #[test]
    fn parse_without_credentials() {
        let parameters = "socks5://123.45.67.89:1080".parse::<ProxyParameters>().unwrap();
        assert_eq!(parameters.proxy_type, ProxyType::Socks5);
        assert_eq!(parameters.addr, Some(("123.45.67.89", 1080).into()));
        assert_eq!(parameters.addr.as_ref().unwrap().get_type(), crate::protocol::AddressType::IPv4);
        assert_eq!(parameters.credentials, None);
        assert_eq!(parameters.to_string(), "socks5://123.45.67.89:1080");

        let parameters = "socks5://proxy.example.com:1080".parse::<ProxyParameters>().unwrap();

        assert_eq!(parameters.proxy_type, ProxyType::Socks5);
        assert_eq!(parameters.addr, Some(("proxy.example.com", 1080).into()));
        assert_eq!(parameters.credentials, None);
        assert_eq!(parameters.to_string(), "socks5://proxy.example.com:1080");

        let parameters = "http://proxy.example.com:8080".parse::<ProxyParameters>().unwrap();
        assert_eq!(parameters.proxy_type, ProxyType::Http);
        assert_eq!(parameters.addr, Some(("proxy.example.com", 8080).into()));
        assert_eq!(parameters.credentials, None);
        assert_eq!(parameters.to_string(), "http://proxy.example.com:8080");

        let parameters = "http://proxy.example.com".parse::<ProxyParameters>().unwrap();
        assert_eq!(parameters.proxy_type, ProxyType::Http);
        assert_eq!(parameters.addr, Some(("proxy.example.com", 80).into()));
        assert_eq!(parameters.credentials, None);
        assert_eq!(parameters.to_string(), "http://proxy.example.com:80");

        assert!("socks5://proxy.example.com".parse::<ProxyParameters>().is_err());
    }

    #[test]
    fn parse_with_credentials() {
        let parameters = "socks5://user:password@proxy.example.com:1080".parse::<ProxyParameters>().unwrap();
        assert_eq!(parameters.proxy_type, ProxyType::Socks5);
        assert_eq!(parameters.addr, Some(("proxy.example.com", 1080).into()));
        assert_eq!(parameters.credentials, Some(UserKey::new("user", "password")));
        assert_eq!(parameters.to_string(), "socks5://user:password@proxy.example.com:1080");

        let parameters = "socks5://user@123.45.67.89:1080".parse::<ProxyParameters>().unwrap();
        assert_eq!(parameters.proxy_type, ProxyType::Socks5);
        assert_eq!(parameters.addr, Some(("123.45.67.89", 1080).into()));
        assert_eq!(parameters.credentials, Some(UserKey::new("user", "")));
        assert_eq!(parameters.to_string(), "socks5://user@123.45.67.89:1080");

        let parameters = "socks5://:password@123.45.67.89:1080".parse::<ProxyParameters>().unwrap();
        assert_eq!(parameters.proxy_type, ProxyType::Socks5);
        assert_eq!(parameters.addr, Some(("123.45.67.89", 1080).into()));
        assert_eq!(parameters.credentials, Some(UserKey::new("", "password")));
        assert_eq!(parameters.to_string(), "socks5://:password@123.45.67.89:1080");
    }

    #[test]
    fn parse_with_percent_encoded_credentials() {
        let parameters = "socks5://user%40name:pa%24%24@proxy.example.com:1080"
            .parse::<ProxyParameters>()
            .unwrap();

        assert_eq!(parameters.proxy_type, ProxyType::Socks5);
        assert_eq!(parameters.addr, Some(("proxy.example.com", 1080).into()));
        assert_eq!(parameters.credentials, Some(UserKey::new("user@name", "pa$$")));
        assert_eq!(parameters.to_string(), "socks5://user%40name:pa%24%24@proxy.example.com:1080");
    }

    #[test]
    fn parse_none_proxy() {
        let parameters = "none".parse::<ProxyParameters>().unwrap();

        assert_eq!(parameters.proxy_type, ProxyType::None);
        assert_eq!(parameters.addr, None);
        assert_eq!(parameters.credentials, None);
        assert_eq!(parameters.to_string(), "none");
        assert!(url::Url::try_from(parameters).is_err());
    }

    #[test]
    fn parse_invalid_proxy_type() {
        let err = "ftp://proxy.example.com:21".parse::<ProxyParameters>().unwrap_err();
        assert!(format!("{err}").contains("invalid proxy type"));
    }

    #[test]
    fn try_from_string_and_url_work() {
        let s = "http://proxy.example.com:8080".to_string();
        let parameters = ProxyParameters::try_from(s.clone()).unwrap();
        assert_eq!(parameters.proxy_type, ProxyType::Http);
        assert_eq!(parameters.addr, Some(("proxy.example.com", 8080).into()));
        assert_eq!(parameters.to_string(), "http://proxy.example.com:8080");

        let url = url::Url::parse(&s).unwrap();
        let parameters = ProxyParameters::try_from(url).unwrap();
        assert_eq!(parameters.proxy_type, ProxyType::Http);
        assert_eq!(parameters.addr, Some(("proxy.example.com", 8080).into()));
    }

    #[test]
    fn proxy_parameters_new_and_into_url() {
        let parameters = ProxyParameters::new(
            ProxyType::Socks5,
            Some(("proxy.example.com", 1080).into()),
            Some(UserKey::new("user", "password")),
        );

        let url = url::Url::try_from(parameters.clone()).unwrap();
        assert_eq!(url.as_str(), "socks5://user:password@proxy.example.com:1080");
        assert_eq!(parameters.to_string(), "socks5://user:password@proxy.example.com:1080");
    }

    #[cfg(feature = "serde")]
    #[test]
    fn serde_string_round_trip() {
        let parameters: ProxyParameters = serde_json::from_str(r#""mixed://127.0.0.1:3080""#).unwrap();
        assert_eq!(parameters.to_string(), "mixed://127.0.0.1:3080");
        assert_eq!(serde_json::to_string(&parameters).unwrap(), r#""mixed://127.0.0.1:3080""#);

        let no_proxy: ProxyParameters = serde_json::from_str(r#""none""#).unwrap();
        assert_eq!(no_proxy.proxy_type, ProxyType::None);
        assert_eq!(serde_json::to_string(&no_proxy).unwrap(), r#""none""#);
    }
}