telemt/src/transport/middle_proxy/config_updater.rs

use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::Duration;

use httpdate;
use tokio::sync::watch;
use tracing::{debug, info, warn};

use crate::config::ProxyConfig;
use crate::error::Result;

use super::MePool;
use super::secret::download_proxy_secret;
use crate::crypto::SecureRandom;
use std::time::SystemTime;

async fn retry_fetch(url: &str) -> Option<ProxyConfigData> {
    let delays = [1u64, 5, 15];
    for (i, d) in delays.iter().enumerate() {
        match fetch_proxy_config(url).await {
            Ok(cfg) => return Some(cfg),
            Err(e) => {
                if i == delays.len() - 1 {
                    warn!(error = %e, url, "fetch_proxy_config failed");
                } else {
                    debug!(error = %e, url, "fetch_proxy_config retrying");
                    tokio::time::sleep(Duration::from_secs(*d)).await;
                }
            }
        }
    }
    None
}

#[derive(Debug, Clone, Default)]
pub struct ProxyConfigData {
    pub map: HashMap<i32, Vec<(IpAddr, u16)>>,
    pub default_dc: Option<i32>,
}

fn parse_host_port(s: &str) -> Option<(IpAddr, u16)> {
    if let Some(bracket_end) = s.rfind(']') {
        if s.starts_with('[') && bracket_end + 1 < s.len() && s.as_bytes().get(bracket_end + 1) == Some(&b':') {
            let host = &s[1..bracket_end];
            let port_str = &s[bracket_end + 2..];
            let ip = host.parse::<IpAddr>().ok()?;
            let port = port_str.parse::<u16>().ok()?;
            return Some((ip, port));
        }
    }

    let idx = s.rfind(':')?;
    let host = &s[..idx];
    let port_str = &s[idx + 1..];
    let ip = host.parse::<IpAddr>().ok()?;
    let port = port_str.parse::<u16>().ok()?;
    Some((ip, port))
}

fn parse_proxy_line(line: &str) -> Option<(i32, IpAddr, u16)> {
    // Accepts lines like:
    // proxy_for 4 91.108.4.195:8888;
    // proxy_for 2 [2001:67c:04e8:f002::d]:80;
    // proxy_for 2 2001:67c:04e8:f002::d:80;
    let trimmed = line.trim();
    if !trimmed.starts_with("proxy_for") {
        return None;
    }
    // Capture everything between dc and trailing ';'
    let without_prefix = trimmed.trim_start_matches("proxy_for").trim();
    let mut parts = without_prefix.split_whitespace();
    let dc_str = parts.next()?;
    let rest = parts.next()?;
    let host_port = rest.trim_end_matches(';');
    let dc = dc_str.parse::<i32>().ok()?;
    let (ip, port) = parse_host_port(host_port)?;
    Some((dc, ip, port))
}

pub async fn fetch_proxy_config(url: &str) -> Result<ProxyConfigData> {
    let resp = reqwest::get(url)
        .await
        .map_err(|e| crate::error::ProxyError::Proxy(format!("fetch_proxy_config GET failed: {e}")))?
        ;

    if let Some(date) = resp.headers().get(reqwest::header::DATE) {
        if let Ok(date_str) = date.to_str() {
            if let Ok(server_time) = httpdate::parse_http_date(date_str) {
                if let Ok(skew) = SystemTime::now().duration_since(server_time).or_else(|e| {
                    server_time.duration_since(SystemTime::now()).map_err(|_| e)
                }) {
                    let skew_secs = skew.as_secs();
                    if skew_secs > 60 {
                        warn!(skew_secs, "Time skew >60s detected from fetch_proxy_config Date header");
                    } else if skew_secs > 30 {
                        warn!(skew_secs, "Time skew >30s detected from fetch_proxy_config Date header");
                    }
                }
            }
        }
    }

    let text = resp
        .text()
        .await
        .map_err(|e| crate::error::ProxyError::Proxy(format!("fetch_proxy_config read failed: {e}")))?;

    let mut map: HashMap<i32, Vec<(IpAddr, u16)>> = HashMap::new();
    for line in text.lines() {
        if let Some((dc, ip, port)) = parse_proxy_line(line) {
            map.entry(dc).or_default().push((ip, port));
        }
    }

    let default_dc = text
        .lines()
        .find_map(|l| {
            let t = l.trim();
            if let Some(rest) = t.strip_prefix("default") {
                return rest
                    .trim()
                    .trim_end_matches(';')
                    .parse::<i32>()
                    .ok();
            }
            None
        });

    Ok(ProxyConfigData { map, default_dc })
}

async fn run_update_cycle(pool: &Arc<MePool>, rng: &Arc<SecureRandom>, cfg: &ProxyConfig) {
    pool.update_runtime_reinit_policy(
        cfg.general.hardswap,
        cfg.general.me_pool_drain_ttl_secs,
        cfg.general.effective_me_pool_force_close_secs(),
        cfg.general.me_pool_min_fresh_ratio,
    );

    let mut maps_changed = false;

    // Update proxy config v4
    let cfg_v4 = retry_fetch("https://core.telegram.org/getProxyConfig").await;
    if let Some(cfg_v4) = cfg_v4 {
        let changed = pool.update_proxy_maps(cfg_v4.map.clone(), None).await;
        if let Some(dc) = cfg_v4.default_dc {
            pool.default_dc
                .store(dc, std::sync::atomic::Ordering::Relaxed);
        }
        if changed {
            maps_changed = true;
            info!("ME config updated (v4)");
        } else {
            debug!("ME config v4 unchanged");
        }
    }

    // Update proxy config v6 (optional)
    let cfg_v6 = retry_fetch("https://core.telegram.org/getProxyConfigV6").await;
    if let Some(cfg_v6) = cfg_v6 {
        let changed = pool.update_proxy_maps(HashMap::new(), Some(cfg_v6.map)).await;
        if changed {
            maps_changed = true;
            info!("ME config updated (v6)");
        } else {
            debug!("ME config v6 unchanged");
        }
    }

    if maps_changed {
        pool.zero_downtime_reinit_after_map_change(rng.as_ref())
            .await;
    }

    pool.reset_stun_state();

    // Update proxy-secret
    match download_proxy_secret().await {
        Ok(secret) => {
            if pool.update_secret(secret).await {
                info!("proxy-secret updated and pool reconnect scheduled");
            }
        }
        Err(e) => warn!(error = %e, "proxy-secret update failed"),
    }
}

pub async fn me_config_updater(
    pool: Arc<MePool>,
    rng: Arc<SecureRandom>,
    mut config_rx: watch::Receiver<Arc<ProxyConfig>>,
) {
    let mut update_every_secs = config_rx
        .borrow()
        .general
        .effective_update_every_secs()
        .max(1);
    let mut update_every = Duration::from_secs(update_every_secs);
    let mut next_tick = tokio::time::Instant::now() + update_every;
    info!(update_every_secs, "ME config updater started");

    loop {
        let sleep = tokio::time::sleep_until(next_tick);
        tokio::pin!(sleep);

        tokio::select! {
            _ = &mut sleep => {
                let cfg = config_rx.borrow().clone();
                run_update_cycle(&pool, &rng, cfg.as_ref()).await;
                let refreshed_secs = cfg.general.effective_update_every_secs().max(1);
                if refreshed_secs != update_every_secs {
                    info!(
                        old_update_every_secs = update_every_secs,
                        new_update_every_secs = refreshed_secs,
                        "ME config updater interval changed"
                    );
                    update_every_secs = refreshed_secs;
                    update_every = Duration::from_secs(update_every_secs);
                }
                next_tick = tokio::time::Instant::now() + update_every;
            }
            changed = config_rx.changed() => {
                if changed.is_err() {
                    warn!("ME config updater stopped: config channel closed");
                    break;
                }
                let cfg = config_rx.borrow().clone();
                pool.update_runtime_reinit_policy(
                    cfg.general.hardswap,
                    cfg.general.me_pool_drain_ttl_secs,
                    cfg.general.effective_me_pool_force_close_secs(),
                    cfg.general.me_pool_min_fresh_ratio,
                );
                let new_secs = cfg.general.effective_update_every_secs().max(1);
                if new_secs == update_every_secs {
                    continue;
                }

                if new_secs < update_every_secs {
                    info!(
                        old_update_every_secs = update_every_secs,
                        new_update_every_secs = new_secs,
                        "ME config updater interval decreased, running immediate refresh"
                    );
                    update_every_secs = new_secs;
                    update_every = Duration::from_secs(update_every_secs);
                    run_update_cycle(&pool, &rng, cfg.as_ref()).await;
                    next_tick = tokio::time::Instant::now() + update_every;
                } else {
                    info!(
                        old_update_every_secs = update_every_secs,
                        new_update_every_secs = new_secs,
                        "ME config updater interval increased"
                    );
                    update_every_secs = new_secs;
                    update_every = Duration::from_secs(update_every_secs);
                    next_tick = tokio::time::Instant::now() + update_every;
                }
            }
        }
    }
}

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

    #[test]
    fn parse_ipv6_bracketed() {
        let line = "proxy_for 2 [2001:67c:04e8:f002::d]:80;";
        let res = parse_proxy_line(line).unwrap();
        assert_eq!(res.0, 2);
        assert_eq!(res.1, "2001:67c:04e8:f002::d".parse::<IpAddr>().unwrap());
        assert_eq!(res.2, 80);
    }

    #[test]
    fn parse_ipv6_plain() {
        let line = "proxy_for 2 2001:67c:04e8:f002::d:80;";
        let res = parse_proxy_line(line).unwrap();
        assert_eq!(res.0, 2);
        assert_eq!(res.1, "2001:67c:04e8:f002::d".parse::<IpAddr>().unwrap());
        assert_eq!(res.2, 80);
    }

    #[test]
    fn parse_ipv4() {
        let line = "proxy_for 4 91.108.4.195:8888;";
        let res = parse_proxy_line(line).unwrap();
        assert_eq!(res.0, 4);
        assert_eq!(res.1, "91.108.4.195".parse::<IpAddr>().unwrap());
        assert_eq!(res.2, 8888);
    }
}