EasyTier/easytier/src/dns/system/windows.rs

use std::net::IpAddr;
use std::process::Command;

use std::io;
use winreg::RegKey;

use crate::common::ifcfg::RegistryManager;

use super::{SystemConfig, SystemConfigurator};

pub fn is_windows_10_or_better() -> io::Result<bool> {
    let hklm = winreg::enums::HKEY_LOCAL_MACHINE;
    let key_path = "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion";
    let key = winreg::RegKey::predef(hklm).open_subkey(key_path)?;

    // check CurrentMajorVersionNumber, which only exists on Windows 10 and later
    let value_name = "CurrentMajorVersionNumber";
    key.get_raw_value(value_name).map(|_| true)
}

// 假设 interface_guid 是你的网络接口 GUID
#[derive(Clone)]
pub struct InterfaceControl {
    interface_guid: String,
}

impl InterfaceControl {
    // 构造函数
    pub fn new(interface_guid: &str) -> Self {
        InterfaceControl {
            interface_guid: interface_guid.to_string(),
        }
    }

    // 删除注册表值（模拟 delValue）
    fn delete_value(key: &RegKey, value_name: &str) -> io::Result<()> {
        match key.delete_value(value_name) {
            Ok(_) => Ok(()),
            Err(e) => {
                if matches!(e.kind(), io::ErrorKind::NotFound) {
                    Ok(()) // 忽略不存在的值
                } else {
                    Err(e)
                }
            }
        }
    }

    pub fn set_primary_dns(&self, resolvers: &[IpAddr], domains: &[String]) -> io::Result<()> {
        let (ipsv4, ipsv6): (Vec<String>, Vec<String>) = resolvers
            .iter()
            .map(|ip| ip.to_string())
            .partition(|ip| ip.contains('.'));

        let dom_strs: Vec<String> = domains
            .iter()
            .map(|d| d.trim_end_matches('.').to_string())
            .collect();

        // IPv4 处理
        if let Ok(key4) = RegistryManager::open_interface_key(
            &self.interface_guid,
            RegistryManager::IPV4_TCPIP_INTERFACE_PREFIX,
        ) {
            if ipsv4.is_empty() {
                Self::delete_value(&key4, "NameServer")?;
            } else {
                key4.set_value("NameServer", &ipsv4.join(","))?;
            }

            if dom_strs.is_empty() {
                Self::delete_value(&key4, "SearchList")?;
            } else {
                key4.set_value("SearchList", &dom_strs.join(","))?;
            }

            // 禁用 LLMNR（通过 DisableMulticast）
            key4.set_value("EnableMulticast", &0u32)?;
        }

        // IPv6 处理
        if let Ok(key6) = RegistryManager::open_interface_key(
            &self.interface_guid,
            RegistryManager::IPV6_TCPIP_INTERFACE_PREFIX,
        ) {
            if ipsv6.is_empty() {
                Self::delete_value(&key6, "NameServer")?;
            } else {
                key6.set_value("NameServer", &ipsv6.join(","))?;
            }

            if dom_strs.is_empty() {
                Self::delete_value(&key6, "SearchList")?;
            } else {
                key6.set_value("SearchList", &dom_strs.join(","))?;
            }
            key6.set_value("EnableMulticast", &0u32)?;
        }

        Ok(())
    }

    fn flush_dns(&self) -> io::Result<()> {
        // 刷新 DNS 缓存
        let output = Command::new("ipconfig")
            .arg("/flushdns")
            .output()
            .expect("failed to execute process");
        if !output.status.success() {
            return Err(io::Error::other("Failed to flush DNS cache"));
        }
        Ok(())
    }

    // re-register DNS
    pub fn re_register_dns(&self) -> io::Result<()> {
        // ipconfig /registerdns
        let output = Command::new("ipconfig")
            .arg("/registerdns")
            .output()
            .expect("failed to execute process");
        if !output.status.success() {
            return Err(io::Error::other("Failed to register DNS"));
        }
        Ok(())
    }
}

#[derive(Clone)]
pub struct WindowsDNSManager {
    tun_dev_name: String,
    interface_control: InterfaceControl,
}

impl WindowsDNSManager {
    pub fn new(tun_dev_name: &str) -> io::Result<Self> {
        let interface_guid = RegistryManager::find_interface_guid(tun_dev_name)?;
        Ok(WindowsDNSManager {
            tun_dev_name: tun_dev_name.to_string(),
            interface_control: InterfaceControl::new(&interface_guid),
        })
    }

    pub fn set_primary_dns(&self, resolvers: &[IpAddr], domains: &[String]) -> io::Result<()> {
        self.interface_control.set_primary_dns(resolvers, domains)?;
        self.interface_control.flush_dns()?;
        Ok(())
    }
}

impl SystemConfigurator for WindowsDNSManager {
    fn set_dns(&self, cfg: &SystemConfig) -> io::Result<()> {
        self.set_primary_dns(
            &cfg.nameservers
                .iter()
                .map(|s| s.parse::<IpAddr>().unwrap())
                .collect::<Vec<_>>(),
            &cfg.match_domains,
        )?;
        Ok(())
    }

    fn clean(&self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(all(test, target_os = "windows", feature = "magic-dns", feature = "tun"))]
mod tests {
    use cidr::Ipv4Inet;
    use std::net::IpAddr;

    #[tokio::test]
    async fn test_windows_set_primary_server() {
        use std::{str::FromStr as _, sync::Arc, time::Duration};

        use crate::dns::{
            config::DNS_DEFAULT_ADDRESSES,
            tests::{prepare_env, start_dns_node},
        };
        use crate::instance::proxy_cidrs_monitor::ProxyCidrsMonitor;
        use crate::instance::virtual_nic::NicCtx;
        use crate::peers::peer_manager::PeerManager;

        let tun_ip = Ipv4Inet::from_str("10.144.144.10/24").unwrap();
        let (peer_mgr, virtual_nic): (Arc<PeerManager>, NicCtx) =
            prepare_env("test1", tun_ip).await;
        let tun_name = virtual_nic.ifname().await.unwrap();

        // prepare_env does not run full Instance::run, so start the monitor explicitly in test.
        let _monitor = ProxyCidrsMonitor::new(peer_mgr.clone(), peer_mgr.get_global_ctx()).start();

        let mut dns_node = start_dns_node(peer_mgr, virtual_nic);

        println!("dev_name: {}", tun_name);
        let fake_ip = match DNS_DEFAULT_ADDRESSES[0].addr.ip() {
            IpAddr::V4(ip) => ip,
            IpAddr::V6(ip) => panic!("unexpected ipv6 default dns address in test: {ip}"),
        };

        // Windows may take a while to attach the test IP; wait until ping succeeds.
        let now = std::time::Instant::now();
        let mut ping_ready = false;
        while now.elapsed() < Duration::from_secs(5) {
            tokio::time::sleep(Duration::from_secs(1)).await;
            if let Ok(o) = tokio::process::Command::new("ping")
                .arg("-n")
                .arg("1")
                .arg("-w")
                .arg("100")
                .arg(fake_ip.to_string())
                .output()
                .await
                && o.status.success()
            {
                ping_ready = true;
                break;
            }
        }
        if !ping_ready {
            tracing::warn!(
                "dns test endpoint {} did not respond to ping in time; continue with dns checks",
                fake_ip
            );
        }

        // First verify the DNS node can answer queries when explicitly targeting it.
        let direct_lookup_deadline = std::time::Instant::now() + Duration::from_secs(3);
        let mut direct_lookup_ok = false;
        let mut last_direct_output = String::new();

        while std::time::Instant::now() < direct_lookup_deadline {
            let result = tokio::time::timeout(
                Duration::from_secs(3),
                tokio::process::Command::new("nslookup")
                    .arg("test1.et.net")
                    .arg(fake_ip.to_string())
                    .output(),
            )
            .await;

            if let Ok(Ok(ret)) = result {
                let output = String::from_utf8_lossy(&ret.stdout).to_string();
                println!("direct nslookup output: {}", output);
                if ret.status.success() && output.contains("10.144.144.10") {
                    direct_lookup_ok = true;
                    break;
                }
                last_direct_output = output;
            }

            tokio::time::sleep(Duration::from_millis(500)).await;
        }

        assert!(
            direct_lookup_ok,
            "direct nslookup against dns server did not return expected IP in time; last output: {last_direct_output}"
        );

        let dns_mgr = super::WindowsDNSManager::new(&tun_name).unwrap();
        println!("dev_name: {}", tun_name);
        println!("guid: {}", dns_mgr.interface_control.interface_guid);

        dns_mgr
            .interface_control
            .set_primary_dns(&[fake_ip.into()], &[".et.net.".to_string()])
            .unwrap();
        dns_mgr.interface_control.flush_dns().unwrap();

        tracing::info!("check dns record with nslookup");

        // nslookup should eventually return 10.144.144.10 after system DNS setting propagation.
        let lookup_deadline = std::time::Instant::now() + Duration::from_secs(30);
        let mut lookup_ok = false;
        let mut last_output = String::new();

        while std::time::Instant::now() < lookup_deadline {
            let result = tokio::time::timeout(
                Duration::from_secs(3),
                tokio::process::Command::new("nslookup")
                    .arg("test1.et.net")
                    .output(),
            )
            .await;

            if let Ok(Ok(ret)) = result {
                let output = String::from_utf8_lossy(&ret.stdout).to_string();
                println!("nslookup output: {}", output);
                if ret.status.success() && output.contains("10.144.144.10") {
                    lookup_ok = true;
                    break;
                }
                last_output = output;
            }

            tokio::time::sleep(Duration::from_millis(500)).await;
        }

        assert!(
            lookup_ok,
            "nslookup did not return expected IP in time; last output: {last_output}"
        );

        dns_node.stop().await.unwrap();
    }
}