wx-cli/src/daemon/mod.rs

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<String>) -> 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<String>) -> 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<String> = 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<String, String> {
    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());
}