hottub/src/util/proxy.rs

use std::fmt;
use std::str::FromStr;

use serde::Serialize;
use url::Url;
use wreq::Proxy as WreqProxy;

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct Proxy {
    pub protocol: String,
    pub host: String,
    pub port: u16,
    pub username: Option<String>,
    pub password: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProxyParseError {
    message: String,
}

impl ProxyParseError {
    fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }
}

impl fmt::Display for ProxyParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.message)
    }
}

impl std::error::Error for ProxyParseError {}

impl fmt::Display for Proxy {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match (&self.username, &self.password) {
            (Some(username), Some(password)) => write!(
                f,
                "{}://{}:{}@{}:{}",
                self.protocol, username, password, self.host, self.port
            ),
            (Some(username), None) => write!(
                f,
                "{}://{}@{}:{}",
                self.protocol, username, self.host, self.port
            ),
            (None, Some(_)) => write!(f, "{}://{}:{}", self.protocol, self.host, self.port),
            (None, None) => write!(f, "{}://{}:{}", self.protocol, self.host, self.port),
        }
    }
}

impl Proxy {
    /// `protocol://host:port`, without credentials: safe to show or log.
    pub fn redacted(&self) -> String {
        format!("{}://{}:{}", self.protocol, self.host, self.port)
    }

    pub fn parse(input: &str) -> Result<Self, ProxyParseError> {
        input.parse()
    }

    pub fn to_wreq_proxy(&self) -> Result<WreqProxy, ProxyParseError> {
        let base_url = format!("{}://{}:{}", self.protocol, self.host, self.port);
        let proxy_url = Url::parse(&base_url)
            .map_err(|e| ProxyParseError::new(format!("invalid proxy url: {e}")))?;

        let mut proxy = WreqProxy::all(proxy_url.as_str())
            .map_err(|e| ProxyParseError::new(format!("failed to build wreq proxy: {e}")))?;

        if let Some(username) = &self.username {
            let password = self.password.as_deref().unwrap_or("");
            proxy = proxy.basic_auth(username, password);
        }

        Ok(proxy)
    }
}

impl FromStr for Proxy {
    type Err = ProxyParseError;

    fn from_str(input: &str) -> Result<Self, Self::Err> {
        let trimmed = input.trim();
        if trimmed.is_empty() {
            return Err(ProxyParseError::new("proxy string is empty"));
        }

        let with_scheme = if trimmed.contains("://") {
            trimmed.to_string()
        } else {
            format!("http://{trimmed}")
        };

        let url = Url::parse(&with_scheme)
            .map_err(|e| ProxyParseError::new(format!("invalid proxy url: {e}")))?;
        if !(url.path().is_empty() || url.path() == "/")
            || url.query().is_some()
            || url.fragment().is_some()
        {
            return Err(ProxyParseError::new(format!(
                "proxy url must not include path, query, or fragment: {:?}",
                input
            )));
        }

        let host = url
            .host_str()
            .ok_or_else(|| ProxyParseError::new("proxy url is missing host"))?
            .to_string();

        let port = url.port().unwrap_or(80);

        Ok(Proxy {
            protocol: url.scheme().to_string(),
            host,
            port,
            username: match url.username() {
                "" => None,
                username => Some(username.to_string()),
            },
            password: url.password().map(|password| password.to_string()),
        })
    }
}