ws2tcp-local 0.1.4

use std::{
    collections::HashSet,
    fs::{self, FileTimes},
    path::{Path, PathBuf},
    time::{Duration, SystemTime},
};

use anyhow::{Context, Result, anyhow, bail};
use base64::{Engine, engine::general_purpose::STANDARD};
use reqwest::{
    Client,
    header::{HeaderMap, LAST_MODIFIED},
};
use tracing::{info, warn};

use crate::cli::ProxyMode;

const GFWLIST_URL: &str = "https://gitlab.com/gfwlist/gfwlist/raw/master/gfwlist.txt";
const CACHE_DIR_NAME: &str = "ws2tcp-local";
const GFWLIST_CACHE_FILE: &str = "gfwlist.txt";

#[derive(Debug, Clone)]
pub(crate) enum RoutingRules {
    Domains {
        rules: DomainRules,
        custom_domain_rules: Option<PathBuf>,
    },
    GlobalProxy,
    AllProxyFallback,
}

impl RoutingRules {
    pub(crate) async fn load(proxy_mode: ProxyMode, custom_domain_rules: Option<&Path>) -> Self {
        if proxy_mode == ProxyMode::Global {
            info!("using global proxy mode; skipping proxy routing rule download");
            return Self::GlobalProxy;
        }

        match Self::download_and_parse(custom_domain_rules).await {
            Ok(rules) => {
                info!(
                    url = GFWLIST_URL,
                    custom_domain_rules =
                        custom_domain_rules.map(|path| path.display().to_string()),
                    domain_count = rules.len(),
                    "loaded proxy routing rules"
                );
                Self::Domains {
                    rules,
                    custom_domain_rules: custom_domain_rules.map(Path::to_path_buf),
                }
            }
            Err(err) => {
                warn!(
                    url = GFWLIST_URL,
                    custom_domain_rules = custom_domain_rules.map(|path| path.display().to_string()),
                    error = %format_args!("{err:#}"),
                    "failed to load proxy routing rules; proxying all domains"
                );
                Self::AllProxyFallback
            }
        }
    }

    pub(crate) fn should_proxy_host(&self, host: &str) -> bool {
        match self {
            Self::Domains { rules, .. } => rules.matches(host),
            Self::GlobalProxy | Self::AllProxyFallback => true,
        }
    }

    fn mode(&self) -> &'static str {
        match self {
            Self::Domains { .. } => "auto",
            Self::GlobalProxy => "global",
            Self::AllProxyFallback => "all-proxy",
        }
    }

    pub(crate) fn describe(&self) -> String {
        match self {
            Self::Domains {
                rules,
                custom_domain_rules: Some(path),
            } => format!(
                "{} domains from {} plus custom rules from {}",
                rules.len(),
                GFWLIST_URL,
                path.display()
            ),
            Self::Domains {
                rules,
                custom_domain_rules: None,
            } => format!("{} domains from {}", rules.len(), GFWLIST_URL),
            Self::GlobalProxy => "all domains via proxy; proxy mode is global".to_owned(),
            Self::AllProxyFallback => {
                format!("all domains via proxy; failed to load {GFWLIST_URL}")
            }
        }
    }

    async fn download_and_parse(custom_domain_rules: Option<&Path>) -> Result<DomainRules> {
        let body = load_gfwlist_body().await?;

        let mut rules = parse_gfwlist(&body)?;
        if let Some(path) = custom_domain_rules {
            let custom_domains = read_custom_domain_rules(path)?;
            let custom_count = custom_domains.len();
            rules.extend(custom_domains);
            info!(
                path = %path.display(),
                custom_domain_count = custom_count,
                "merged custom proxy routing rules"
            );
        }

        Ok(rules)
    }
}

async fn load_gfwlist_body() -> Result<Vec<u8>> {
    let cache_path = gfwlist_cache_path()?;
    let client = Client::new();
    let remote_modified = fetch_remote_last_modified(&client).await?;

    if let Some(remote_modified) = remote_modified
        && is_cache_current(&cache_path, remote_modified)?
    {
        info!(
            cache_path = %cache_path.display(),
            url = GFWLIST_URL,
            "using cached gfwlist"
        );
        return fs::read(&cache_path)
            .with_context(|| format!("failed to read cached gfwlist {}", cache_path.display()));
    }

    let response = client
        .get(GFWLIST_URL)
        .send()
        .await
        .with_context(|| format!("failed to download {GFWLIST_URL}"))?
        .error_for_status()
        .with_context(|| format!("failed to download {GFWLIST_URL}"))?;
    let downloaded_modified = parse_last_modified(response.headers()).or(remote_modified);
    let body = response
        .bytes()
        .await
        .context("failed to read gfwlist response body")?;
    write_gfwlist_cache(&cache_path, &body, downloaded_modified)?;

    Ok(body.to_vec())
}

async fn fetch_remote_last_modified(client: &Client) -> Result<Option<SystemTime>> {
    let response = client
        .head(GFWLIST_URL)
        .send()
        .await
        .with_context(|| format!("failed to check remote gfwlist timestamp {GFWLIST_URL}"))?
        .error_for_status()
        .with_context(|| format!("failed to check remote gfwlist timestamp {GFWLIST_URL}"))?;

    Ok(parse_last_modified(response.headers()))
}

fn parse_last_modified(headers: &HeaderMap) -> Option<SystemTime> {
    headers
        .get(LAST_MODIFIED)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| httpdate::parse_http_date(value).ok())
}

fn gfwlist_cache_path() -> Result<PathBuf> {
    let cache_dir = user_cache_dir()?;

    Ok(cache_dir.join(CACHE_DIR_NAME).join(GFWLIST_CACHE_FILE))
}

#[cfg(windows)]
fn user_cache_dir() -> Result<PathBuf> {
    if let Some(path) = std::env::var_os("LOCALAPPDATA") {
        return Ok(PathBuf::from(path));
    }

    let profile = std::env::var_os("USERPROFILE")
        .context("USERPROFILE is not set; cannot locate gfwlist cache")?;
    Ok(PathBuf::from(profile).join("AppData").join("Local"))
}

#[cfg(target_os = "macos")]
fn user_cache_dir() -> Result<PathBuf> {
    let home = std::env::var_os("HOME").context("HOME is not set; cannot locate gfwlist cache")?;
    Ok(PathBuf::from(home).join("Library").join("Caches"))
}

#[cfg(all(not(windows), not(target_os = "macos")))]
fn user_cache_dir() -> Result<PathBuf> {
    match std::env::var_os("XDG_CACHE_HOME") {
        Some(path) => Ok(PathBuf::from(path)),
        None => {
            let home =
                std::env::var_os("HOME").context("HOME is not set; cannot locate gfwlist cache")?;
            Ok(PathBuf::from(home).join(".cache"))
        }
    }
}

fn is_cache_current(cache_path: &Path, remote_modified: SystemTime) -> Result<bool> {
    let metadata = match fs::metadata(cache_path) {
        Ok(metadata) => metadata,
        Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(false),
        Err(err) => {
            return Err(err)
                .with_context(|| format!("failed to read gfwlist cache {}", cache_path.display()));
        }
    };
    let cache_modified = metadata.modified().with_context(|| {
        format!(
            "failed to read gfwlist cache timestamp {}",
            cache_path.display()
        )
    })?;

    Ok(system_times_match_to_second(
        cache_modified,
        remote_modified,
    ))
}

fn write_gfwlist_cache(
    cache_path: &Path,
    body: &[u8],
    remote_modified: Option<SystemTime>,
) -> Result<()> {
    if let Some(parent) = cache_path.parent() {
        fs::create_dir_all(parent).with_context(|| {
            format!(
                "failed to create gfwlist cache directory {}",
                parent.display()
            )
        })?;
    }

    fs::write(cache_path, body)
        .with_context(|| format!("failed to write gfwlist cache {}", cache_path.display()))?;

    if let Some(remote_modified) = remote_modified {
        fs::File::options()
            .write(true)
            .open(cache_path)
            .and_then(|file| file.set_times(FileTimes::new().set_modified(remote_modified)))
            .with_context(|| {
                format!(
                    "failed to update gfwlist cache timestamp {}",
                    cache_path.display()
                )
            })?;
    }

    info!(
        cache_path = %cache_path.display(),
        url = GFWLIST_URL,
        "updated gfwlist cache"
    );
    Ok(())
}

fn system_times_match_to_second(left: SystemTime, right: SystemTime) -> bool {
    left.duration_since(SystemTime::UNIX_EPOCH)
        .ok()
        .map(truncate_to_second)
        == right
            .duration_since(SystemTime::UNIX_EPOCH)
            .ok()
            .map(truncate_to_second)
}

fn truncate_to_second(duration: Duration) -> Duration {
    Duration::from_secs(duration.as_secs())
}

#[derive(Debug, Clone)]
pub(crate) struct DomainRules {
    domains: HashSet<String>,
}

impl DomainRules {
    fn new(domains: HashSet<String>) -> Result<Self> {
        if domains.is_empty() {
            bail!("gfwlist did not contain any usable domain rules");
        }

        Ok(Self { domains })
    }

    fn len(&self) -> usize {
        self.domains.len()
    }

    fn extend(&mut self, domains: HashSet<String>) {
        self.domains.extend(domains);
    }

    fn matches(&self, host: &str) -> bool {
        let host = normalize_host_for_match(host);
        if host.is_empty() {
            return false;
        }

        if self.domains.contains(&host) {
            return true;
        }

        host.match_indices('.')
            .any(|(idx, _)| self.domains.contains(&host[idx + 1..]))
    }
}

fn parse_gfwlist(encoded: &[u8]) -> Result<DomainRules> {
    let compact: Vec<u8> = encoded
        .iter()
        .copied()
        .filter(|byte| !byte.is_ascii_whitespace())
        .collect();
    let decoded = STANDARD
        .decode(compact)
        .context("failed to decode gfwlist")?;
    let text = String::from_utf8(decoded).context("decoded gfwlist is not valid UTF-8")?;
    parse_gfwlist_text(&text)
}

fn parse_gfwlist_text(text: &str) -> Result<DomainRules> {
    let domains = text
        .lines()
        .filter_map(parse_proxy_rule_domain)
        .collect::<HashSet<_>>();

    DomainRules::new(domains)
}

fn read_custom_domain_rules(path: &Path) -> Result<HashSet<String>> {
    let text = fs::read_to_string(path)
        .with_context(|| format!("failed to read custom domain rules {}", path.display()))?;
    Ok(parse_custom_domain_rules_text(&text))
}

fn parse_custom_domain_rules_text(text: &str) -> HashSet<String> {
    text.lines()
        .filter_map(parse_custom_domain_rule_domain)
        .collect()
}

fn parse_custom_domain_rule_domain(line: &str) -> Option<String> {
    let rule = line.split('#').next().unwrap_or_default().trim();
    let domain = normalize_host_for_match(rule);
    is_domain_like(&domain).then_some(domain)
}

fn parse_proxy_rule_domain(line: &str) -> Option<String> {
    let rule = line.strip_prefix("||").or_else(|| line.strip_prefix('.'))?;
    if rule.contains('*') {
        return None;
    }

    let domain = rule
        .split(['/', '^', '$'])
        .next()
        .unwrap_or_default()
        .trim_matches('.');
    let domain = normalize_host_for_match(domain);
    is_domain_like(&domain).then_some(domain)
}

fn normalize_host_for_match(host: &str) -> String {
    host.trim().trim_matches('.').to_ascii_lowercase()
}

fn is_domain_like(domain: &str) -> bool {
    if domain.is_empty() || domain.contains(':') || domain.parse::<std::net::IpAddr>().is_ok() {
        return false;
    }

    domain
        .bytes()
        .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'.')
}

pub(crate) fn host_from_authority(authority: &str) -> Result<&str> {
    if let Some(rest) = authority.strip_prefix('[') {
        return rest
            .split_once("]:")
            .map(|(host, _)| host)
            .ok_or_else(|| anyhow!("IPv6 authority must be [host]:port"));
    }

    authority
        .rsplit_once(':')
        .map(|(host, _)| host)
        .ok_or_else(|| anyhow!("authority must include :port"))
}

impl std::fmt::Display for RoutingRules {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.mode())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use base64::engine::general_purpose::STANDARD;

    #[test]
    fn parses_gfwlist_domain_rules() {
        let text = "\
! comment
||example.com
||example.net/path
||example.org^
||example.edu$third-party
||wild*.blocked.test
.leading-dot.example
|http://ignored.example
";
        let encoded = STANDARD.encode(text);
        let rules = parse_gfwlist(encoded.as_bytes()).unwrap();

        assert!(rules.matches("example.com"));
        assert!(rules.matches("www.example.com"));
        assert!(rules.matches("example.net"));
        assert!(rules.matches("a.example.org"));
        assert!(rules.matches("example.edu"));
        assert!(rules.matches("www.leading-dot.example"));
        assert!(!rules.matches("wild.blocked.test"));
        assert!(!rules.matches("ignored.example"));
    }

    #[test]
    fn matches_case_insensitively_and_on_suffix_boundary() {
        let rules = parse_gfwlist_text("||example.com\n").unwrap();

        assert!(rules.matches("WWW.Example.Com."));
        assert!(!rules.matches("badexample.com"));
    }

    #[test]
    fn parses_custom_domain_rules() {
        let domains = parse_custom_domain_rules_text(
            "\
# One Squid dstdomain entry per line.
.paypal.com
.www.paypal.com

.googleadservices.com # inline comment
127.0.0.1
bad:domain
",
        );
        let rules = DomainRules::new(domains).unwrap();

        assert!(rules.matches("paypal.com"));
        assert!(rules.matches("checkout.paypal.com"));
        assert!(rules.matches("www.paypal.com"));
        assert!(rules.matches("pagead.googleadservices.com"));
        assert!(!rules.matches("127.0.0.1"));
        assert!(!rules.matches("bad:domain"));
    }

    #[test]
    fn parses_last_modified_header() {
        let mut headers = HeaderMap::new();
        headers.insert(
            LAST_MODIFIED,
            "Wed, 21 Oct 2015 07:28:00 GMT".parse().unwrap(),
        );

        assert_eq!(
            parse_last_modified(&headers).unwrap(),
            SystemTime::UNIX_EPOCH + Duration::from_secs(1_445_412_480)
        );
    }

    #[test]
    fn compares_timestamps_to_second_precision() {
        let timestamp = SystemTime::UNIX_EPOCH + Duration::from_secs(42);

        assert!(system_times_match_to_second(
            timestamp + Duration::from_millis(900),
            timestamp
        ));
        assert!(!system_times_match_to_second(
            timestamp + Duration::from_secs(1),
            timestamp
        ));
    }

    #[test]
    fn extracts_host_from_authority() {
        assert_eq!(
            host_from_authority("example.com:443").unwrap(),
            "example.com"
        );
        assert_eq!(
            host_from_authority("[2001:db8::1]:443").unwrap(),
            "2001:db8::1"
        );
    }

    #[test]
    fn global_proxy_matches_every_host() {
        let rules = RoutingRules::GlobalProxy;

        assert!(rules.should_proxy_host("example.com"));
        assert_eq!(rules.to_string(), "global");
        assert_eq!(
            rules.describe(),
            "all domains via proxy; proxy mode is global"
        );
    }

    #[tokio::test]
    async fn global_proxy_load_skips_rule_files() {
        let rules = RoutingRules::load(
            ProxyMode::Global,
            Some(Path::new("/definitely/missing/custom-domains.txt")),
        )
        .await;

        assert!(matches!(rules, RoutingRules::GlobalProxy));
    }
}