pub mod cache; pub mod query; pub mod server; use anyhow::Result; use std::collections::HashMap; use std::sync::Arc; use tracing::{info, warn}; use crate::config; /// daemon 入口 /// /// 当 WX_DAEMON_MODE 环境变量设置时,main() 调用此函数 pub fn run() { let rt = tokio::runtime::Runtime::new().expect("无法创建 tokio runtime"); if let Err(e) = rt.block_on(start_daemon(None)) { tracing::error!(error = %e, "启动失败"); std::process::exit(1); } } /// 从 CLI `wx daemon start [--tcp ADDR]` 调用 /// /// 查找当前可执行文件路径,设置 WX_DAEMON_MODE=1,后台启动新进程。 /// tracing 已在子进程 main() 中直接写入 daemon.log,无需重定向 stdout/stderr。 pub fn run_start(tcp_addr: Option) -> Result<()> { let exe = std::env::current_exe()?; let mut cmd = std::process::Command::new(&exe); cmd.env("WX_DAEMON_MODE", "1"); if let Some(addr) = &tcp_addr { cmd.env("WX_DAEMON_TCP_ADDR", addr); } #[cfg(unix)] { use std::os::unix::process::CommandExt; unsafe { cmd.pre_exec(|| { libc::setsid(); Ok(()) }) }; } let child = cmd.spawn()?; let pid = child.id(); info!("已启动 daemon 进程 (PID {})", pid); Ok(()) } /// daemon 核心启动逻辑(被 run() 和 WX_DAEMON_MODE 路径共享) #[tracing::instrument(name = "daemon.startup", skip_all)] pub async fn start_daemon(tcp_addr: Option) -> Result<()> { // 确保工作目录存在 let cli_dir = config::cli_dir(); tokio::fs::create_dir_all(&cli_dir).await?; tokio::fs::create_dir_all(config::cache_dir()).await?; let pid = std::process::id(); // 注册 SIGTERM / SIGINT 处理 setup_signal_handler().await; info!("wx-daemon 启动 (PID {})", pid); // 加载配置 let cfg = config::load_config()?; info!(db_dir = %cfg.db_dir.display(), "配置加载完成"); // 加载密钥 let keys_content = tokio::fs::read_to_string(&cfg.keys_file) .await .map_err(|e| anyhow::anyhow!("读取密钥文件 {:?} 失败: {}", cfg.keys_file, e))?; let keys_raw: serde_json::Value = serde_json::from_str(&keys_content)?; let all_keys = extract_keys(&keys_raw); info!("密钥数量: {}", all_keys.len()); // 初始化 DbCache let db = Arc::new(cache::DbCache::new(cfg.db_dir.clone(), all_keys.clone()).await?); // 收集消息 DB 列表 let msg_db_keys: Vec = all_keys .keys() .filter(|k| { let k = k.replace('\\', "/"); k.contains("message/message_") && k.ends_with(".db") && !k.contains("_fts") && !k.contains("_resource") }) .cloned() .collect(); // 预热:加载联系人 + 解密 session.db info!("开始预热..."); let names_raw = query::load_names(&*db).await.unwrap_or_else(|e| { warn!(error = %e, "加载联系人失败,使用空联系人表"); query::Names { map: HashMap::new(), md5_to_uname: HashMap::new(), msg_db_keys: Vec::new(), verify_flags: HashMap::new(), } }); let mut names = names_raw; names.msg_db_keys = msg_db_keys; let _ = db.get("session/session.db").await; let _ = db.get("sns/sns.db").await; info!("预热完成,联系人 {} 个", names.map.len()); // 包一层内部 Arc let names_arc = Arc::new(tokio::sync::RwLock::new(Arc::new(names))); // 检查环境变量中的 TCP 地址(WX_DAEMON_MODE 路径下通过 env 传入) let effective_tcp_addr = tcp_addr.or_else(|| std::env::var("WX_DAEMON_TCP_ADDR").ok()); // 启动 IPC server(阻塞) server::serve(Arc::clone(&db), Arc::clone(&names_arc), effective_tcp_addr.as_deref()).await?; // 正常退出时清理(signal 路径下由 cleanup_and_exit 处理,不会走到这里) cleanup_ipc_files(); Ok(()) } /// 从 all_keys.json 提取 rel_key -> enc_key 映射 /// /// 兼容两种格式: /// - `{ "rel/path.db": { "enc_key": "hex" } }`(Python 版原生格式) /// - `{ "rel/path.db": "hex" }`(简化格式) fn extract_keys(json: &serde_json::Value) -> HashMap { let mut result = HashMap::new(); if let Some(obj) = json.as_object() { for (k, v) in obj { if k.starts_with('_') { continue; } let enc_key = if let Some(s) = v.as_str() { s.to_string() } else if let Some(obj2) = v.as_object() { obj2.get("enc_key") .and_then(|e| e.as_str()) .unwrap_or_default() .to_string() } else { continue; }; if !enc_key.is_empty() { // 统一路径分隔符 let rel = k.replace('\\', "/"); result.insert(rel, enc_key); } } } result } /// 设置信号处理(Unix: SIGTERM/SIGINT) async fn setup_signal_handler() { #[cfg(unix)] tokio::spawn(async move { use tokio::signal::unix::{signal, SignalKind}; let mut term = signal(SignalKind::terminate()).expect("无法监听 SIGTERM"); let mut int = signal(SignalKind::interrupt()).expect("无法监听 SIGINT"); tokio::select! { _ = term.recv() => {}, _ = int.recv() => {}, } cleanup_and_exit(); }); } #[cfg(unix)] fn cleanup_and_exit() { cleanup_ipc_files(); std::process::exit(0); } fn cleanup_ipc_files() { let _ = std::fs::remove_file(config::sock_path()); let _ = std::fs::remove_file(config::pid_path()); }