Refactor application structure and simplify implementation

src-tauri/src/subscription.rs

@@ -1,8 +1,7 @@
 use crate::models::{decode_percent_encoded_utf8, SubscriptionCache, SubscriptionServer};
 use base64::{engine::general_purpose, Engine};
-use reqwest::redirect;
 use serde_json::{json, Map, Value};
-use std::net::{IpAddr, Ipv6Addr};
+use std::net::{IpAddr, Ipv6Addr, SocketAddr, ToSocketAddrs};
 use std::time::Duration;
 use std::time::{SystemTime, UNIX_EPOCH};
 use url::Url;
@@ -11,6 +10,7 @@ const SUPPORTED_PROXY_TYPES: &[&str] = &["vless", "vmess", "trojan", "shadowsock
 const DEFAULT_APP_NAME: &str = "ProxyWarden";
 const SUBSCRIPTION_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
 const SUBSCRIPTION_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
+const SUBSCRIPTION_MAX_REDIRECTS: usize = 5;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct SubscriptionError {
@@ -155,69 +155,129 @@ pub fn fetch_subscription_with_identity_and_policy(
 ) -> Result<SubscriptionCache, SubscriptionError> {
     let parsed_url =
         Url::parse(url).map_err(|_| SubscriptionError::new("Invalid subscription URL"))?;
-    validate_subscription_fetch_url(&parsed_url, policy)?;
+    let mut current_url = parsed_url;
 
-    let redirect_policy = redirect::Policy::custom(move |attempt| {
-        if validate_subscription_fetch_url(attempt.url(), policy).is_ok() {
-            attempt.follow()
-        } else {
-            attempt.stop()
+    for redirect_count in 0..=SUBSCRIPTION_MAX_REDIRECTS {
+        validate_subscription_fetch_url(&current_url, policy)?;
+        let client = subscription_client_for_url(&current_url, policy)?;
+        let mut request = client.get(current_url.clone());
+
+        for (name, value) in identity.request_headers_without_device_hwid() {
+            request = request.header(name, value);
         }
-    });
-    let client = reqwest::blocking::Client::builder()
+        if let Some(device_hwid) = identity
+            .device_hwid
+            .as_deref()
+            .map(str::trim)
+            .filter(|value| !value.is_empty())
+        {
+            request = request.header("x-hwid", device_hwid);
+        }
+
+        let response = request.send().map_err(|error| {
+            SubscriptionError::new(format!("Subscription request failed: {error}"))
+        })?;
+        let status = response.status();
+        if status.is_redirection() {
+            if redirect_count == SUBSCRIPTION_MAX_REDIRECTS {
+                return Err(SubscriptionError::new(
+                    "Subscription request exceeded redirect limit",
+                ));
+            }
+            let location = response
+                .headers()
+                .get(reqwest::header::LOCATION)
+                .and_then(|value| value.to_str().ok())
+                .ok_or_else(|| {
+                    SubscriptionError::new("Subscription redirect has no valid Location header")
+                })?;
+            current_url = current_url
+                .join(location)
+                .map_err(|_| SubscriptionError::new("Subscription redirect URL is invalid"))?;
+            continue;
+        }
+        if !status.is_success() {
+            return Err(SubscriptionError::new(format!(
+                "Subscription request failed: HTTP {}",
+                status.as_u16()
+            )));
+        }
+
+        let user_info = parse_user_info(
+            response
+                .headers()
+                .get("subscription-userinfo")
+                .and_then(|value| value.to_str().ok()),
+        );
+        let body = response.text().map_err(|error| {
+            SubscriptionError::new(format!("Subscription body read failed: {error}"))
+        })?;
+        let parsed = parse_subscription_body(&body)?;
+
+        return Ok(SubscriptionCache {
+            config: parsed.config,
+            servers: parsed.servers,
+            user_info,
+            fetched_at: now_timestamp(),
+        });
+    }
+
+    Err(SubscriptionError::new(
+        "Subscription request could not complete",
+    ))
+}
+
+fn subscription_client_for_url(
+    parsed_url: &Url,
+    policy: SubscriptionFetchPolicy,
+) -> Result<reqwest::blocking::Client, SubscriptionError> {
+    let mut builder = reqwest::blocking::Client::builder()
         .connect_timeout(SUBSCRIPTION_CONNECT_TIMEOUT)
         .timeout(SUBSCRIPTION_REQUEST_TIMEOUT)
-        .redirect(redirect_policy)
-        .build()
-        .map_err(|error| {
-            SubscriptionError::new(format!("Subscription client setup failed: {error}"))
-        })?;
-    let mut request = client.get(parsed_url);
+        .redirect(reqwest::redirect::Policy::none());
 
-    for (name, value) in identity.request_headers_without_device_hwid() {
-        request = request.header(name, value);
+    if !policy.allow_unsafe_local_urls {
+        let host = parsed_url
+            .host_str()
+            .ok_or_else(|| SubscriptionError::new("Subscription URL has no host"))?;
+        if host.parse::<IpAddr>().is_err() {
+            let port = parsed_url
+                .port_or_known_default()
+                .ok_or_else(|| SubscriptionError::new("Subscription URL has no resolvable port"))?;
+            let addresses = (host, port)
+                .to_socket_addrs()
+                .map_err(|error| {
+                    SubscriptionError::new(format!(
+                        "Subscription host DNS resolution failed: {error}"
+                    ))
+                })?
+                .collect::<Vec<_>>();
+            validate_resolved_subscription_addresses(&addresses)?;
+            builder = builder.resolve_to_addrs(host, &addresses);
+        }
     }
 
-    if let Some(device_hwid) = identity
-        .device_hwid
-        .as_deref()
-        .map(str::trim)
-        .filter(|value| !value.is_empty())
-    {
-        request = request.header("x-hwid", device_hwid);
-    }
-
-    let response = request
-        .send()
-        .map_err(|error| SubscriptionError::new(format!("Subscription request failed: {error}")))?;
-
-    let status = response.status();
-    if !status.is_success() {
-        return Err(SubscriptionError::new(format!(
-            "Subscription request failed: HTTP {}",
-            status.as_u16()
-        )));
-    }
-
-    let user_info = parse_user_info(
-        response
-            .headers()
-            .get("subscription-userinfo")
-            .and_then(|value| value.to_str().ok()),
-    );
-    let body = response.text().map_err(|error| {
-        SubscriptionError::new(format!("Subscription body read failed: {error}"))
-    })?;
-    let parsed = parse_subscription_body(&body)?;
-
-    Ok(SubscriptionCache {
-        config: parsed.config,
-        servers: parsed.servers,
-        user_info,
-        fetched_at: now_timestamp(),
+    builder.build().map_err(|error| {
+        SubscriptionError::new(format!("Subscription client setup failed: {error}"))
     })
 }
 
+pub fn validate_resolved_subscription_addresses(
+    addresses: &[SocketAddr],
+) -> Result<(), SubscriptionError> {
+    if addresses.is_empty() {
+        return Err(SubscriptionError::new(
+            "Subscription host DNS resolution returned no addresses",
+        ));
+    }
+    if addresses.iter().any(|address| is_unsafe_ip(address.ip())) {
+        return Err(SubscriptionError::new(
+            "Subscription host resolves to a local, private, link-local, multicast, or metadata address",
+        ));
+    }
+    Ok(())
+}
+
 fn validate_subscription_fetch_url(
     parsed_url: &Url,
     policy: SubscriptionFetchPolicy,
@@ -286,23 +346,171 @@ fn parse_link_subscription(body: &str) -> Result<Value, SubscriptionError> {
     let links = decoded
         .lines()
         .map(str::trim)
-        .filter(|line| line.starts_with("vless://"))
+        .filter(|line| {
+            ["vless://", "trojan://", "ss://", "vmess://"]
+                .iter()
+                .any(|scheme| line.starts_with(scheme))
+        })
         .collect::<Vec<_>>();
 
     if links.is_empty() {
         return Err(SubscriptionError::new(
-            "Subscription does not contain JSON config or VLESS links",
+            "Subscription does not contain JSON config or supported VLESS, VMess, Trojan, or Shadowsocks links",
         ));
     }
 
     let outbounds = links
         .into_iter()
-        .map(parse_vless_url)
+        .map(|link| {
+            if link.starts_with("vless://") {
+                parse_vless_url(link)
+            } else if link.starts_with("trojan://") {
+                parse_trojan_url(link)
+            } else if link.starts_with("ss://") {
+                parse_shadowsocks_url(link)
+            } else {
+                parse_vmess_url(link)
+            }
+        })
         .collect::<Result<Vec<_>, _>>()?;
 
     Ok(json!({ "outbounds": outbounds }))
 }
 
+fn parse_trojan_url(raw_url: &str) -> Result<Value, SubscriptionError> {
+    let parsed = Url::parse(raw_url).map_err(|_| SubscriptionError::new("Invalid Trojan URL"))?;
+    let password = parsed.username().trim().to_string();
+    let server = parsed.host_str().map(str::to_string).unwrap_or_default();
+    let server_port = parsed.port_or_known_default().unwrap_or(443);
+    if password.is_empty() || server.is_empty() {
+        return Err(SubscriptionError::new(
+            "Trojan URL misses password, host or port",
+        ));
+    }
+    let tag = parsed
+        .fragment()
+        .map(decode_percent_encoded_utf8)
+        .unwrap_or_else(|| "trojan-out".to_string());
+    let server_name = query_value(&parsed, "sni").unwrap_or_else(|| server.clone());
+
+    Ok(json!({
+        "type": "trojan",
+        "tag": tag,
+        "server": server,
+        "server_port": server_port,
+        "password": password,
+        "tls": {
+            "enabled": true,
+            "server_name": server_name
+        }
+    }))
+}
+
+fn parse_shadowsocks_url(raw_url: &str) -> Result<Value, SubscriptionError> {
+    let parsed =
+        Url::parse(raw_url).map_err(|_| SubscriptionError::new("Invalid Shadowsocks URL"))?;
+    let server = parsed.host_str().map(str::to_string).unwrap_or_default();
+    let server_port = parsed.port().unwrap_or(8388);
+    let credentials = match parsed.password() {
+        Some(password) => format!("{}:{password}", parsed.username()),
+        None => decode_base64_text(parsed.username()).ok_or_else(|| {
+            SubscriptionError::new("Shadowsocks credentials are not valid base64")
+        })?,
+    };
+    let (method, password) = credentials
+        .split_once(':')
+        .ok_or_else(|| SubscriptionError::new("Shadowsocks URL misses method or password"))?;
+    if method.trim().is_empty() || password.is_empty() || server.is_empty() {
+        return Err(SubscriptionError::new(
+            "Shadowsocks URL misses method, password, host or port",
+        ));
+    }
+    let tag = parsed
+        .fragment()
+        .map(decode_percent_encoded_utf8)
+        .unwrap_or_else(|| "shadowsocks-out".to_string());
+
+    Ok(json!({
+        "type": "shadowsocks",
+        "tag": tag,
+        "server": server,
+        "server_port": server_port,
+        "method": method,
+        "password": password
+    }))
+}
+
+fn parse_vmess_url(raw_url: &str) -> Result<Value, SubscriptionError> {
+    let payload = raw_url
+        .strip_prefix("vmess://")
+        .and_then(|value| value.split('#').next())
+        .ok_or_else(|| SubscriptionError::new("Invalid VMess URL"))?;
+    let decoded = decode_base64_text(payload)
+        .ok_or_else(|| SubscriptionError::new("VMess payload is not valid base64"))?;
+    let source: Value = serde_json::from_str(&decoded)
+        .map_err(|_| SubscriptionError::new("VMess payload is not valid JSON"))?;
+    let server = source
+        .get("add")
+        .and_then(Value::as_str)
+        .unwrap_or_default();
+    let server_port = source
+        .get("port")
+        .and_then(|value| value.as_u64().or_else(|| value.as_str()?.parse().ok()))
+        .and_then(|value| u16::try_from(value).ok())
+        .unwrap_or(443);
+    let uuid = source.get("id").and_then(Value::as_str).unwrap_or_default();
+    if server.is_empty() || uuid.is_empty() {
+        return Err(SubscriptionError::new(
+            "VMess payload misses host, port or uuid",
+        ));
+    }
+    let tag = source
+        .get("ps")
+        .and_then(Value::as_str)
+        .map(decode_percent_encoded_utf8)
+        .unwrap_or_else(|| "vmess-out".to_string());
+    let security = source
+        .get("scy")
+        .and_then(Value::as_str)
+        .filter(|value| !value.is_empty())
+        .unwrap_or("auto");
+    let mut outbound = json!({
+        "type": "vmess",
+        "tag": tag,
+        "server": server,
+        "server_port": server_port,
+        "uuid": uuid,
+        "security": security
+    });
+    if source.get("tls").and_then(Value::as_str) == Some("tls") {
+        let server_name = source
+            .get("sni")
+            .or_else(|| source.get("host"))
+            .and_then(Value::as_str)
+            .filter(|value| !value.is_empty())
+            .unwrap_or(server);
+        outbound["tls"] = json!({ "enabled": true, "server_name": server_name });
+    }
+    if source.get("net").and_then(Value::as_str) == Some("ws") {
+        let path = source
+            .get("path")
+            .and_then(Value::as_str)
+            .filter(|value| !value.is_empty())
+            .unwrap_or("/");
+        let host = source
+            .get("host")
+            .and_then(Value::as_str)
+            .filter(|value| !value.is_empty());
+        outbound["transport"] = json!({
+            "type": "ws",
+            "path": path,
+            "headers": host.map(|host| json!({ "Host": host })).unwrap_or_else(|| json!({}))
+        });
+    }
+
+    Ok(outbound)
+}
+
 fn parse_vless_url(raw_url: &str) -> Result<Value, SubscriptionError> {
     if !raw_url.starts_with("vless://") {
         return Err(SubscriptionError::new("VLESS URL must start with vless://"));
@@ -399,6 +607,7 @@ fn server_from_outbound(outbound: &Value) -> Option<SubscriptionServer> {
         .unwrap_or_else(|| format!("{server_type}-{server}"));
 
     Some(SubscriptionServer {
+        id: outbound_server_id(outbound),
         tag,
         server_type,
         server,
@@ -406,6 +615,14 @@ fn server_from_outbound(outbound: &Value) -> Option<SubscriptionServer> {
     })
 }
 
+fn outbound_server_id(outbound: &Value) -> String {
+    let bytes = serde_json::to_vec(outbound).unwrap_or_default();
+    let hash = bytes.iter().fold(0xcbf29ce484222325_u64, |hash, byte| {
+        (hash ^ u64::from(*byte)).wrapping_mul(0x100000001b3)
+    });
+    format!("pw-{hash:016x}")
+}
+
 fn maybe_decode_base64(content: &str) -> String {
     let compact = content.split_whitespace().collect::<String>();
     if compact.is_empty()
@@ -419,7 +636,11 @@ fn maybe_decode_base64(content: &str) -> String {
     for engine in [general_purpose::STANDARD, general_purpose::URL_SAFE] {
         if let Ok(decoded) = engine.decode(compact.as_bytes()) {
             if let Ok(decoded) = String::from_utf8(decoded) {
-                if decoded.contains("vless://") || decoded.contains('{') {
+                if ["vless://", "vmess://", "trojan://", "ss://"]
+                    .iter()
+                    .any(|scheme| decoded.contains(scheme))
+                    || decoded.contains('{')
+                {
                     return decoded;
                 }
             }
@@ -429,6 +650,23 @@ fn maybe_decode_base64(content: &str) -> String {
     content.to_string()
 }
 
+fn decode_base64_text(value: &str) -> Option<String> {
+    let value = value.trim();
+    for engine in [
+        general_purpose::STANDARD,
+        general_purpose::STANDARD_NO_PAD,
+        general_purpose::URL_SAFE,
+        general_purpose::URL_SAFE_NO_PAD,
+    ] {
+        if let Ok(decoded) = engine.decode(value.as_bytes()) {
+            if let Ok(decoded) = String::from_utf8(decoded) {
+                return Some(decoded);
+            }
+        }
+    }
+    None
+}
+
 fn query_value(url: &Url, key: &str) -> Option<String> {
     url.query_pairs()
         .find(|(name, _)| name == key)