diff --git a/Cargo.lock b/Cargo.lock index 79fde39..9b79cbb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1079,6 +1079,7 @@ dependencies = [ "tracing-subscriber", "url", "webpki-roots 0.26.11", + "windows-service", ] [[package]] @@ -2416,6 +2417,12 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "widestring" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72069c3113ab32ab29e5584db3c6ec55d416895e60715417b5b883a357c3e471" + [[package]] name = "winapi-util" version = "0.1.11" @@ -2431,6 +2438,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-service" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "857224b3b211c6f3616921f081ee54721ee3ad2ace2fac6a6337e032f7b4dcf2" +dependencies = [ + "bitflags 2.10.0", + "widestring", + "windows-sys 0.61.2", +] + [[package]] name = "windows-sys" version = "0.45.0" diff --git a/Cargo.toml b/Cargo.toml index 29f7577..490e7f3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,6 +41,10 @@ rustls = { version = "0.23", features = ["ring"] } webpki-roots = "0.26" url = "2" quinn = { version = "0.11", features = ["runtime-tokio", "rustls"] } + +# Windows service support (Windows only, harmless on other platforms) +# Windows 服务支持（仅 Windows，其他平台无害） +windows-service = { version = "0.8" } [dev-dependencies] criterion = "0.5" ctor = "0.2" @@ -97,4 +101,3 @@ strip = true # 详情请参考 GitHub Actions 工作流 # [target.aarch64-unknown-freebsd] # linker = "aarch64-unknown-freebsd12.0-gcc" - diff --git a/README.md b/README.md index 6fd0a4f..5423e69 100644 --- a/README.md +++ b/README.md @@ -90,32 +90,212 @@ OPTIONS: -V, --version Show version ``` -### systemd Service +### Service Management (Cross-Platform) -Create `/etc/systemd/system/kixdns.service`: +KixDNS includes built-in service management that supports all major init systems. +The `service` subcommand works on Windows, Linux, and FreeBSD. + +#### Quick Commands + +```bash +# Install as a system service (auto-detects init system) +sudo ./target/release/kixdns service install + +# Install with custom config +sudo ./target/release/kixdns service install \ + --config /etc/kixdns/pipeline.json \ + --listener-label edge-internal + +# Uninstall the service +sudo ./target/release/kixdns service uninstall + +# Run as service (used by init system internally) +kixdns service run +``` + +> **Note**: `service install` and `service uninstall` require root/Administrator +> privileges. `service run` is called automatically by the init system. + +#### Auto-Detection Logic + +The init system is auto-detected in the following order: + +| Condition | Detected System | +|-----------|----------------| +| `#[cfg(target_os = "windows")]` | Windows SCM | +| `/sbin/procd` exists | Procd (OpenWrt) | +| `/usr/lib/systemd/systemd` or `/lib/systemd/systemd` exists | systemd | +| `/sbin/openrc-run` exists | OpenRC | +| `#[cfg(target_os = "freebsd")]` | BSD rc.d | +| Nothing matched | Unknown (manual setup required) | + +#### Service Scripts + +
+systemd (Linux — most distributions) + +Installed to: `/etc/systemd/system/kixdns.service` ```ini [Unit] -Description=KixDNS +Description=KixDNS DNS Server +Documentation=https://github.com/kixdns/kixdns After=network.target +Wants=nss-lookup.target [Service] Type=simple -ExecStart=/usr/local/bin/kixdns --config /etc/kixdns/pipeline.json +ExecStart=/usr/local/bin/kixdns run \ + --config /etc/kixdns/pipeline.json \ + --listener-label default Restart=on-failure +RestartSec=5 +User=nobody +Group=nogroup +AmbientCapabilities=CAP_NET_BIND_SERVICE +NoNewPrivileges=true LimitNOFILE=65536 +PrivateTmp=true +ProtectSystem=full +ProtectHome=true [Install] WantedBy=multi-user.target ``` ```bash -sudo install -m 0755 target/release/kixdns /usr/local/bin/kixdns -sudo mkdir -p /etc/kixdns -sudo cp config/pipeline.json /etc/kixdns/ -sudo systemctl daemon-reload -sudo systemctl enable --now kixdns +# Manual commands +systemctl daemon-reload +systemctl enable --now kixdns +systemctl status kixdns +journalctl -u kixdns -f +``` +
+ +
+OpenRC (Gentoo, Alpine Linux) + +Installed to: `/etc/init.d/kixdns` + +```sh +#!/sbin/openrc-run + +name="KixDNS DNS Server" +description="High-performance async DNS server" +command="/usr/local/bin/kixdns" +command_args="run --config /etc/kixdns/pipeline.json --listener-label default" +command_user="nobody:nogroup" +pidfile="/run/${RC_SVCNAME}.pid" +command_background=false + +depend() { + need net + use dns logger +} +``` + +```bash +# Manual commands +rc-update add kixdns default +rc-service kixdns start +rc-service kixdns status +``` +
+ +
+Procd (OpenWrt) + +Installed to: `/etc/init.d/kixdns` + +```sh +#!/bin/sh /etc/rc.common + +USE_PROCD=1 + +start_service() { + procd_open_instance + procd_set_param command "/usr/local/bin/kixdns" + procd_append_param command run --config /etc/kixdns/pipeline.json --listener-label default + procd_set_param user nobody + procd_set_param respawn 3600 5 0 + procd_set_param stdout 1 + procd_set_param stderr 1 + procd_close_instance +} + +service_triggers() { + procd_add_reload_trigger "kixdns" +} +``` + +```bash +# Manual commands +/etc/init.d/kixdns enable +/etc/init.d/kixdns start +/etc/init.d/kixdns status +``` +
+ +
+BSD rc.d (FreeBSD) + +Installed to: `/usr/local/etc/rc.d/kixdns` + +```sh +#!/bin/sh +# +# PROVIDE: kixdns +# REQUIRE: NETWORKING SERVERS +# KEYWORD: shutdown + +. /etc/rc.subr + +name="kixdns" +rcvar="kixdns_enable" + +load_rc_config $name + +: ${kixdns_enable:="NO"} +: ${kixdns_config:="/usr/local/etc/kixdns/pipeline.json"} +: ${kixdns_listener_label:="default"} +: ${kixdns_user:="nobody"} + +command="/usr/local/bin/kixdns" +command_args="run --config ${kixdns_config} --listener-label ${kixdns_listener_label}" +command_user="${kixdns_user}" +pidfile="/var/run/kixdns.pid" + +run_rc_command "$1" +``` + +```bash +# Manual commands +echo 'kixdns_enable="YES"' >> /etc/rc.conf +service kixdns start +service kixdns status +``` +
+ +
+Windows (SCM — Service Control Manager) + +```bash +# Install (Administrator PowerShell) +kixdns.exe service install --config config/pipeline.json + +# Start +sc start KixDNS + +# Query status +sc query KixDNS + +# Stop +sc stop KixDNS + +# Uninstall (Administrator PowerShell) +kixdns.exe service uninstall ``` +
## Configuration diff --git a/README.zh-CN.md b/README.zh-CN.md index 80f6f7f..9010351 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -90,32 +90,211 @@ OPTIONS: -V, --version 显示版本 ``` -### systemd 服务 +### 服务管理（跨平台） -创建 `/etc/systemd/system/kixdns.service`： +KixDNS 内置服务管理功能，支持所有主流 init 系统。 +`service` 子命令在 Windows、Linux 和 FreeBSD 上均可使用。 + +#### 快速命令 + +```bash +# 安装为系统服务（自动检测 init 系统） +sudo ./target/release/kixdns service install + +# 使用自定义配置安装 +sudo ./target/release/kixdns service install \ + --config /etc/kixdns/pipeline.json \ + --listener-label edge-internal + +# 卸载服务 +sudo ./target/release/kixdns service uninstall + +# 以服务方式运行（供 init 系统内部调用） +kixdns service run +``` + +> **注意**：安装和卸载需要 root/管理员权限。`service run` 由 init 系统自动调用。 + +#### 自动检测逻辑 + +init 系统按如下顺序自动检测： + +| 检测条件 | 判定结果 | +|----------|----------| +| `#[cfg(target_os = "windows")]` | Windows SCM | +| `/sbin/procd` 存在 | Procd (OpenWrt) | +| `/usr/lib/systemd/systemd` 或 `/lib/systemd/systemd` 存在 | systemd | +| `/sbin/openrc-run` 存在 | OpenRC | +| `#[cfg(target_os = "freebsd")]` | BSD rc.d | +| 以上均不匹配 | 未知（需要手动配置） | + +#### 服务脚本 + +
+systemd（Linux — 大多数发行版） + +安装到：`/etc/systemd/system/kixdns.service` ```ini [Unit] -Description=KixDNS +Description=KixDNS DNS Server +Documentation=https://github.com/kixdns/kixdns After=network.target +Wants=nss-lookup.target [Service] Type=simple -ExecStart=/usr/local/bin/kixdns --config /etc/kixdns/pipeline.json +ExecStart=/usr/local/bin/kixdns run \ + --config /etc/kixdns/pipeline.json \ + --listener-label default Restart=on-failure +RestartSec=5 +User=nobody +Group=nogroup +AmbientCapabilities=CAP_NET_BIND_SERVICE +NoNewPrivileges=true LimitNOFILE=65536 +PrivateTmp=true +ProtectSystem=full +ProtectHome=true [Install] WantedBy=multi-user.target ``` ```bash -sudo install -m 0755 target/release/kixdns /usr/local/bin/kixdns -sudo mkdir -p /etc/kixdns -sudo cp config/pipeline.json /etc/kixdns/ -sudo systemctl daemon-reload -sudo systemctl enable --now kixdns +# 手动命令 +systemctl daemon-reload +systemctl enable --now kixdns +systemctl status kixdns +journalctl -u kixdns -f +``` +
+ +
+OpenRC（Gentoo、Alpine Linux） + +安装到：`/etc/init.d/kixdns` + +```sh +#!/sbin/openrc-run + +name="KixDNS DNS Server" +description="High-performance async DNS server" +command="/usr/local/bin/kixdns" +command_args="run --config /etc/kixdns/pipeline.json --listener-label default" +command_user="nobody:nogroup" +pidfile="/run/${RC_SVCNAME}.pid" +command_background=false + +depend() { + need net + use dns logger +} +``` + +```bash +# 手动命令 +rc-update add kixdns default +rc-service kixdns start +rc-service kixdns status +``` +
+ +
+Procd（OpenWrt） + +安装到：`/etc/init.d/kixdns` + +```sh +#!/bin/sh /etc/rc.common + +USE_PROCD=1 + +start_service() { + procd_open_instance + procd_set_param command "/usr/local/bin/kixdns" + procd_append_param command run --config /etc/kixdns/pipeline.json --listener-label default + procd_set_param user nobody + procd_set_param respawn 3600 5 0 + procd_set_param stdout 1 + procd_set_param stderr 1 + procd_close_instance +} + +service_triggers() { + procd_add_reload_trigger "kixdns" +} +``` + +```bash +# 手动命令 +/etc/init.d/kixdns enable +/etc/init.d/kixdns start +/etc/init.d/kixdns status +``` +
+ +
+BSD rc.d（FreeBSD） + +安装到：`/usr/local/etc/rc.d/kixdns` + +```sh +#!/bin/sh +# +# PROVIDE: kixdns +# REQUIRE: NETWORKING SERVERS +# KEYWORD: shutdown + +. /etc/rc.subr + +name="kixdns" +rcvar="kixdns_enable" + +load_rc_config $name + +: ${kixdns_enable:="NO"} +: ${kixdns_config:="/usr/local/etc/kixdns/pipeline.json"} +: ${kixdns_listener_label:="default"} +: ${kixdns_user:="nobody"} + +command="/usr/local/bin/kixdns" +command_args="run --config ${kixdns_config} --listener-label ${kixdns_listener_label}" +command_user="${kixdns_user}" +pidfile="/var/run/kixdns.pid" + +run_rc_command "$1" +``` + +```bash +# 手动命令 +echo 'kixdns_enable="YES"' >> /etc/rc.conf +service kixdns start +service kixdns status +``` +
+ +
+Windows（SCM — 服务控制管理器） + +```bash +# 安装（管理员 PowerShell） +kixdns.exe service install --config config/pipeline.json + +# 启动 +sc start KixDNS + +# 查询状态 +sc query KixDNS + +# 停止 +sc stop KixDNS + +# 卸载（管理员 PowerShell） +kixdns.exe service uninstall ``` +
## 配置 diff --git a/src/lib.rs b/src/lib.rs index c1b5ccc..5f6d46e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,10 +1,897 @@ pub mod cache; pub mod config; pub mod engine; +pub mod error_utils; pub mod lock; pub mod matcher; pub mod proto_utils; -pub mod watcher; +pub mod service; pub mod socket_utils; -pub mod error_utils; +pub mod watcher; + +// ============================================================================ +// Re-exports & shared server logic +// ============================================================================ + +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream, UdpSocket}; +use tokio::sync::oneshot; +use tracing::{error, info, debug, warn}; +use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt}; + +use crate::config::load_config; +use crate::engine::{Engine, FastPathResponse}; +use crate::matcher::RuntimePipelineConfig; + +// ============================================================================ +// run_dns_server — console mode (no external shutdown signal) +// ============================================================================ + +/// 运行 DNS 服务器（控制台模式）/ Run DNS server (console mode) +pub async fn run_dns_server( + config: PathBuf, + listener_label: String, + debug: bool, + udp_workers_count: usize, +) -> anyhow::Result<()> { + init_tracing(debug); + + // Install default CryptoProvider for rustls + // 安装 rustls 的默认 CryptoProvider + // This is required for rustls 0.23+ when multiple crypto backends are available + // 当有多个加密后端可用时，rustls 0.23+ 需要此调用 + rustls::crypto::ring::default_provider() + .install_default() + .expect("failed to install rustls crypto provider"); + + let cfg = load_config(&config).context("load initial config")?; + let cfg = RuntimePipelineConfig::from_config(cfg).context("compile matchers")?; + let bind_addr: SocketAddr = cfg.settings.bind_udp.parse().context("parse bind addr")?; + let bind_tcp: SocketAddr = cfg.settings.bind_tcp.parse().context("parse tcp bind addr")?; + + let engine = Engine::new(cfg, listener_label.clone()); + + watcher::spawn(config.clone(), engine.clone()); + + // UDP worker 数量：默认为 CPU 核心数，最少 1 个 / UDP worker count: defaults to CPU core count, minimum 1 + let udp_workers_final = if udp_workers_count > 0 { + udp_workers_count + } else { + num_cpus::get() + }; + + info!(bind_udp = %bind_addr, bind_tcp = %bind_tcp, udp_workers_count = udp_workers_final, "dns server started"); + + let mut all_handles: Vec> = Vec::new(); + + #[cfg(unix)] + { + // ✅ OpenBSD 兼容性方案：双 socket（IPv4 + IPv6）+ 零拷贝 recv_buf_from + // ✅ OpenBSD compatibility: dual sockets (IPv4 + IPv6) + zero-copy recv_buf_from + // 为每个地址族创建独立的 socket 和 workers，避免 sockaddr 大小断言失败 + // Create separate sockets and workers for each address family to avoid sockaddr size assertion failures + + // 根据配置地址决定创建哪种 socket / Determine which socket type to create based on config + // IPv6 unspecified address (::) 需要同时创建 IPv4 和 IPv6 socket + // IPv6 other addresses 只创建 IPv6 socket + // IPv4 addresses 只创建 IPv4 socket + let needs_ipv4 = bind_addr.is_ipv4() + || (bind_addr.is_ipv6() && bind_addr.ip().is_unspecified()); + let needs_ipv6 = bind_addr.is_ipv6(); + + if needs_ipv4 { + let workers_per_family = if needs_ipv6 { + udp_workers_final.div_ceil(2) + } else { + udp_workers_final + }; + spawn_ipv4_udp_workers(bind_addr, workers_per_family, engine.clone(), &mut all_handles)?; + } + + if needs_ipv6 { + let workers_per_family = if needs_ipv4 { + udp_workers_final.div_ceil(2) + } else { + udp_workers_final + }; + spawn_ipv6_udp_workers(bind_addr, workers_per_family, engine.clone(), &mut all_handles)?; + } + } + + #[cfg(not(unix))] + { + // Non-Unix: create a single shared socket and spawn workers that share it / 非 Unix：创建单个共享套接字并生成共享它的工作线程 + // Use socket2 to set buffer sizes / 使用 socket2 设置缓冲区大小 + use socket2::{Domain, Protocol, Socket, Type}; + let domain = if bind_addr.is_ipv4() { + Domain::IPV4 + } else { + Domain::IPV6 + }; + let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP)).context("create socket")?; + + // ✅ Windows 上设置 IPV6_V6ONLY=0 以支持双栈，与 Linux 行为一致 + // ✅ On Windows, set IPV6_V6ONLY=0 for dual-stack support, consistent with Linux behavior + if domain == Domain::IPV6 { + if let Err(e) = socket.set_only_v6(false) { + debug!( + "failed to set IPV6_V6ONLY=0: {}, IPv4 may not work on [::] bind", + e + ); + } else { + info!("UDP IPv6 socket set to dual-stack mode (IPV6_V6ONLY=0)"); + } + } + + // Set buffer sizes to prevent packet loss under load + // Try 4MB first, then fall back to 1MB if it fails + let desired_size = 4 * 1024 * 1024; + let fallback_size = 1024 * 1024; + + if let Err(e) = socket.set_recv_buffer_size(desired_size) { + debug!( + "failed to set udp recv buffer to {} bytes: {}, trying {}", + desired_size, e, fallback_size + ); + let _ = socket.set_recv_buffer_size(fallback_size); + } + if let Err(e) = socket.set_send_buffer_size(desired_size) { + debug!( + "failed to set udp send buffer to {} bytes: {}, trying {}", + desired_size, e, fallback_size + ); + let _ = socket.set_send_buffer_size(fallback_size); + } + + socket.set_nonblocking(true).context("set nonblocking")?; + socket.bind(&bind_addr.into()).context("bind socket")?; + + let udp_socket = Arc::new(UdpSocket::from_std(socket.into()).context("from_std")?); + for worker_id in 0..udp_workers_final { + let engine = engine.clone(); + let socket = Arc::clone(&udp_socket); + let handle = tokio::spawn(async move { + if let Err(err) = run_udp_worker(worker_id, socket, engine).await { + error!(worker_id, error = %err, "udp worker exited"); + } + }); + all_handles.push(handle); + } + } + + // TCP listener / TCP 监听器 + // ✅ 双 socket 方案，与 UDP 行为一致 / Dual-socket approach, consistent with UDP + let needs_ipv4_tcp = bind_tcp.is_ipv4() + || (bind_tcp.is_ipv6() && bind_tcp.ip().is_unspecified()); + + // --- 启动 IPv4 TCP 监听 / Start IPv4 TCP listener --- + if needs_ipv4_tcp { + let addr = if bind_tcp.is_ipv4() { + bind_tcp + } else { + SocketAddr::new( + std::net::IpAddr::V4(std::net::Ipv4Addr::new(0, 0, 0, 0)), + bind_tcp.port(), + ) + }; + // 纯 IPv4 绑定，不受 bindv6only 影响 / Pure IPv4 bind, unaffected by bindv6only + let listener = TcpListener::bind(addr).await.context("bind ipv4 tcp")?; + let engine = engine.clone(); + let h = tokio::spawn(async move { + if let Err(err) = run_tcp(listener, engine).await { + error!(error = %err, "ipv4 tcp server exited"); + } + }); + all_handles.push(h); + } + + // --- 启动 IPv6 TCP 监听 / Start IPv6 TCP listener --- + if bind_tcp.is_ipv6() { + use socket2::{Domain, Protocol, Socket, Type}; + let socket = Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP))?; + + // ⭐️ 核心：强制 IPV6_V6ONLY=1，避免和 IPv4 监听器冲突 + // ⭐️ Key: force IPV6_V6ONLY=1 to avoid conflict with IPv4 listener + socket.set_only_v6(true).context("set ipv6 only for kixdns")?; + socket.set_reuse_address(true)?; + + socket.bind(&bind_tcp.into()).context("bind ipv6 tcp socket")?; + socket.listen(128)?; + socket.set_nonblocking(true)?; + + let listener = TcpListener::from_std(socket.into())?; + let engine = engine.clone(); + let h = tokio::spawn(async move { + if let Err(err) = run_tcp(listener, engine).await { + error!(error = %err, "ipv6 tcp server exited"); + } + }); + all_handles.push(h); + } + + // 等待所有任务 / Wait for all tasks + for h in all_handles { + let _ = h.await; + } + + Ok(()) +} + +// ============================================================================ +// run_dns_server_with_shutdown — service mode (can be stopped via oneshot) +// ============================================================================ + +/// 运行 DNS 服务器并支持通过 oneshot 通道优雅关闭 +/// Run DNS server with graceful shutdown via a oneshot channel +/// +/// This function is called by the service module when running as a Windows +/// service. The `shutdown_rx` receiver is signalled when SCM requests stop. +/// 此函数在作为 Windows 服务运行时由 service 模块调用。 +/// 当 SCM 请求停止时，`shutdown_rx` 接收器会收到信号。 +pub async fn run_dns_server_with_shutdown( + config: PathBuf, + listener_label: String, + shutdown_rx: oneshot::Receiver<()>, +) -> anyhow::Result<()> { + // Init tracing — no debug flag, use env var or default to info + // 初始化 tracing —— 无 debug 标志，使用环境变量或默认 info 级别 + init_tracing(false); + + let cfg = load_config(&config).context("load initial config")?; + let cfg = RuntimePipelineConfig::from_config(cfg).context("compile matchers")?; + let bind_addr: SocketAddr = cfg.settings.bind_udp.parse().context("parse bind addr")?; + let bind_tcp: SocketAddr = cfg.settings.bind_tcp.parse().context("parse tcp bind addr")?; + + let engine = Engine::new(cfg, listener_label.clone()); + + watcher::spawn(config.clone(), engine.clone()); + + let udp_workers_final = num_cpus::get().max(1); + + info!( + bind_udp = %bind_addr, + bind_tcp = %bind_tcp, + udp_workers_count = udp_workers_final, + "dns server started (service mode)" + ); + + let mut all_handles: Vec> = Vec::new(); + + #[cfg(unix)] + { + let needs_ipv4 = bind_addr.is_ipv4() + || (bind_addr.is_ipv6() && bind_addr.ip().is_unspecified()); + let needs_ipv6 = bind_addr.is_ipv6(); + + if needs_ipv4 { + let workers_per_family = if needs_ipv6 { + udp_workers_final.div_ceil(2) + } else { + udp_workers_final + }; + spawn_ipv4_udp_workers(bind_addr, workers_per_family, engine.clone(), &mut all_handles)?; + } + + if needs_ipv6 { + let workers_per_family = if needs_ipv4 { + udp_workers_final.div_ceil(2) + } else { + udp_workers_final + }; + spawn_ipv6_udp_workers(bind_addr, workers_per_family, engine.clone(), &mut all_handles)?; + } + } + + #[cfg(not(unix))] + { + use socket2::{Domain, Protocol, Socket, Type}; + let domain = if bind_addr.is_ipv4() { + Domain::IPV4 + } else { + Domain::IPV6 + }; + let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP)).context("create socket")?; + + if domain == Domain::IPV6 { + if let Err(e) = socket.set_only_v6(false) { + debug!("failed to set IPV6_V6ONLY=0: {}, IPv4 may not work on [::] bind", e); + } else { + info!("UDP IPv6 socket set to dual-stack mode (IPV6_V6ONLY=0)"); + } + } + + let desired_size = 4 * 1024 * 1024; + let fallback_size = 1024 * 1024; + + if let Err(e) = socket.set_recv_buffer_size(desired_size) { + debug!("failed to set udp recv buffer to {} bytes: {}, trying {}", desired_size, e, fallback_size); + let _ = socket.set_recv_buffer_size(fallback_size); + } + if let Err(e) = socket.set_send_buffer_size(desired_size) { + debug!("failed to set udp send buffer to {} bytes: {}, trying {}", desired_size, e, fallback_size); + let _ = socket.set_send_buffer_size(fallback_size); + } + + socket.set_nonblocking(true).context("set nonblocking")?; + socket.bind(&bind_addr.into()).context("bind socket")?; + + let udp_socket = Arc::new(UdpSocket::from_std(socket.into()).context("from_std")?); + for worker_id in 0..udp_workers_final { + let engine = engine.clone(); + let socket = Arc::clone(&udp_socket); + let handle = tokio::spawn(async move { + if let Err(err) = run_udp_worker(worker_id, socket, engine).await { + error!(worker_id, error = %err, "udp worker exited"); + } + }); + all_handles.push(handle); + } + } + + // TCP listener / TCP 监听器 + let needs_ipv4_tcp = bind_tcp.is_ipv4() + || (bind_tcp.is_ipv6() && bind_tcp.ip().is_unspecified()); + + if needs_ipv4_tcp { + let addr = if bind_tcp.is_ipv4() { + bind_tcp + } else { + SocketAddr::new( + std::net::IpAddr::V4(std::net::Ipv4Addr::new(0, 0, 0, 0)), + bind_tcp.port(), + ) + }; + let listener = TcpListener::bind(addr).await.context("bind ipv4 tcp")?; + let engine = engine.clone(); + let h = tokio::spawn(async move { + if let Err(err) = run_tcp(listener, engine).await { + error!(error = %err, "ipv4 tcp server exited"); + } + }); + all_handles.push(h); + } + + if bind_tcp.is_ipv6() { + use socket2::{Domain, Protocol, Socket, Type}; + let socket = Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP))?; + socket.set_only_v6(true).context("set ipv6 only for kixdns")?; + socket.set_reuse_address(true)?; + socket.bind(&bind_tcp.into()).context("bind ipv6 tcp socket")?; + socket.listen(128)?; + socket.set_nonblocking(true)?; + let listener = TcpListener::from_std(socket.into())?; + let engine = engine.clone(); + let h = tokio::spawn(async move { + if let Err(err) = run_tcp(listener, engine).await { + error!(error = %err, "ipv6 tcp server exited"); + } + }); + all_handles.push(h); + } + + // Wait for shutdown signal OR all tasks to complete + // 等待关闭信号或所有任务完成 + tokio::select! { + _ = shutdown_rx => { + info!("Service shutdown signal received, stopping server..."); + // Tasks will be dropped when we return, so no need to explicitly cancel + // 返回时任务会被自动丢弃，无需显式取消 + } + _ = async { + for h in &mut all_handles { + let _ = h.await; + } + } => { + info!("All server tasks completed (unexpected)"); + } + } + + Ok(()) +} + +// ============================================================================ +// Tracing initialisation +// ============================================================================ + +fn init_tracing(debug: bool) { + // 默认仅保留错误日志以平衡性能与可观测性，除非显式指定 + // Default to error-level logging to balance performance with observability unless explicitly enabled + let fmt_layer = fmt::layer() + .with_target(false) + .with_ansi(false) + .with_level(debug); + + let level = if debug { "debug" } else { "error" }; + let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(level)); + tracing_subscriber::registry() + .with(filter) + .with(fmt_layer) + .init(); +} + +// ============================================================================ +// UDP worker spawning (Unix — SO_REUSEPORT) +// ============================================================================ + +// 为 IPv4 地址创建并启动 UDP workers / Create and spawn UDP workers for IPv4 address +#[cfg(unix)] +fn spawn_ipv4_udp_workers( + bind_addr: SocketAddr, + worker_count: usize, + engine: Engine, + all_handles: &mut Vec>, +) -> anyhow::Result<()> { + let ipv4_addr: SocketAddr = if bind_addr.is_ipv4() { + bind_addr + } else { + // 预编译的常量地址，避免 unwrap / Precompiled constant address, avoid unwrap + // 使用配置中的端口号而非硬编码 / Use port from config instead of hardcoded + SocketAddr::new( + std::net::IpAddr::V4(std::net::Ipv4Addr::new(0, 0, 0, 0)), + bind_addr.port(), + ) + }; + + info!(bind_addr = %ipv4_addr, workers = worker_count, "Starting IPv4 UDP workers"); + + for worker_id in 0..worker_count { + let engine = engine.clone(); + let std_socket = create_reuseport_udp_socket(ipv4_addr) + .with_context(|| format!("create ipv4 udp socket for worker {}", worker_id))?; + let socket = UdpSocket::from_std(std_socket)?; + let handle = tokio::spawn(async move { + if let Err(err) = run_udp_worker(worker_id, Arc::new(socket), engine).await { + error!(worker_id, error = %err, "IPv4 udp worker exited"); + } + }); + all_handles.push(handle); + } + + Ok(()) +} + +// 为 IPv6 地址创建并启动 UDP workers / Create and spawn UDP workers for IPv6 address +#[cfg(unix)] +fn spawn_ipv6_udp_workers( + bind_addr: SocketAddr, + worker_count: usize, + engine: Engine, + all_handles: &mut Vec>, +) -> anyhow::Result<()> { + let ipv6_addr: SocketAddr = if bind_addr.is_ipv6() { + bind_addr + } else { + // 预编译的常量地址，避免 unwrap / Precompiled constant address, avoid unwrap + // 使用配置中的端口号而非硬编码 / Use port from config instead of hardcoded + SocketAddr::new( + std::net::IpAddr::V6(std::net::Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)), + bind_addr.port(), + ) + }; + + info!(bind_addr = %ipv6_addr, workers = worker_count, "Starting IPv6 UDP workers"); + + for worker_id in 0..worker_count { + let engine = engine.clone(); + let std_socket = create_reuseport_udp_socket(ipv6_addr) + .with_context(|| format!("create ipv6 udp socket for worker {}", worker_id))?; + let socket = UdpSocket::from_std(std_socket)?; + let handle = tokio::spawn(async move { + if let Err(err) = run_udp_worker(worker_id, Arc::new(socket), engine).await { + error!(worker_id, error = %err, "IPv6 udp worker exited"); + } + }); + all_handles.push(handle); + } + + Ok(()) +} + +// 在 Unix 上创建带 SO_REUSEPORT 的 UDP socket；非 Unix 使用标准绑定 / Create UDP socket with SO_REUSEPORT on Unix; use standard binding on non-Unix +#[cfg(unix)] +fn create_reuseport_udp_socket(addr: SocketAddr) -> anyhow::Result { + use socket2::{Domain, Protocol, Socket, Type}; + let domain = if addr.is_ipv4() { + Domain::IPV4 + } else { + Domain::IPV6 + }; + let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?; + socket.set_reuse_address(true)?; + + // ✅ OpenBSD/FreeBSD 安全措施：为 IPv6 socket 显式设置 IPV6_V6ONLY=1 + // ✅ OpenBSD/FreeBSD safety: explicitly set IPV6_V6ONLY=1 for IPv6 sockets + // 双 socket 方案下，IPv6 socket 只处理 IPv6 流量，确保地址族一致性，避免 sockaddr 大小断言失败 + // With dual-socket approach, IPv6 socket only handles IPv6 traffic, ensuring address family consistency + // 这使得我们可以安全地使用零拷贝的 recv_buf_from + // This allows us to safely use zero-copy recv_buf_from + if domain == Domain::IPV6 { + if let Err(e) = crate::socket_utils::set_ipv6_v6only(&socket, true) { + tracing::warn!( + "Failed to set IPV6_V6ONLY=1: {}, this may cause issues on OpenBSD", + e + ); + } + } + + // Try to set SO_REUSEPORT via safe wrapper / 尝试通过安全封装设置 SO_REUSEPORT + if let Err(e) = crate::socket_utils::set_reuseport(&socket, true) { + // Log warning if SO_REUSEPORT fails / SO_REUSEPORT 失败时记录警告 + tracing::warn!("SO_REUSEPORT failed: {}, falling back to shared socket", e); + } + + // Set buffer sizes to prevent packet loss under load + // Try 4MB first, then fall back to 1MB if it fails + let desired_size = 4 * 1024 * 1024; + let fallback_size = 1024 * 1024; + + if let Err(e) = socket.set_recv_buffer_size(desired_size) { + debug!( + "failed to set udp recv buffer to {} bytes: {}, trying {}", + desired_size, e, fallback_size + ); + let _ = socket.set_recv_buffer_size(fallback_size); + } + if let Err(e) = socket.set_send_buffer_size(desired_size) { + debug!( + "failed to set udp send buffer to {} bytes: {}, trying {}", + desired_size, e, fallback_size + ); + let _ = socket.set_send_buffer_size(fallback_size); + } + + socket.set_nonblocking(true)?; + socket.bind(&addr.into())?; + Ok(socket.into()) +} + +// ============================================================================ +// UDP worker +// ============================================================================ + +/// 高性能 UDP worker：直接在接收循环中处理请求，避免 spawn 开销 / High-performance UDP worker: process requests directly in receive loop, avoiding spawn overhead +async fn run_udp_worker( + worker_id: usize, + socket: Arc, + engine: Engine, +) -> anyhow::Result<()> { + // 预分配缓冲区 / Pre-allocate buffer + // 使用 BytesMut 避免 Bytes::copy_from_slice 的内存分配 / Use BytesMut to avoid memory allocation in Bytes::copy_from_slice + use bytes::BytesMut; + let mut buf = BytesMut::with_capacity(4096); + // 复用发送缓冲区：用于缓存命中时 patch TXID，避免每包堆分配 / Reuse send buffer to patch TXID on cache hits, avoiding per-packet heap allocation + let mut send_buf = BytesMut::with_capacity(512); + + // 自适应流控：每 100 个请求检查一次是否需要调整 permits + // Adaptive flow control: check if adjustment needed every 100 requests + let mut request_count = 0u32; + + info!(worker_id, "UDP worker started"); + + loop { + // 确保有足够的空间 / Ensure sufficient space + if buf.capacity() < 4096 { + buf.reserve(4096 - buf.len()); + } + + // ✅ 使用 tokio 的 recv_buf_from 配合 BytesMut，实现零拷贝的高性能接收 + // ✅ Use tokio's recv_buf_from with BytesMut for zero-copy high-performance reception + // 由于使用双 socket 方案（IPv4 + IPv6 分离），不会出现混合地址族的 sockaddr 问题 + // Since we use dual-socket approach (IPv4 + IPv6 separated), no mixed address family sockaddr issues + match socket.recv_buf_from(&mut buf).await { + Ok((_len, peer)) => { + // 零拷贝获取 Bytes / Zero-copy obtain Bytes + let packet_bytes = buf.split().freeze(); + + // 每 100 个请求检查一次流控调整 / Check flow control adjustment every 100 requests + request_count += 1; + if request_count >= 100 { + request_count = 0; + engine.adjust_flow_control(); // Now synchronous with atomic CAS + } + + // ✅ 优化：使用 handle_packet_fast 避免重复解析 + // ✅ Optimization: Use handle_packet_fast to avoid re-parsing + // 如果缓存命中，直接返回；如果缓存未命中，返回预解析的数据 + // If cache hit, return directly; if cache miss, return pre-parsed data + match engine.handle_packet_fast(&packet_bytes, peer) { + Ok(Some(FastPathResponse::Direct(bytes))) => { + // 已包含正确 TXID，可直接发送 / Already contains correct TXID + let _ = socket.send_to(&bytes, peer).await; + } + Ok(Some(FastPathResponse::CacheHit { + cached, + tx_id, + inserted_at, + })) => { + // 复用 send_buf：copy + patch TXID / Reuse send_buf: copy + patch TXID + send_buf.clear(); + if send_buf.capacity() < cached.len() { + send_buf.reserve(cached.len() - send_buf.capacity()); + } + send_buf.extend_from_slice(&cached); + + // RFC 1035 §5.2: Patch TTL based on residence time / 根据停留时间修正 TTL + let elapsed = inserted_at.elapsed().as_secs() as u32; + if elapsed > 0 { + crate::proto_utils::patch_all_ttls(&mut send_buf, elapsed); + } + + if send_buf.len() >= 2 { + let id_bytes = tx_id.to_be_bytes(); + send_buf[0] = id_bytes[0]; + send_buf[1] = id_bytes[1]; + } + let _ = socket.send_to(&send_buf, peer).await; + } + Ok(Some(FastPathResponse::AsyncNeeded { + qname, + qtype, + qclass, + tx_id, + edns_present, + pipeline_id, + })) => { + // 缓存未命中，使用预解析的数据避免重复解析 + // Cache miss, use pre-parsed data to avoid re-parsing + let permit_mgr = Arc::clone(&engine.permit_manager); + let timeout_ms = engine.get_request_timeout_ms(); + let timeout_dur = Duration::from_millis(timeout_ms); + + // 非阻塞式 try_acquire，避免在接收循环中 await / Non-blocking try_acquire to avoid await in receive loop + if let Some(permit) = permit_mgr.try_acquire() { + let engine = engine.clone(); + let socket = Arc::clone(&socket); + let packet_bytes = packet_bytes.clone(); + tokio::spawn(async move { + let _permit = permit; // 自动释放 / Auto-release on drop + + // ✅ 传递预解析数据给 handle_packet_internal，避免重复解析 + // ✅ Pass pre-parsed data to handle_packet_internal to avoid re-parsing + match tokio::time::timeout( + timeout_dur, + engine.handle_packet_internal_with_pre_parsed( + &packet_bytes, + peer, + false, + qname, + qtype, + qclass, + tx_id, + edns_present, + pipeline_id, + ), + ) + .await + { + Ok(Ok(resp)) => { + let _ = socket.send_to(&resp, peer).await; + } + Ok(Err(e)) => { + debug!(error = %e, "handle_packet error"); + } + Err(_) => { + warn!( + timeout_ms, + upstream_timeout_ms = engine.get_upstream_timeout_ms(), + "request timeout after hedge and fallback exhausted" + ); + } + } + }); + } + } + Ok(None) => { + // 快速解析失败，回退到完整处理 + // Fast parse failed, fallback to full processing + let permit_mgr = Arc::clone(&engine.permit_manager); + let timeout_ms = engine.get_request_timeout_ms(); + let timeout_dur = Duration::from_millis(timeout_ms); + + // 非阻塞式 try_acquire，避免在接收循环中 await / Non-blocking try_acquire to avoid await in receive loop + if let Some(permit) = permit_mgr.try_acquire() { + let engine = engine.clone(); + let socket = Arc::clone(&socket); + let packet_bytes = packet_bytes.clone(); + tokio::spawn(async move { + let _permit = permit; // 自动释放 / Auto-release on drop + match tokio::time::timeout( + timeout_dur, + engine.handle_packet(&packet_bytes, peer), + ) + .await + { + Ok(Ok(resp)) => { + let _ = socket.send_to(&resp, peer).await; + } + Ok(Err(e)) => { + debug!(error = %e, "handle_packet error"); + } + Err(_) => { + warn!( + timeout_ms, + upstream_timeout_ms = engine.get_upstream_timeout_ms(), + "request timeout" + ); + } + } + }); + } + } + Err(_) => { + // 解析错误，忽略 / Parse error, ignore + } + } + } + Err(e) => { + // 接收错误，清除缓冲区并继续 / Receive error, clear buffer and continue + debug!(error = %e, "UDP recv error"); + buf.clear(); + } + } + } +} + +// ============================================================================ +// TCP +// ============================================================================ + +async fn run_tcp(listener: TcpListener, engine: Engine) -> anyhow::Result<()> { + loop { + let (stream, peer) = listener.accept().await?; + let engine = engine.clone(); + tokio::spawn(async move { + let _ = handle_tcp_conn(stream, peer, engine).await; + }); + } +} + +async fn handle_tcp_conn( + mut stream: TcpStream, + peer: SocketAddr, + engine: Engine, +) -> anyhow::Result<()> { + const MAX_TCP_FRAME: usize = 64 * 1024; + let mut len_buf = [0u8; 2]; + + // ✅ 获取整体请求超时（包含 hedge + TCP fallback） + // ✅ Get overall request timeout (including hedge + TCP fallback) + let timeout_ms = engine.get_request_timeout_ms(); + + // Reusable buffer to avoid per-frame heap allocation / 可复用缓冲区，避免每帧堆分配 + // 使用 BytesMut 以支持零拷贝操作 / Use BytesMut for zero-copy operations + let mut buf = bytes::BytesMut::with_capacity(MAX_TCP_FRAME); + + loop { + if let Err(err) = stream.read_exact(&mut len_buf).await { + if err.kind() != std::io::ErrorKind::UnexpectedEof { + return Err(err.into()); + } + return Ok(()); + } + let frame_len = u16::from_be_bytes(len_buf) as usize; + if frame_len == 0 || frame_len > MAX_TCP_FRAME { + return Ok(()); + } + + // Reuse buffer: resize to exact frame length / 复用缓冲区：调整到精确帧长度 + // resize() is safe and efficient - it only initializes new bytes if growing + buf.clear(); + buf.resize(frame_len, 0); + if stream.read_exact(&mut buf).await.is_err() { + return Ok(()); + } + + // ✅ 优化：使用 handle_packet_fast 进行快速路径检查 + // ✅ Optimization: Use handle_packet_fast for fast path check + // 统一 UDP 和 TCP 的行为，避免重复解析 + // Unify UDP and TCP behavior to avoid re-parsing + let packet_bytes = buf.split().freeze(); + let timeout_dur = Duration::from_millis(timeout_ms); + + let resp = match engine.handle_packet_fast(&packet_bytes, peer) { + Ok(Some(FastPathResponse::Direct(bytes))) => { + // 快速路径命中：直接返回 / Fast path hit: return directly + bytes + } + Ok(Some(FastPathResponse::CacheHit { + cached, + tx_id, + inserted_at, + })) => { + // 缓存命中：patch TXID / Cache hit: patch TXID + let mut resp_buf = bytes::BytesMut::with_capacity(cached.len()); + resp_buf.extend_from_slice(&cached); + + // RFC 1035 §5.2: Patch TTL based on residence time / 根据停留时间修正 TTL + let elapsed = inserted_at.elapsed().as_secs() as u32; + if elapsed > 0 { + crate::proto_utils::patch_all_ttls(&mut resp_buf, elapsed); + } + + if resp_buf.len() >= 2 { + let id_bytes = tx_id.to_be_bytes(); + resp_buf[0] = id_bytes[0]; + resp_buf[1] = id_bytes[1]; + } + resp_buf.freeze() + } + Ok(Some(FastPathResponse::AsyncNeeded { + qname, + qtype, + qclass, + tx_id, + edns_present, + pipeline_id, + })) => { + // 缓存未命中：使用预解析数据避免重复解析 + // Cache miss: use pre-parsed data to avoid re-parsing + match tokio::time::timeout( + timeout_dur, + engine.handle_packet_internal_with_pre_parsed( + &packet_bytes, + peer, + false, + qname, + qtype, + qclass, + tx_id, + edns_present, + pipeline_id, + ), + ) + .await + { + Ok(Ok(r)) => r, + Ok(Err(_)) => return Ok(()), + Err(_) => { + warn!( + timeout_ms, + upstream_timeout_ms = engine.get_upstream_timeout_ms(), + "TCP request timeout after hedge and fallback exhausted" + ); + return Ok(()); // 关闭连接 / Close connection + } + } + } + Ok(None) => { + // 快速解析失败，回退到完整处理 + // Fast parse failed, fallback to full processing + match tokio::time::timeout(timeout_dur, engine.handle_packet(&packet_bytes, peer)) + .await + { + Ok(Ok(r)) => r, + Ok(Err(_)) => return Ok(()), + Err(_) => { + warn!( + timeout_ms, + upstream_timeout_ms = engine.get_upstream_timeout_ms(), + "TCP request timeout" + ); + return Ok(()); // 关闭连接 / Close connection + } + } + } + Err(_) => { + // 解析错误，关闭连接 / Parse error, close connection + return Ok(()); + } + }; + if resp.len() <= u16::MAX as usize { + let len_bytes = (resp.len() as u16).to_be_bytes(); + if stream.write_all(&len_bytes).await.is_err() { + return Ok(()); + } + if stream.write_all(&resp).await.is_err() { + return Ok(()); + } + } + } +} diff --git a/src/main.rs b/src/main.rs index 6e167c3..cfe7e93 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,19 +1,12 @@ -use std::net::SocketAddr; use std::path::PathBuf; -use std::sync::Arc; -use std::time::Duration; use anyhow::Context; use clap::{Parser, Subcommand}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::{TcpListener, TcpStream, UdpSocket}; -use tracing::{error, info, debug, warn}; -use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt}; +use tracing::error; -use kixdns::config::load_config; -use kixdns::engine::{Engine, FastPathResponse}; -use kixdns::matcher::RuntimePipelineConfig; -use kixdns::watcher; +// ============================================================================ +// CLI argument definitions +// ============================================================================ #[derive(Parser, Debug)] #[command(author, version, about = "KixDNS async DNS with hot-reload pipelines", long_about = None)] @@ -51,668 +44,128 @@ enum Commands { #[arg(short = 'f', long = "filter")] filter: Option, }, + /// Manage system service (install/uninstall/run) / 系统服务管理（安装/卸载/运行） + Service { + #[command(subcommand)] + action: ServiceAction, + }, } -#[tokio::main] -async fn main() -> anyhow::Result<()> { +#[derive(Subcommand, Debug)] +enum ServiceAction { + /// Install KixDNS as a system service / 安装为系统服务 + Install { + /// 配置文件路径（JSON） / Config file path (JSON) + #[arg(short = 'c', long = "config", default_value = "config/pipeline.json")] + config: PathBuf, + /// 监听实例标签 / Listener instance label + #[arg(long = "listener-label", default_value = "default")] + listener_label: String, + }, + /// Uninstall KixDNS system service / 卸载系统服务 + Uninstall, + /// Run as system service (used by init system, not for direct use) / 以服务方式运行（供 init 系统内部使用） + Run { + /// 配置文件路径（JSON） / Config file path (JSON) + #[arg(short = 'c', long = "config", default_value = "config/pipeline.json")] + config: PathBuf, + /// 监听实例标签 / Listener instance label + #[arg(long = "listener-label", default_value = "default")] + listener_label: String, + }, +} + +// ============================================================================ +// Entry point +// ============================================================================ +// +// We do NOT use #[tokio::main] here because the Windows service path must call +// `service_dispatcher::start` from the main thread (it blocks forever). We +// manually create the tokio runtime only when running in console mode. +// +// 我们在此处不使用 #[tokio::main]，因为 Windows 服务路径须在主线程调用 +// `service_dispatcher::start`（它会永久阻塞）。仅在控制台模式下才手动创建 tokio 运行时。 +// ============================================================================ +fn main() -> anyhow::Result<()> { let args = Args::parse(); match args.command { - Some(Commands::ConvertGeoIp { input, output, filter }) => { - // Convert GeoIP .dat to MMDB - let filter_countries: Option> = filter - .map(|f| f.split(',').map(|s| s.trim().to_uppercase()).collect()); - - let filter_slice = filter_countries.as_deref(); - - match kixdns::matcher::geoip::GeoIpManager::convert_dat_to_mmdb(&input, &output, filter_slice) { - Ok(stats) => { - println!("Conversion completed successfully:\n{}", stats); - Ok(()) - } - Err(e) => { - error!("Conversion failed: {}", e); - Err(e) - } - } - } - Some(Commands::Run { config, listener_label, debug, udp_workers_count }) => { - run_dns_server(config, listener_label, debug, udp_workers_count).await + // ---- System service management ---- + Some(Commands::Service { action }) => match action { + ServiceAction::Install { + config, + listener_label, + } => kixdns::service::install_service(&config, &listener_label), + ServiceAction::Uninstall => kixdns::service::uninstall_service(), + ServiceAction::Run { + config, + listener_label, + } => { + // Blocks on Windows (SCM dispatcher), sets up signal handlers on Unix. + // Windows：阻塞主线程（SCM 分发）； + // Unix：设置信号处理。 + kixdns::service::run_service(config, listener_label) + } + }, + + // ---- Convert GeoIP ---- + Some(Commands::ConvertGeoIp { + input, + output, + filter, + }) => convert_geoip(input, output, filter), + + // ---- Run DNS server (console mode) ---- + Some(Commands::Run { + config, + listener_label, + debug, + udp_workers_count, + }) => { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .context("failed to build tokio runtime")?; + rt.block_on(kixdns::run_dns_server( + config, + listener_label, + debug, + udp_workers_count, + )) } + + // ---- Default: no subcommand → run DNS server ---- None => { - // No subcommand provided - run DNS server with defaults - run_dns_server( + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .context("failed to build tokio runtime")?; + rt.block_on(kixdns::run_dns_server( PathBuf::from("config/pipeline.json"), "default".to_string(), false, 0, - ).await + )) } } } -/// 运行 DNS 服务器 / Run DNS server -/// 提取公共逻辑以消除代码重复 / Extract common logic to eliminate code duplication -async fn run_dns_server( - config: PathBuf, - listener_label: String, - debug: bool, - udp_workers_count: usize, -) -> anyhow::Result<()> { - // Run DNS server - init_tracing(debug); - - // Install default CryptoProvider for rustls - // 安装 rustls 的默认 CryptoProvider - // This is required for rustls 0.23+ when multiple crypto backends are available - // 当有多个加密后端可用时，rustls 0.23+ 需要此调用 - rustls::crypto::ring::default_provider() - .install_default() - .expect("failed to install rustls crypto provider"); - - let cfg = load_config(&config).context("load initial config")?; - let cfg = RuntimePipelineConfig::from_config(cfg).context("compile matchers")?; - let bind_addr: SocketAddr = cfg.settings.bind_udp.parse().context("parse bind addr")?; - let bind_tcp: SocketAddr = cfg - .settings - .bind_tcp - .parse() - .context("parse tcp bind addr")?; - - let engine = Engine::new(cfg, listener_label.clone()); - - watcher::spawn(config.clone(), engine.clone()); - - // UDP worker 数量：默认为 CPU 核心数，最少 1 个 / UDP worker count: defaults to CPU core count, minimum 1 - let udp_workers_final = if udp_workers_count > 0 { - udp_workers_count - } else { - num_cpus::get() - }; - - info!(bind_udp = %bind_addr, bind_tcp = %bind_tcp, udp_workers_count = udp_workers_final, "dns server started"); - - let mut all_handles: Vec> = Vec::new(); - - #[cfg(unix)] - { - // ✅ OpenBSD 兼容性方案：双 socket（IPv4 + IPv6）+ 零拷贝 recv_buf_from - // ✅ OpenBSD compatibility: dual sockets (IPv4 + IPv6) + zero-copy recv_buf_from - // 为每个地址族创建独立的 socket 和 workers，避免 sockaddr 大小断言失败 - // Create separate sockets and workers for each address family to avoid sockaddr size assertion failures - - // ✅ OpenBSD 兼容性方案：双 socket（IPv4 + IPv6）+ 零拷贝 recv_buf_from - // ✅ OpenBSD compatibility: dual sockets (IPv4 + IPv6) + zero-copy recv_buf_from - // 为每个地址族创建独立的 socket 和 workers，避免 sockaddr 大小断言失败 - // Create separate sockets and workers for each address family to avoid sockaddr size assertion failures - - // 根据配置地址决定创建哪种 socket / Determine which socket type to create based on config - // IPv6 unspecified address (::) 需要同时创建 IPv4 和 IPv6 socket - // IPv6 other addresses 只创建 IPv6 socket - // IPv4 addresses 只创建 IPv4 socket - let needs_ipv4 = bind_addr.is_ipv4() || - (bind_addr.is_ipv6() && bind_addr.ip().is_unspecified()); - let needs_ipv6 = bind_addr.is_ipv6(); - - if needs_ipv4 { - let workers_per_family = if needs_ipv6 { - udp_workers_final.div_ceil(2) - } else { - udp_workers_final - }; - spawn_ipv4_udp_workers(bind_addr, workers_per_family, engine.clone(), &mut all_handles)?; - } - - if needs_ipv6 { - let workers_per_family = if needs_ipv4 { - udp_workers_final.div_ceil(2) - } else { - udp_workers_final - }; - spawn_ipv6_udp_workers(bind_addr, workers_per_family, engine.clone(), &mut all_handles)?; - } - } - - #[cfg(not(unix))] - { - // Non-Unix: create a single shared socket and spawn workers that share it / 非 Unix：创建单个共享套接字并生成共享它的工作线程 - // Use socket2 to set buffer sizes / 使用 socket2 设置缓冲区大小 - use socket2::{Domain, Protocol, Socket, Type}; - let domain = if bind_addr.is_ipv4() { - Domain::IPV4 - } else { - Domain::IPV6 - }; - let socket = - Socket::new(domain, Type::DGRAM, Some(Protocol::UDP)).context("create socket")?; - - // ✅ Windows 上设置 IPV6_V6ONLY=0 以支持双栈，与 Linux 行为一致 - // ✅ On Windows, set IPV6_V6ONLY=0 for dual-stack support, consistent with Linux behavior - if domain == Domain::IPV6 { - if let Err(e) = socket.set_only_v6(false) { - debug!("failed to set IPV6_V6ONLY=0: {}, IPv4 may not work on [::] bind", e); - } else { - info!("UDP IPv6 socket set to dual-stack mode (IPV6_V6ONLY=0)"); - } - } - - // Set buffer sizes to prevent packet loss under load - // Try 4MB first, then fall back to 1MB if it fails - let desired_size = 4 * 1024 * 1024; - let fallback_size = 1024 * 1024; - - if let Err(e) = socket.set_recv_buffer_size(desired_size) { - debug!("failed to set udp recv buffer to {} bytes: {}, trying {}", desired_size, e, fallback_size); - let _ = socket.set_recv_buffer_size(fallback_size); - } - if let Err(e) = socket.set_send_buffer_size(desired_size) { - debug!("failed to set udp send buffer to {} bytes: {}, trying {}", desired_size, e, fallback_size); - let _ = socket.set_send_buffer_size(fallback_size); - } - - socket.set_nonblocking(true).context("set nonblocking")?; - socket.bind(&bind_addr.into()).context("bind socket")?; - - let udp_socket = Arc::new(UdpSocket::from_std(socket.into()).context("from_std")?); - for worker_id in 0..udp_workers_final { - let engine = engine.clone(); - let socket = Arc::clone(&udp_socket); - let handle = tokio::spawn(async move { - if let Err(err) = run_udp_worker(worker_id, socket, engine).await { - error!(worker_id, error = %err, "udp worker exited"); - } - }); - all_handles.push(handle); - } - } - - // TCP listener / TCP 监听器 - // ✅ 双 socket 方案，与 UDP 行为一致 / Dual-socket approach, consistent with UDP - let needs_ipv4_tcp = bind_tcp.is_ipv4() || (bind_tcp.is_ipv6() && bind_tcp.ip().is_unspecified()); - - // --- 启动 IPv4 TCP 监听 / Start IPv4 TCP listener --- - if needs_ipv4_tcp { - let addr = if bind_tcp.is_ipv4() { - bind_tcp - } else { - SocketAddr::new(std::net::IpAddr::V4(std::net::Ipv4Addr::new(0, 0, 0, 0)), bind_tcp.port()) - }; - // 纯 IPv4 绑定，不受 bindv6only 影响 / Pure IPv4 bind, unaffected by bindv6only - let listener = TcpListener::bind(addr).await.context("bind ipv4 tcp")?; - let engine = engine.clone(); - let h = tokio::spawn(async move { - if let Err(err) = run_tcp(listener, engine).await { - error!(error = %err, "ipv4 tcp server exited"); - } - }); - all_handles.push(h); - } - - // --- 启动 IPv6 TCP 监听 / Start IPv6 TCP listener --- - if bind_tcp.is_ipv6() { - use socket2::{Domain, Protocol, Socket, Type}; - let socket = Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP))?; - - // ⭐️ 核心：强制 IPV6_V6ONLY=1，避免和 IPv4 监听器冲突 - // ⭐️ Key: force IPV6_V6ONLY=1 to avoid conflict with IPv4 listener - socket.set_only_v6(true).context("set ipv6 only for kixdns")?; - socket.set_reuse_address(true)?; - - socket.bind(&bind_tcp.into()).context("bind ipv6 tcp socket")?; - socket.listen(128)?; - socket.set_nonblocking(true)?; - - let listener = TcpListener::from_std(socket.into())?; - let engine = engine.clone(); - let h = tokio::spawn(async move { - if let Err(err) = run_tcp(listener, engine).await { - error!(error = %err, "ipv6 tcp server exited"); - } - }); - all_handles.push(h); - } - - // 等待所有任务 / Wait for all tasks - for h in all_handles { - let _ = h.await; - } - +// ============================================================================ +// GeoIP conversion (synchronous helper) +// ============================================================================ +fn convert_geoip(input: PathBuf, output: PathBuf, filter: Option) -> anyhow::Result<()> { + let filter_countries: Option> = + filter.map(|f| f.split(',').map(|s| s.trim().to_uppercase()).collect()); + let filter_slice = filter_countries.as_deref(); + + match kixdns::matcher::geoip::GeoIpManager::convert_dat_to_mmdb(&input, &output, filter_slice) { + Ok(stats) => { + println!("Conversion completed successfully:\n{}", stats); Ok(()) -} - -fn init_tracing(debug: bool) { - // 默认仅保留错误日志以平衡性能与可观测性，除非显式指定 - // Default to error-level logging to balance performance with observability unless explicitly enabled - let fmt_layer = fmt::layer() - .with_target(false) - .with_ansi(false) - .with_level(debug); - - let level = if debug { "debug" } else { "error" }; - let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(level)); - tracing_subscriber::registry() - .with(filter) - .with(fmt_layer) - .init(); -} - -// 为 IPv4 地址创建并启动 UDP workers / Create and spawn UDP workers for IPv4 address -#[cfg(unix)] -fn spawn_ipv4_udp_workers( - bind_addr: SocketAddr, - worker_count: usize, - engine: Engine, - all_handles: &mut Vec>, -) -> anyhow::Result<()> { - let ipv4_addr: SocketAddr = if bind_addr.is_ipv4() { - bind_addr - } else { - // 预编译的常量地址，避免 unwrap / Precompiled constant address, avoid unwrap - // 使用配置中的端口号而非硬编码 / Use port from config instead of hardcoded - SocketAddr::new(std::net::IpAddr::V4(std::net::Ipv4Addr::new(0, 0, 0, 0)), bind_addr.port()) - }; - - info!(bind_addr = %ipv4_addr, workers = worker_count, "Starting IPv4 UDP workers"); - - for worker_id in 0..worker_count { - let engine = engine.clone(); - let std_socket = create_reuseport_udp_socket(ipv4_addr) - .with_context(|| format!("create ipv4 udp socket for worker {}", worker_id))?; - let socket = UdpSocket::from_std(std_socket)?; - let handle = tokio::spawn(async move { - if let Err(err) = run_udp_worker(worker_id, Arc::new(socket), engine).await { - error!(worker_id, error = %err, "IPv4 udp worker exited"); - } - }); - all_handles.push(handle); - } - - Ok(()) -} - -// 为 IPv6 地址创建并启动 UDP workers / Create and spawn UDP workers for IPv6 address -#[cfg(unix)] -fn spawn_ipv6_udp_workers( - bind_addr: SocketAddr, - worker_count: usize, - engine: Engine, - all_handles: &mut Vec>, -) -> anyhow::Result<()> { - let ipv6_addr: SocketAddr = if bind_addr.is_ipv6() { - bind_addr - } else { - // 预编译的常量地址，避免 unwrap / Precompiled constant address, avoid unwrap - // 使用配置中的端口号而非硬编码 / Use port from config instead of hardcoded - SocketAddr::new(std::net::IpAddr::V6(std::net::Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)), bind_addr.port()) - }; - - info!(bind_addr = %ipv6_addr, workers = worker_count, "Starting IPv6 UDP workers"); - - for worker_id in 0..worker_count { - let engine = engine.clone(); - let std_socket = create_reuseport_udp_socket(ipv6_addr) - .with_context(|| format!("create ipv6 udp socket for worker {}", worker_id))?; - let socket = UdpSocket::from_std(std_socket)?; - let handle = tokio::spawn(async move { - if let Err(err) = run_udp_worker(worker_id, Arc::new(socket), engine).await { - error!(worker_id, error = %err, "IPv6 udp worker exited"); - } - }); - all_handles.push(handle); - } - - Ok(()) -} - -// 在 Unix 上创建带 SO_REUSEPORT 的 UDP socket；非 Unix 使用标准绑定 / Create UDP socket with SO_REUSEPORT on Unix; use standard binding on non-Unix -#[cfg(unix)] -fn create_reuseport_udp_socket(addr: SocketAddr) -> anyhow::Result { - use socket2::{Domain, Protocol, Socket, Type}; - let domain = if addr.is_ipv4() { - Domain::IPV4 - } else { - Domain::IPV6 - }; - let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?; - socket.set_reuse_address(true)?; - - // ✅ OpenBSD/FreeBSD 安全措施：为 IPv6 socket 显式设置 IPV6_V6ONLY=1 - // ✅ OpenBSD/FreeBSD safety: explicitly set IPV6_V6ONLY=1 for IPv6 sockets - // 双 socket 方案下，IPv6 socket 只处理 IPv6 流量，确保地址族一致性，避免 sockaddr 大小断言失败 - // With dual-socket approach, IPv6 socket only handles IPv6 traffic, ensuring address family consistency - // 这使得我们可以安全地使用零拷贝的 recv_buf_from - // This allows us to safely use zero-copy recv_buf_from - if domain == Domain::IPV6 { - if let Err(e) = kixdns::socket_utils::set_ipv6_v6only(&socket, true) { - tracing::warn!("Failed to set IPV6_V6ONLY=1: {}, this may cause issues on OpenBSD", e); } - } - - // Try to set SO_REUSEPORT via safe wrapper / 尝试通过安全封装设置 SO_REUSEPORT - if let Err(e) = kixdns::socket_utils::set_reuseport(&socket, true) { - // Log warning if SO_REUSEPORT fails / SO_REUSEPORT 失败时记录警告 - tracing::warn!("SO_REUSEPORT failed: {}, falling back to shared socket", e); - } - - // Set buffer sizes to prevent packet loss under load - // Try 4MB first, then fall back to 1MB if it fails - let desired_size = 4 * 1024 * 1024; - let fallback_size = 1024 * 1024; - - if let Err(e) = socket.set_recv_buffer_size(desired_size) { - debug!("failed to set udp recv buffer to {} bytes: {}, trying {}", desired_size, e, fallback_size); - let _ = socket.set_recv_buffer_size(fallback_size); - } - if let Err(e) = socket.set_send_buffer_size(desired_size) { - debug!("failed to set udp send buffer to {} bytes: {}, trying {}", desired_size, e, fallback_size); - let _ = socket.set_send_buffer_size(fallback_size); - } - - socket.set_nonblocking(true)?; - socket.bind(&addr.into())?; - Ok(socket.into()) -} - -/// 高性能 UDP worker：直接在接收循环中处理请求，避免 spawn 开销 / High-performance UDP worker: process requests directly in receive loop, avoiding spawn overhead -async fn run_udp_worker( - worker_id: usize, - socket: Arc, - engine: Engine, -) -> anyhow::Result<()> { - // 预分配缓冲区 / Pre-allocate buffer - // 使用 BytesMut 避免 Bytes::copy_from_slice 的内存分配 / Use BytesMut to avoid memory allocation in Bytes::copy_from_slice - use bytes::BytesMut; - let mut buf = BytesMut::with_capacity(4096); - // 复用发送缓冲区：用于缓存命中时 patch TXID，避免每包堆分配 / Reuse send buffer to patch TXID on cache hits, avoiding per-packet heap allocation - let mut send_buf = BytesMut::with_capacity(512); - - // 自适应流控：每 100 个请求检查一次是否需要调整 permits - // Adaptive flow control: check if adjustment needed every 100 requests - let mut request_count = 0u32; - - info!(worker_id, "UDP worker started"); - - loop { - // 确保有足够的空间 / Ensure sufficient space - if buf.capacity() < 4096 { - buf.reserve(4096 - buf.len()); - } - - // ✅ 使用 tokio 的 recv_buf_from 配合 BytesMut，实现零拷贝的高性能接收 - // ✅ Use tokio's recv_buf_from with BytesMut for zero-copy high-performance reception - // 由于使用双 socket 方案（IPv4 + IPv6 分离），不会出现混合地址族的 sockaddr 问题 - // Since we use dual-socket approach (IPv4 + IPv6 separated), no mixed address family sockaddr issues - match socket.recv_buf_from(&mut buf).await { - Ok((_len, peer)) => { - // 零拷贝获取 Bytes / Zero-copy obtain Bytes - let packet_bytes = buf.split().freeze(); - - // 每 100 个请求检查一次流控调整 / Check flow control adjustment every 100 requests - request_count += 1; - if request_count >= 100 { - request_count = 0; - engine.adjust_flow_control(); // Now synchronous with atomic CAS - } - - // ✅ 优化：使用 handle_packet_fast 避免重复解析 - // ✅ Optimization: Use handle_packet_fast to avoid re-parsing - // 如果缓存命中，直接返回；如果缓存未命中，返回预解析的数据 - // If cache hit, return directly; if cache miss, return pre-parsed data - match engine.handle_packet_fast(&packet_bytes, peer) { - Ok(Some(FastPathResponse::Direct(bytes))) => { - // 已包含正确 TXID，可直接发送 / Already contains correct TXID - let _ = socket.send_to(&bytes, peer).await; - } - Ok(Some(FastPathResponse::CacheHit { cached, tx_id, inserted_at })) => { - // 复用 send_buf：copy + patch TXID / Reuse send_buf: copy + patch TXID - send_buf.clear(); - if send_buf.capacity() < cached.len() { - send_buf.reserve(cached.len() - send_buf.capacity()); - } - send_buf.extend_from_slice(&cached); - - // RFC 1035 §5.2: Patch TTL based on residence time / 根据停留时间修正 TTL - let elapsed = inserted_at.elapsed().as_secs() as u32; - if elapsed > 0 { - kixdns::proto_utils::patch_all_ttls(&mut send_buf, elapsed); - } - - if send_buf.len() >= 2 { - let id_bytes = tx_id.to_be_bytes(); - send_buf[0] = id_bytes[0]; - send_buf[1] = id_bytes[1]; - } - let _ = socket.send_to(&send_buf, peer).await; - } - Ok(Some(FastPathResponse::AsyncNeeded { qname, qtype, qclass, tx_id, edns_present, pipeline_id })) => { - // 缓存未命中，使用预解析的数据避免重复解析 - // Cache miss, use pre-parsed data to avoid re-parsing - let permit_mgr = Arc::clone(&engine.permit_manager); - let timeout_ms = engine.get_request_timeout_ms(); - let timeout_dur = Duration::from_millis(timeout_ms); - - // 非阻塞式 try_acquire，避免在接收循环中 await / Non-blocking try_acquire to avoid await in receive loop - if let Some(permit) = permit_mgr.try_acquire() { - let engine = engine.clone(); - let socket = Arc::clone(&socket); - let packet_bytes = packet_bytes.clone(); - tokio::spawn(async move { - let _permit = permit; // 自动释放 / Auto-release on drop - - // ✅ 传递预解析数据给 handle_packet_internal，避免重复解析 - // ✅ Pass pre-parsed data to handle_packet_internal to avoid re-parsing - match tokio::time::timeout( - timeout_dur, - engine.handle_packet_internal_with_pre_parsed( - &packet_bytes, - peer, - false, - qname, - qtype, - qclass, - tx_id, - edns_present, - pipeline_id, - ) - ).await { - Ok(Ok(resp)) => { - let _ = socket.send_to(&resp, peer).await; - } - Ok(Err(e)) => { - debug!(error = %e, "handle_packet error"); - } - Err(_) => { - warn!( - timeout_ms, - upstream_timeout_ms = engine.get_upstream_timeout_ms(), - "request timeout after hedge and fallback exhausted" - ); - } - } - }); - } - } - Ok(None) => { - // 快速解析失败，回退到完整处理 - // Fast parse failed, fallback to full processing - let permit_mgr = Arc::clone(&engine.permit_manager); - let timeout_ms = engine.get_request_timeout_ms(); - let timeout_dur = Duration::from_millis(timeout_ms); - - // 非阻塞式 try_acquire，避免在接收循环中 await / Non-blocking try_acquire to avoid await in receive loop - if let Some(permit) = permit_mgr.try_acquire() { - let engine = engine.clone(); - let socket = Arc::clone(&socket); - let packet_bytes = packet_bytes.clone(); - tokio::spawn(async move { - let _permit = permit; // 自动释放 / Auto-release on drop - match tokio::time::timeout(timeout_dur, engine.handle_packet(&packet_bytes, peer)).await { - Ok(Ok(resp)) => { - let _ = socket.send_to(&resp, peer).await; - } - Ok(Err(e)) => { - debug!(error = %e, "handle_packet error"); - } - Err(_) => { - warn!( - timeout_ms, - upstream_timeout_ms = engine.get_upstream_timeout_ms(), - "request timeout" - ); - } - } - }); - } - } - Err(_) => { - // 解析错误，忽略 / Parse error, ignore - } - } - } - Err(e) => { - // 接收错误，清除缓冲区并继续 / Receive error, clear buffer and continue - debug!(error = %e, "UDP recv error"); - buf.clear(); - } - } - } -} - -async fn run_tcp(listener: TcpListener, engine: Engine) -> anyhow::Result<()> { - loop { - let (stream, peer) = listener.accept().await?; - let engine = engine.clone(); - tokio::spawn(async move { - let _ = handle_tcp_conn(stream, peer, engine).await; - }); - } -} - -async fn handle_tcp_conn( - mut stream: TcpStream, - peer: SocketAddr, - engine: Engine, -) -> anyhow::Result<()> { - const MAX_TCP_FRAME: usize = 64 * 1024; - let mut len_buf = [0u8; 2]; - - // ✅ 获取整体请求超时（包含 hedge + TCP fallback） - // ✅ Get overall request timeout (including hedge + TCP fallback) - let timeout_ms = engine.get_request_timeout_ms(); - - // Reusable buffer to avoid per-frame heap allocation / 可复用缓冲区，避免每帧堆分配 - // 使用 BytesMut 以支持零拷贝操作 / Use BytesMut for zero-copy operations - let mut buf = bytes::BytesMut::with_capacity(MAX_TCP_FRAME); - - loop { - if let Err(err) = stream.read_exact(&mut len_buf).await { - if err.kind() != std::io::ErrorKind::UnexpectedEof { - return Err(err.into()); - } - return Ok(()); - } - let frame_len = u16::from_be_bytes(len_buf) as usize; - if frame_len == 0 || frame_len > MAX_TCP_FRAME { - return Ok(()); - } - - // Reuse buffer: resize to exact frame length / 复用缓冲区：调整到精确帧长度 - // resize() is safe and efficient - it only initializes new bytes if growing - buf.clear(); - buf.resize(frame_len, 0); - if stream.read_exact(&mut buf).await.is_err() { - return Ok(()); - } - - // ✅ 优化：使用 handle_packet_fast 进行快速路径检查 - // ✅ Optimization: Use handle_packet_fast for fast path check - // 统一 UDP 和 TCP 的行为，避免重复解析 - // Unify UDP and TCP behavior to avoid re-parsing - let packet_bytes = buf.split().freeze(); - let timeout_dur = Duration::from_millis(timeout_ms); - - let resp = match engine.handle_packet_fast(&packet_bytes, peer) { - Ok(Some(FastPathResponse::Direct(bytes))) => { - // 快速路径命中：直接返回 / Fast path hit: return directly - bytes - } - Ok(Some(FastPathResponse::CacheHit { cached, tx_id, inserted_at })) => { - // 缓存命中：patch TXID / Cache hit: patch TXID - let mut resp_buf = bytes::BytesMut::with_capacity(cached.len()); - resp_buf.extend_from_slice(&cached); - - // RFC 1035 §5.2: Patch TTL based on residence time / 根据停留时间修正 TTL - let elapsed = inserted_at.elapsed().as_secs() as u32; - if elapsed > 0 { - kixdns::proto_utils::patch_all_ttls(&mut resp_buf, elapsed); - } - - if resp_buf.len() >= 2 { - let id_bytes = tx_id.to_be_bytes(); - resp_buf[0] = id_bytes[0]; - resp_buf[1] = id_bytes[1]; - } - resp_buf.freeze() - } - Ok(Some(FastPathResponse::AsyncNeeded { qname, qtype, qclass, tx_id, edns_present, pipeline_id })) => { - // 缓存未命中：使用预解析数据避免重复解析 - // Cache miss: use pre-parsed data to avoid re-parsing - match tokio::time::timeout( - timeout_dur, - engine.handle_packet_internal_with_pre_parsed( - &packet_bytes, - peer, - false, - qname, - qtype, - qclass, - tx_id, - edns_present, - pipeline_id, - ) - ).await { - Ok(Ok(r)) => r, - Ok(Err(_)) => return Ok(()), - Err(_) => { - warn!( - timeout_ms, - upstream_timeout_ms = engine.get_upstream_timeout_ms(), - "TCP request timeout after hedge and fallback exhausted" - ); - return Ok(()); // 关闭连接 / Close connection - } - } - } - Ok(None) => { - // 快速解析失败，回退到完整处理 - // Fast parse failed, fallback to full processing - match tokio::time::timeout(timeout_dur, engine.handle_packet(&packet_bytes, peer)).await { - Ok(Ok(r)) => r, - Ok(Err(_)) => return Ok(()), - Err(_) => { - warn!( - timeout_ms, - upstream_timeout_ms = engine.get_upstream_timeout_ms(), - "TCP request timeout" - ); - return Ok(()); // 关闭连接 / Close connection - } - } - } - Err(_) => { - // 解析错误，关闭连接 / Parse error, close connection - return Ok(()); - } - }; - - if resp.len() <= u16::MAX as usize { - let len_bytes = (resp.len() as u16).to_be_bytes(); - if stream.write_all(&len_bytes).await.is_err() { - return Ok(()); - } - if stream.write_all(&resp).await.is_err() { - return Ok(()); - } + Err(e) => { + error!("Conversion failed: {}", e); + Err(e) } } } diff --git a/src/service/mod.rs b/src/service/mod.rs new file mode 100644 index 0000000..2d17b0e --- /dev/null +++ b/src/service/mod.rs @@ -0,0 +1,151 @@ +//! Cross-platform service management for KixDNS +//! +//! Provides the ability to install, uninstall, and run KixDNS as a system service +//! on Windows (SCM), Linux (systemd, OpenRC, Procd), and FreeBSD (BSD rc.d). +//! +//! 跨平台服务管理：支持在 Windows (SCM)、Linux (systemd/OpenRC/Procd) +//! 和 FreeBSD (BSD rc.d) 上安装、卸载和运行服务。 + +#[cfg(target_os = "windows")] +pub mod windows; + +#[cfg(not(target_os = "windows"))] +pub mod unix; + +// ============================================================================ +// Common types +// ============================================================================ + +/// Detected init system on the current platform. +/// 检测到的当前平台的 init 系统。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InitSystem { + /// Windows Service Control Manager / Windows 服务控制管理器 + #[cfg(target_os = "windows")] + Scm, + /// systemd (most modern Linux distributions) + Systemd, + /// OpenRC (Gentoo, Alpine, etc.) + OpenRc, + /// Procd (OpenWrt) + Procd, + /// BSD rc.d (FreeBSD, etc.) + BsdRc, + /// Unknown — cannot auto-detect + Unknown, +} + +impl InitSystem { + /// Human-readable description of this init system. + pub fn description(&self) -> &'static str { + match self { + #[cfg(target_os = "windows")] + InitSystem::Scm => "Windows Service Control Manager", + InitSystem::Systemd => "systemd (Linux)", + InitSystem::OpenRc => "OpenRC (Gentoo/Alpine Linux)", + InitSystem::Procd => "Procd (OpenWrt)", + InitSystem::BsdRc => "BSD rc.d (FreeBSD)", + InitSystem::Unknown => "Unknown init system", + } + } +} + +/// Detect the init system on the current machine. +/// +/// Detection order: +/// 1. Windows → SCM +/// 2. procd (/sbin/procd) → OpenWrt +/// 3. systemd (/usr/lib/systemd/systemd) → systemd +/// 4. openrc-run (/sbin/openrc-run) → OpenRC +/// 5. target_os = "freebsd" → BSD rc.d +/// 6. Otherwise → Unknown +/// +/// 检测当前机器上的 init 系统。 +pub fn detect_init_system() -> InitSystem { + #[cfg(target_os = "windows")] + { + return InitSystem::Scm; + } + + #[cfg(not(target_os = "windows"))] + { + // Procd (OpenWrt) + if std::path::Path::new("/sbin/procd").exists() { + return InitSystem::Procd; + } + // systemd + if std::path::Path::new("/usr/lib/systemd/systemd").exists() + || std::path::Path::new("/lib/systemd/systemd").exists() + { + return InitSystem::Systemd; + } + // OpenRC + if std::path::Path::new("/sbin/openrc-run").exists() { + return InitSystem::OpenRc; + } + // FreeBSD + #[cfg(target_os = "freebsd")] + { + return InitSystem::BsdRc; + } + InitSystem::Unknown + } +} + +// ============================================================================ +// Top-level API (delegates to platform-specific modules) +// ============================================================================ + +/// Install KixDNS as a system service. +/// +/// On Windows: registers with SCM. +/// On Unix: auto-detects the init system and installs the appropriate service file. +/// +/// 安装 KixDNS 为系统服务。 +pub fn install_service(config: &std::path::PathBuf, listener_label: &str) -> anyhow::Result<()> { + #[cfg(target_os = "windows")] + { + windows::install_service(config, listener_label) + } + + #[cfg(not(target_os = "windows"))] + { + unix::install_service(config, listener_label) + } +} + +/// Uninstall KixDNS system service. +/// +/// On Windows: removes from SCM database. +/// On Unix: auto-detects the init system and removes the appropriate service file. +/// +/// 卸载 KixDNS 系统服务。 +pub fn uninstall_service() -> anyhow::Result<()> { + #[cfg(target_os = "windows")] + { + windows::uninstall_service() + } + + #[cfg(not(target_os = "windows"))] + { + unix::uninstall_service() + } +} + +/// Run KixDNS as a service (platform-specific). +/// +/// On Windows: dispatches via SCM with proper event handling. +/// On Unix: runs the DNS server with a signal-based shutdown handler. +/// +/// 以服务方式运行 KixDNS（平台相关）。 +pub fn run_service(config: std::path::PathBuf, listener_label: String) -> anyhow::Result<()> { + #[cfg(target_os = "windows")] + { + windows::run_service_dispatcher(config, listener_label) + } + + #[cfg(not(target_os = "windows"))] + { + unix::run_service(config, listener_label) + } +} diff --git a/src/service/unix.rs b/src/service/unix.rs new file mode 100644 index 0000000..8da272c --- /dev/null +++ b/src/service/unix.rs @@ -0,0 +1,485 @@ +//! Unix service management for KixDNS +//! +//! Supports: +//! - systemd (most modern Linux distributions) +//! - OpenRC (Gentoo, Alpine Linux) +//! - Procd (OpenWrt) +//! - BSD rc.d (FreeBSD) +//! +//! Unix 服务管理：支持 systemd、OpenRC、Procd 和 BSD rc.d。 + +use std::io::Write; +use std::path::{Path, PathBuf}; + +use anyhow::Context; +use tokio::sync::oneshot; +use tracing::info; + +use crate::service::InitSystem; +use crate::service::detect_init_system; + +// ============================================================================ +// Service file templates +// ============================================================================ + +/// Generate a systemd service unit. +/// +/// Output: `/etc/systemd/system/kixdns.service` +fn systemd_unit(exe: &Path, config: &Path, listener_label: &str) -> String { + format!( + r#"# KixDNS systemd service unit +# Auto-generated by kixdns service install — do not edit manually. +# +# Install: kixdns service install +# Uninstall: kixdns service uninstall +# Start: systemctl start kixdns +# Enable: systemctl enable kixdns + +[Unit] +Description=KixDNS DNS Server +Documentation=https://github.com/kixdns/kixdns +After=network.target +Wants=nss-lookup.target + +[Service] +Type=simple +ExecStart={exe} run \ + --config {config} \ + --listener-label {label} +Restart=on-failure +RestartSec=5 +User=nobody +Group=nogroup +AmbientCapabilities=CAP_NET_BIND_SERVICE +NoNewPrivileges=true +LimitNOFILE=65536 +PrivateTmp=true +ProtectSystem=full +ProtectHome=true + +[Install] +WantedBy=multi-user.target +"#, + exe = exe.display(), + config = config.display(), + label = listener_label, + ) +} + +/// Generate an OpenRC init script. +/// +/// Output: `/etc/init.d/kixdns` +fn openrc_script(exe: &Path, config: &Path, listener_label: &str) -> String { + format!( + r#"#!/sbin/openrc-run +# KixDNS OpenRC init script +# Auto-generated by kixdns service install — do not edit manually. +# +# Install: kixdns service install +# Uninstall: kixdns service uninstall +# Start: rc-service kixdns start +# or: /etc/init.d/kixdns start + +name="KixDNS DNS Server" +description="High-performance async DNS server with hot-reloadable pipeline configuration" +command="{exe}" +command_args="run --config {config} --listener-label {label}" +command_user="nobody:nogroup" +pidfile="/run/${{RC_SVCNAME}}.pid" +command_background=false + +depend() {{ + need net + use dns logger +}} +"#, + exe = exe.display(), + config = config.display(), + label = listener_label, + ) +} + +/// Generate a Procd init script (OpenWrt). +/// +/// Output: `/etc/init.d/kixdns` +fn procd_script(exe: &Path, config: &Path, listener_label: &str) -> String { + format!( + r#"#!/bin/sh /etc/rc.common +# KixDNS Procd init script (OpenWrt) +# Auto-generated by kixdns service install — do not edit manually. +# +# Install: kixdns service install +# Uninstall: kixdns service uninstall +# Start: /etc/init.d/kixdns start +# Enable: /etc/init.d/kixdns enable + +USE_PROCD=1 + +start_service() {{ + procd_open_instance + procd_set_param command "{exe}" + procd_append_param command run --config {config} --listener-label {label} + procd_set_param user nobody + procd_set_param respawn 3600 5 0 + procd_set_param stdout 1 + procd_set_param stderr 1 + procd_close_instance +}} + +stop_service() {{ + # The service will be stopped automatically by procd + : +}} + +service_triggers() {{ + procd_add_reload_trigger "kixdns" +}} +"#, + exe = exe.display(), + config = config.display(), + label = listener_label, + ) +} + +/// Generate a FreeBSD rc.d script. +/// +/// Output: `/usr/local/etc/rc.d/kixdns` +fn bsd_rc_script(exe: &Path, config: &Path, listener_label: &str) -> String { + format!( + r#"#!/bin/sh +# +# PROVIDE: kixdns +# REQUIRE: NETWORKING SERVERS +# KEYWORD: shutdown +# +# Auto-generated by kixdns service install — do not edit manually. +# +# Install: kixdns service install +# Uninstall: kixdns service uninstall +# Start: service kixdns start +# Enable: echo 'kixdns_enable="YES"' >> /etc/rc.conf + +. /etc/rc.subr + +name="kixdns" +rcvar="kixdns_enable" + +load_rc_config $name + +: ${{kixdns_enable:="NO"}} +: ${{kixdns_config:="{config}"}} +: ${{kixdns_listener_label:="{label}"}} +: ${{kixdns_user:="nobody"}} + +command="{exe}" +command_args="run --config ${{kixdns_config}} --listener-label ${{kixdns_listener_label}}" +command_user="${{kixdns_user}}" +pidfile="/var/run/kixdns.pid" + +run_rc_command "$1" +"#, + exe = exe.display(), + config = config.display(), + label = listener_label, + ) +} + +// ============================================================================ +// Service file installation paths +// ============================================================================ + +fn systemd_path() -> PathBuf { + PathBuf::from("/etc/systemd/system/kixdns.service") +} + +fn openrc_path() -> PathBuf { + PathBuf::from("/etc/init.d/kixdns") +} + +fn procd_path() -> PathBuf { + PathBuf::from("/etc/init.d/kixdns") +} + +fn bsd_rc_path() -> PathBuf { + PathBuf::from("/usr/local/etc/rc.d/kixdns") +} + +// ============================================================================ +// Install logic +// ============================================================================ + +/// Install KixDNS as a system service on Unix. +/// +/// Auto-detects the init system and installs the appropriate service file. +pub fn install_service(config: &PathBuf, listener_label: &str) -> anyhow::Result<()> { + let init = detect_init_system(); + let exe = std::env::current_exe().context("failed to get current executable path")?; + + println!("→ Detected init system: {}", init.description()); + + match init { + InitSystem::Systemd => install_systemd(&exe, config, listener_label), + InitSystem::OpenRc => install_openrc(&exe, config, listener_label), + InitSystem::Procd => install_procd(&exe, config, listener_label), + InitSystem::BsdRc => install_bsd_rc(&exe, config, listener_label), + InitSystem::Unknown | _ => { + anyhow::bail!( + "Could not auto-detect init system.\n\ + ┌─ Please create a service script manually ──────────────────────┐\n\ + │ │\n\ + │ systemd: Write a .service unit to /etc/systemd/system/ │\n\ + │ OpenRC: Write an init script to /etc/init.d/ │\n\ + │ Procd: Write a procd script to /etc/init.d/ │\n\ + │ FreeBSD: Write a rc.d script to /usr/local/etc/rc.d/ │\n\ + │ │\n\ + │ The binary is at: {} │\n\ + │ Run it with: {} run --config {} │\n\ + │ │\n\ + └──────────────────────────────────────────────────────────────┘\n\ + ", + exe.display(), + exe.display(), + config.display(), + ) + } + } +} + +fn install_systemd(exe: &Path, config: &Path, listener_label: &str) -> anyhow::Result<()> { + let path = systemd_path(); + let content = systemd_unit(exe, config, listener_label); + + write_service_file(&path, &content, 0o644)?; + + // Reload systemd, enable and start + run_command("systemctl", &["daemon-reload"])?; + run_command("systemctl", &["enable", "kixdns"])?; + run_command("systemctl", &["start", "kixdns"])?; + + println!("✓ KixDNS installed as a systemd service."); + println!(" Unit file : {}", path.display()); + println!(" Status : systemctl status kixdns"); + println!(" Logs : journalctl -u kixdns -f"); + Ok(()) +} + +fn install_openrc(exe: &Path, config: &Path, listener_label: &str) -> anyhow::Result<()> { + let path = openrc_path(); + let content = openrc_script(exe, config, listener_label); + + write_service_file(&path, &content, 0o755)?; + + run_command("rc-update", &["add", "kixdns", "default"])?; + run_command("rc-service", &["kixdns", "start"])?; + + println!("✓ KixDNS installed as an OpenRC service."); + println!(" Init script : {}", path.display()); + println!(" Status : rc-service kixdns status"); + Ok(()) +} + +fn install_procd(exe: &Path, config: &Path, listener_label: &str) -> anyhow::Result<()> { + let path = procd_path(); + let content = procd_script(exe, config, listener_label); + + write_service_file(&path, &content, 0o755)?; + + run_command("/etc/init.d/kixdns", &["enable"])?; + run_command("/etc/init.d/kixdns", &["start"])?; + + println!("✓ KixDNS installed as a Procd service (OpenWrt)."); + println!(" Init script : {}", path.display()); + println!(" Status : /etc/init.d/kixdns status"); + Ok(()) +} + +fn install_bsd_rc(exe: &Path, config: &Path, listener_label: &str) -> anyhow::Result<()> { + let path = bsd_rc_path(); + let content = bsd_rc_script(exe, config, listener_label); + + write_service_file(&path, &content, 0o755)?; + + println!("✓ KixDNS installed as a FreeBSD rc.d service."); + println!(" Script : {}", path.display()); + println!(" Enable : echo 'kixdns_enable=\"YES\"' >> /etc/rc.conf"); + println!(" Start : service kixdns start"); + println!( + " Config : Add 'kixdns_config=\"{}\"' to /etc/rc.conf for custom path", + config.display() + ); + Ok(()) +} + +// ============================================================================ +// Uninstall logic +// ============================================================================ + +/// Uninstall KixDNS system service on Unix. +pub fn uninstall_service() -> anyhow::Result<()> { + let init = detect_init_system(); + + println!("→ Detected init system: {}", init.description()); + + match init { + InitSystem::Systemd => uninstall_systemd(), + InitSystem::OpenRc => uninstall_openrc(), + InitSystem::Procd => uninstall_procd(), + InitSystem::BsdRc => uninstall_bsd_rc(), + InitSystem::Unknown | _ => { + anyhow::bail!( + "Could not auto-detect init system.\n\ + Please manually remove the service file you created earlier." + ); + } + } +} + +fn uninstall_systemd() -> anyhow::Result<()> { + let path = systemd_path(); + + // Stop and disable + let _ = run_command("systemctl", &["stop", "kixdns"]); + let _ = run_command("systemctl", &["disable", "kixdns"]); + + // Remove unit file + if path.exists() { + std::fs::remove_file(&path).context("failed to remove systemd unit file")?; + } + let _ = run_command("systemctl", &["daemon-reload"]); + + println!("✓ KixDNS systemd service uninstalled."); + Ok(()) +} + +fn uninstall_openrc() -> anyhow::Result<()> { + let path = openrc_path(); + + let _ = run_command("rc-service", &["kixdns", "stop"]); + let _ = run_command("rc-update", &["delete", "kixdns"]); + + if path.exists() { + std::fs::remove_file(&path).context("failed to remove OpenRC init script")?; + } + + println!("✓ KixDNS OpenRC service uninstalled."); + Ok(()) +} + +fn uninstall_procd() -> anyhow::Result<()> { + let path = procd_path(); + + if path.exists() { + let _ = run_command("/etc/init.d/kixdns", &["stop"]); + let _ = run_command("/etc/init.d/kixdns", &["disable"]); + std::fs::remove_file(&path).context("failed to remove Procd init script")?; + } + + println!("✓ KixDNS Procd service uninstalled."); + Ok(()) +} + +fn uninstall_bsd_rc() -> anyhow::Result<()> { + let path = bsd_rc_path(); + + let _ = run_command("service", &["kixdns", "stop"]); + + if path.exists() { + std::fs::remove_file(&path).context("failed to remove BSD rc.d script")?; + } + + println!("✓ KixDNS BSD rc.d service uninstalled."); + println!(" Note: You may want to remove 'kixdns_enable=\"YES\"' from /etc/rc.conf manually."); + Ok(()) +} + +// ============================================================================ +// Run as service (Unix) +// ============================================================================ + +/// Run the DNS server with signal-based graceful shutdown. +/// +/// Sets up handlers for SIGTERM, SIGINT, and SIGQUIT to trigger +/// a graceful shutdown of the server. +pub fn run_service(config: PathBuf, listener_label: String) -> anyhow::Result<()> { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .context("failed to build tokio runtime for service")?; + + rt.block_on(async move { + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + + // Set up signal handlers inside the tokio runtime + // 在 tokio 运行时内设置信号处理 + tokio::spawn(async move { + wait_for_signal_async().await; + let _ = shutdown_tx.send(()); + }); + + crate::run_dns_server_with_shutdown(config, listener_label, shutdown_rx).await + })?; + + info!("Service stopped gracefully"); + Ok(()) +} + +/// Wait for a termination signal using tokio's signal module. +/// Listens for SIGTERM, SIGINT, and SIGQUIT. +pub async fn wait_for_signal_async() { + use tokio::signal::unix::{SignalKind, signal}; + + let mut term = signal(SignalKind::terminate()).expect("failed to create SIGTERM listener"); + let mut int = signal(SignalKind::interrupt()).expect("failed to create SIGINT listener"); + let mut quit = signal(SignalKind::quit()).expect("failed to create SIGQUIT listener"); + + tokio::select! { + _ = term.recv() => info!("Received SIGTERM"), + _ = int.recv() => info!("Received SIGINT"), + _ = quit.recv() => info!("Received SIGQUIT"), + } +} + +// ============================================================================ +// Helpers +// ============================================================================ + +/// Write a service file to disk, creating parent directories if needed. +fn write_service_file(path: &Path, content: &str, mode: u32) -> anyhow::Result<()> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("failed to create directory: {}", parent.display()))?; + } + + let mut file = std::fs::File::create(path) + .with_context(|| format!("failed to create service file: {}", path.display()))?; + + file.write_all(content.as_bytes()) + .with_context(|| format!("failed to write service file: {}", path.display()))?; + + // Set file permissions + use std::os::unix::fs::PermissionsExt; + file.set_permissions(std::fs::Permissions::from_mode(mode)) + .with_context(|| format!("failed to set permissions on: {}", path.display()))?; + + Ok(()) +} + +/// Run a shell command and return its exit status. +fn run_command(program: &str, args: &[&str]) -> anyhow::Result<()> { + let output = std::process::Command::new(program) + .args(args) + .output() + .with_context(|| format!("failed to execute: {} {}", program, args.join(" ")))?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + anyhow::bail!( + "Command failed ({} {}): {}", + program, + args.join(" "), + stderr.trim() + ); + } + + Ok(()) +} diff --git a/src/service/windows.rs b/src/service/windows.rs new file mode 100644 index 0000000..98247b6 --- /dev/null +++ b/src/service/windows.rs @@ -0,0 +1,246 @@ +//! Windows Service Control Manager (SCM) integration for KixDNS +//! +//! Windows SCM 集成：安装、卸载、运行 Windows 服务。 + +use std::ffi::OsString; +use std::path::PathBuf; +use std::sync::OnceLock; +use std::time::Duration; + +use anyhow::Context; +use tokio::sync::oneshot; +use tracing::info; + +use windows_service::{ + define_windows_service, + service::{ + ServiceAccess, ServiceControl, ServiceControlAccept, ServiceDependency, + ServiceErrorControl, ServiceExitCode, ServiceInfo, ServiceStartType, ServiceState, + ServiceStatus, ServiceType, + }, + service_control_handler::{self, ServiceControlHandlerResult}, + service_dispatcher, + service_manager::{ServiceManager, ServiceManagerAccess}, +}; + +/// Service name registered with the SCM / 在 SCM 中注册的服务名称 +pub const SERVICE_NAME: &str = "KixDNS"; + +/// Display name shown in Services.msc / 在服务管理器中显示的名称 +pub const SERVICE_DISPLAY_NAME: &str = "KixDNS DNS Server"; + +/// Service description / 服务描述 +pub const SERVICE_DESCRIPTION: &str = + "High-performance async DNS server with hot-reloadable pipeline configuration."; + +/// Global configuration: (config_path, listener_label) +/// Set before dispatching the service, read inside the service main function. +static SERVICE_CONFIG: OnceLock<(PathBuf, String)> = OnceLock::new(); + +// --------------------------------------------------------------------------- +// FFI service entry point — generated by the macro +// --------------------------------------------------------------------------- +define_windows_service!(ffi_service_main, my_service_main); + +/// The actual service main function called by the SCM on a background thread. +fn my_service_main(_arguments: Vec) { + // Init tracing inside the service context + let _ = tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .try_init(); + + let (config, listener_label) = match SERVICE_CONFIG.get() { + Some(cfg) => cfg.clone(), + None => { + tracing::error!("SERVICE_CONFIG not set; using defaults"); + (PathBuf::from("config/pipeline.json"), "default".to_string()) + } + }; + + if let Err(e) = run_service(config, listener_label) { + tracing::error!("Service exited with error: {:#}", e); + } +} + +/// Run the DNS server as a Windows service with proper SCM integration. +fn run_service(config: PathBuf, listener_label: String) -> anyhow::Result<()> { + // 1. Install the ring crypto provider (same as normal startup) + rustls::crypto::ring::default_provider() + .install_default() + .expect("failed to install rustls crypto provider"); + + // 2. Register service control event handler + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + let shutdown_tx = std::sync::Mutex::new(Some(shutdown_tx)); + + let status_handle = service_control_handler::register(SERVICE_NAME, move |control_event| { + match control_event { + ServiceControl::Stop | ServiceControl::Shutdown => { + info!("Service received stop/shutdown signal"); + if let Some(tx) = shutdown_tx.lock().unwrap().take() { + let _ = tx.send(()); + } + ServiceControlHandlerResult::NoError + } + // Interrogate is mandatory — must respond NoError + ServiceControl::Interrogate => ServiceControlHandlerResult::NoError, + _ => ServiceControlHandlerResult::NotImplemented, + } + })?; + + // 3. Report SERVICE_RUNNING to SCM + status_handle.set_service_status(ServiceStatus { + service_type: ServiceType::OWN_PROCESS, + current_state: ServiceState::Running, + controls_accepted: ServiceControlAccept::STOP | ServiceControlAccept::SHUTDOWN, + exit_code: ServiceExitCode::NO_ERROR, + checkpoint: 0, + wait_hint: Duration::default(), + process_id: None, + })?; + + info!( + "Service is now running (config={:?}, listener_label={})", + config, listener_label + ); + + // 4. Create tokio runtime and run the DNS server, passing the shutdown signal + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .context("failed to build tokio runtime for service")?; + + rt.block_on(async { + crate::run_dns_server_with_shutdown(config, listener_label, shutdown_rx).await + })?; + + // 5. Report SERVICE_STOPPED + status_handle.set_service_status(ServiceStatus { + service_type: ServiceType::OWN_PROCESS, + current_state: ServiceState::Stopped, + controls_accepted: ServiceControlAccept::empty(), + exit_code: ServiceExitCode::NO_ERROR, + checkpoint: 0, + wait_hint: Duration::default(), + process_id: None, + })?; + + info!("Service stopped gracefully"); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Public API — called from mod.rs +// --------------------------------------------------------------------------- + +/// Install KixDNS as a Windows service. +/// +/// Registers the service with the SCM so that it can be started automatically +/// on boot or manually via `sc start` / Services.msc. +pub fn install_service(config: &PathBuf, listener_label: &str) -> anyhow::Result<()> { + let manager = ServiceManager::local_computer( + None::<&str>, // default database (active) + ServiceManagerAccess::CONNECT | ServiceManagerAccess::CREATE_SERVICE, + ) + .context("failed to open Service Control Manager (try running as Administrator)")?; + + let exe_path = std::env::current_exe().context("failed to get current executable path")?; + + // Build launch arguments: "service run --config --listener-label