path: root/src/subscription.rs

use crate::model::{Entry, EntryStatus, Trojan, display_text};
use anyhow::{Result, anyhow, bail, ensure};
use async_trait::async_trait;
use base64::{Engine, engine::general_purpose};
use percent_encoding::percent_decode_str;
use std::time::Duration;
use url::Url;
use uuid::Uuid;

pub const MAX_SUBSCRIPTION_BYTES: usize = 4 * 1024 * 1024;

#[async_trait]
pub trait SubscriptionSource: Send {
    async fn fetch(&self, url: &str) -> Result<Vec<u8>>;
}

pub struct HttpSource;

pub fn subscription_url(value: &str) -> Result<Url> {
    let url = Url::parse(value).map_err(|_| anyhow!("Invalid subscription URL"))?;
    ensure!(
        matches!(url.scheme(), "http" | "https")
            && url.host_str().is_some()
            && url.username().is_empty()
            && url.password().is_none()
            && url.fragment().is_none(),
        "Use an HTTP or HTTPS subscription URL without userinfo or a fragment"
    );
    Ok(url)
}

#[async_trait]
impl SubscriptionSource for HttpSource {
    async fn fetch(&self, value: &str) -> Result<Vec<u8>> {
        let mut url = subscription_url(value)?;
        let client = reqwest::Client::builder()
            .no_proxy()
            .redirect(reqwest::redirect::Policy::none())
            .timeout(Duration::from_secs(30))
            .connect_timeout(Duration::from_secs(10))
            .user_agent(concat!("Roentgen/", env!("CARGO_PKG_VERSION")))
            .build()
            .map_err(|_| anyhow!("Could not initialize subscription TLS client"))?;
        let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
        for redirects in 0..=5 {
            let mut response = tokio::time::timeout_at(
                deadline,
                client
                    .get(url.clone())
                    .header("Accept", "text/plain")
                    .send(),
            )
            .await
            .map_err(|_| anyhow!("Subscription request timed out"))?
            .map_err(|_| anyhow!("Subscription request failed; check connectivity and TLS"))?;
            if response.status().is_redirection() {
                ensure!(
                    matches!(response.status().as_u16(), 301 | 302 | 303 | 307 | 308),
                    "Unexpected subscription redirect response"
                );
                ensure!(redirects < 5, "Subscription exceeded five redirects");
                let location = response
                    .headers()
                    .get(reqwest::header::LOCATION)
                    .and_then(|v| v.to_str().ok())
                    .ok_or_else(|| anyhow!("Subscription redirect has no valid location"))?;
                let next = url
                    .join(location)
                    .map_err(|_| anyhow!("Invalid subscription redirect"))?;
                subscription_url(next.as_str())?;
                ensure!(
                    !(url.scheme() == "https" && next.scheme() != "https"),
                    "Refusing HTTPS-to-HTTP subscription redirect"
                );
                url = next;
                continue;
            }
            ensure!(
                response.status() == reqwest::StatusCode::OK,
                "Subscription returned HTTP {}",
                response.status().as_u16()
            );
            if let Some(kind) = response
                .headers()
                .get(reqwest::header::CONTENT_TYPE)
                .and_then(|h| h.to_str().ok())
            {
                ensure!(
                    !["html", "json", "yaml"]
                        .iter()
                        .any(|v| kind.to_ascii_lowercase().contains(v)),
                    "Subscription returned an unsupported document format; use a plain or Base64 link list"
                );
            }
            let mut body = Vec::new();
            while let Some(chunk) = tokio::time::timeout_at(deadline, response.chunk())
                .await
                .map_err(|_| anyhow!("Subscription body timed out"))?
                .map_err(|_| anyhow!("Subscription body download failed"))?
            {
                ensure!(
                    body.len() + chunk.len() <= MAX_SUBSCRIPTION_BYTES,
                    "Subscription exceeds the 4 MiB limit"
                );
                body.extend_from_slice(&chunk);
            }
            return Ok(body);
        }
        unreachable!()
    }
}

pub fn import(bytes: &[u8]) -> Result<Vec<Entry>> {
    ensure!(
        bytes.len() <= MAX_SUBSCRIPTION_BYTES,
        "Subscription exceeds the 4 MiB limit"
    );
    let text = std::str::from_utf8(bytes)
        .map_err(|_| anyhow!("Subscription is not UTF-8"))?
        .trim_start_matches('\u{feff}')
        .trim();
    if text.is_empty() {
        return Ok(vec![]);
    }
    let decoded;
    let plain =
        if text.contains("://") || text.contains('\n') && text.lines().any(|l| l.contains("://")) {
            text
        } else {
            let compact: String = text.chars().filter(|c| !c.is_whitespace()).collect();
            let mut result = None;
            for engine in [
                general_purpose::STANDARD,
                general_purpose::STANDARD_NO_PAD,
                general_purpose::URL_SAFE,
                general_purpose::URL_SAFE_NO_PAD,
            ] {
                if let Ok(value) = engine.decode(&compact) {
                    result = Some(value);
                    break;
                }
            }
            decoded = result
                .ok_or_else(|| anyhow!("Subscription is neither a link list nor valid Base64"))?;
            ensure!(
                decoded.len() <= MAX_SUBSCRIPTION_BYTES,
                "Decoded subscription exceeds the limit"
            );
            std::str::from_utf8(&decoded)
                .map_err(|_| anyhow!("Decoded subscription is not UTF-8"))?
                .trim()
        };
    let lower = plain.to_ascii_lowercase();
    ensure!(
        !plain.starts_with(['<', '{', '['])
            && !lower.starts_with("proxies:")
            && !lower.starts_with("proxy-providers:"),
        "Unsupported subscription envelope; use a plain or Base64 link list"
    );
    // A readable empty envelope is accepted, but arbitrary prose is not an entry list.
    ensure!(
        plain.is_empty() || plain.lines().any(|line| line.contains("://")),
        "Subscription does not contain a link list"
    );
    Ok(plain
        .lines()
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .map(parse_entry)
        .collect())
}

fn decode(value: &str) -> Result<String> {
    let bytes = value.as_bytes();
    for (index, byte) in bytes.iter().enumerate() {
        if *byte == b'%' {
            ensure!(
                bytes.get(index + 1).is_some_and(u8::is_ascii_hexdigit)
                    && bytes.get(index + 2).is_some_and(u8::is_ascii_hexdigit),
                "Invalid percent encoding"
            );
        }
    }
    percent_decode_str(value)
        .decode_utf8()
        .map(|v| v.into_owned())
        .map_err(|_| anyhow!("Invalid UTF-8 in URI component"))
}

fn parse_entry(raw: &str) -> Entry {
    let mut entry = Entry {
        id: Uuid::new_v4().to_string(),
        name: "Invalid entry".into(),
        status: EntryStatus::Invalid,
        explanation: None,
        connection: None,
    };
    let result = (|| -> Result<Trojan> {
        ensure!(
            !raw.chars().any(char::is_control),
            "Control characters in connection link"
        );
        let url = Url::parse(raw).map_err(|_| anyhow!("Malformed connection link"))?;
        if let Some(label) = url.fragment() {
            entry.name = display_text(&decode(label)?);
        }
        if url.scheme() != "trojan" {
            entry.status = EntryStatus::Unsupported;
            bail!("Only Trojan TCP + TLS connections are supported");
        }
        ensure!(
            url.password().is_none(),
            "Trojan credentials must be a percent-encoded password"
        );
        let password = decode(url.username())?;
        ensure!(!password.is_empty(), "Missing Trojan password");
        let host = match url
            .host()
            .ok_or_else(|| anyhow!("Missing connection host"))?
        {
            url::Host::Domain(v) => v.to_ascii_lowercase(),
            url::Host::Ipv4(v) => v.to_string(),
            url::Host::Ipv6(v) => v.to_string(),
        };
        let port = url
            .port()
            .ok_or_else(|| anyhow!("Missing connection port"))?;
        ensure!(port != 0, "Connection port must be positive");
        if entry.name == "Invalid entry" || entry.name.is_empty() {
            entry.name = display_text(&host);
        }
        let mut server_name = host.clone();
        let mut alpn = vec![];
        let mut fingerprint = None;
        let mut seen = std::collections::HashSet::new();
        for (key, value) in url.query_pairs() {
            ensure!(seen.insert(key.to_string()), "Duplicate connection option");
            match key.as_ref() {
                "type" if matches!(value.as_ref(), "tcp" | "raw") => {}
                "security" if value == "tls" => {}
                "sni" | "peer" => {
                    ensure!(
                        !value.is_empty() && !value.chars().any(char::is_control),
                        "Invalid TLS server name"
                    );
                    server_name = value.into_owned();
                }
                "alpn" => {
                    alpn = value.split(',').map(str::to_owned).collect();
                    ensure!(
                        alpn.iter().all(|v| !v.is_empty()
                            && v.len() <= 255
                            && !v.chars().any(char::is_control)),
                        "Invalid TLS ALPN list"
                    );
                }
                "fp" => {
                    ensure!(
                        matches!(
                            value.as_ref(),
                            "chrome"
                                | "firefox"
                                | "safari"
                                | "ios"
                                | "android"
                                | "edge"
                                | "360"
                                | "qq"
                                | "random"
                                | "randomized"
                        ),
                        "Unsupported TLS fingerprint"
                    );
                    fingerprint = Some(value.into_owned());
                }
                "allowInsecure" if value == "0" || value == "false" => {}
                "headerType" if value == "none" => {}
                _ => {
                    entry.status = EntryStatus::Unsupported;
                    bail!("Unsupported transport, TLS, or connection option");
                }
            }
        }
        ensure!(
            !(seen.contains("sni") && seen.contains("peer")),
            "Specify only one TLS server name option"
        );
        Ok(Trojan {
            host,
            port,
            password,
            server_name,
            alpn,
            fingerprint,
        })
    })();
    match result {
        Ok(connection) => {
            entry.status = EntryStatus::Supported;
            entry.connection = Some(connection);
        }
        Err(error) => entry.explanation = Some(error.to_string()),
    }
    entry
}