use anyhow::{bail, Context, Result}; use std::io::{BufRead, BufReader, Write}; use std::net::TcpStream; use std::time::Duration; use crate::config; use crate::ipc::{Request, Response}; const STARTUP_TIMEOUT_SECS: u64 = 15; const TCP_CONNECT_TIMEOUT_SECS: u64 = 15; const TCP_RW_TIMEOUT_SECS: u64 = 120; /// 检查 daemon 是否存活 pub fn is_alive(tcp_addr: Option<&str>) -> bool { if let Some(addr) = tcp_addr { return is_alive_tcp(addr); } #[cfg(unix)] { use std::os::unix::net::UnixStream; let sock_path = config::sock_path(); if !sock_path.exists() { return false; } let mut stream = match UnixStream::connect(&sock_path) { Ok(s) => s, Err(_) => return false, }; stream.set_read_timeout(Some(Duration::from_secs(2))).ok(); stream.set_write_timeout(Some(Duration::from_secs(2))).ok(); let req = serde_json::json!({"cmd": "ping"}); if write!(stream, "{}\n", req).is_err() { return false; } let mut line = String::new(); let mut reader = BufReader::new(&stream); if reader.read_line(&mut line).is_err() { return false; } serde_json::from_str::(&line) .ok() .and_then(|v| v.get("pong").and_then(|p| p.as_bool())) .unwrap_or(false) } #[cfg(windows)] { use interprocess::local_socket::{prelude::*, GenericNamespaced, Stream}; // 必须用 interprocess 自己的连接 API,和 server 保持一致 match "wx-cli-daemon".to_ns_name::() { Ok(name) => Stream::connect(name).is_ok(), Err(_) => false, } } #[cfg(not(any(unix, windows)))] { false } } /// TCP liveness check: send ping via TCP, return true if pong received pub fn is_alive_tcp(addr: &str) -> bool { let tcp_addr = match addr.parse() { Ok(a) => a, Err(_) => return false, }; let mut stream = match TcpStream::connect_timeout( &tcp_addr, Duration::from_secs(TCP_CONNECT_TIMEOUT_SECS), ) { Ok(s) => s, Err(_) => return false, }; let _ = stream.set_read_timeout(Some(Duration::from_secs(2))); let _ = stream.set_write_timeout(Some(Duration::from_secs(2))); let req = serde_json::json!({"cmd": "ping"}); if write!(stream, "{}\n", req).is_err() { return false; } let mut reader = BufReader::new(stream); let mut line = String::new(); if reader.read_line(&mut line).is_err() { return false; } serde_json::from_str::(&line) .ok() .and_then(|v| v.get("pong").and_then(|p| p.as_bool())) .unwrap_or(false) } /// 确保 daemon 运行,必要时自动启动 /// 当指定 tcp_addr 时,不会自动启动 daemon(用户显式选择了 TCP 模式) pub fn ensure_daemon(tcp_addr: Option<&str>) -> Result<()> { if is_alive(tcp_addr) { return Ok(()); } // TCP 模式下不自动启动 daemon,直接报错 if tcp_addr.is_some() { let addr = tcp_addr.unwrap(); bail!( "无法连接到 TCP daemon ({}):{}\n请确认 daemon 已通过 `wx daemon start --tcp {}` 启动", addr, std::io::Error::last_os_error(), addr, ); } eprintln!("启动 wx-daemon..."); start_daemon()?; Ok(()) } /// 启动 daemon 进程(自身二进制,设置 WX_DAEMON_MODE=1) /// /// tracing 已在子进程 main() 中直接写入 daemon.log,无需重定向 stdout/stderr。 fn start_daemon() -> Result<()> { let exe = std::env::current_exe().context("无法获取当前可执行文件路径")?; // 预检:当前用户是否能写 ~/.wx-cli/。如果不能,给出可操作的错误信息, // 而不是 spawn 一个注定失败的 daemon 然后超时 15s。 preflight_cli_dir_writable()?; #[cfg(unix)] { use std::os::unix::process::CommandExt; let mut cmd = std::process::Command::new(&exe); cmd.env("WX_DAEMON_MODE", "1") .stdin(std::process::Stdio::null()) .stdout(std::process::Stdio::null()) .stderr(std::process::Stdio::null()); // SAFETY: setsid() 在 fork 后的子进程中调用,使 daemon 脱离控制终端 unsafe { cmd.pre_exec(|| { libc::setsid(); Ok(()) }); } let _ = cmd.spawn().context("无法启动 daemon 进程")?; } #[cfg(windows)] { use std::os::windows::process::CommandExt; let _ = std::process::Command::new(&exe) .env("WX_DAEMON_MODE", "1") .stdin(std::process::Stdio::null()) .stdout(std::process::Stdio::null()) .stderr(std::process::Stdio::null()) .creation_flags(0x00000008) // DETACHED_PROCESS .spawn() .context("无法启动 daemon 进程")?; } let probe = cli_dir.join(".daemon_probe"); match std::fs::File::create(&probe) { Ok(_) => { let _ = std::fs::remove_file(&probe); Ok(()) } Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => { let dir = cli_dir.display(); if cfg!(unix) { bail!( "无法写入 {dir}(权限不足)\n\n\ 这通常是老版本的 `sudo wx init` 把目录属主留成了 root。\n\ 修复:\n\n \ sudo chown -R $(whoami) {dir}\n\n\ (新版已修复此问题,下次 init 不会再发生)", ) } else { bail!("无法写入 {dir}: {e}") } } Err(e) => bail!("无法写入 {}: {}", cli_dir.display(), e), } } /// 向 daemon 发送请求并返回响应 pub fn send(req: Request, tcp_addr: Option<&str>) -> Result { if let Some(addr) = tcp_addr { return send_tcp(req, addr); } ensure_daemon(None)?; #[cfg(unix)] { send_unix(req) } #[cfg(windows)] { send_windows(req) } #[cfg(not(any(unix, windows)))] { bail!("不支持当前平台") } } /// 通过 TCP 发送请求并返回响应 pub fn send_tcp(req: Request, addr: &str) -> Result { let mut stream = TcpStream::connect_timeout( &addr.parse().context("TCP 地址格式无效")?, Duration::from_secs(TCP_CONNECT_TIMEOUT_SECS), ) .context(format!("连接 TCP daemon ({}) 失败", addr))?; stream .set_read_timeout(Some(Duration::from_secs(TCP_RW_TIMEOUT_SECS))) .ok(); stream .set_write_timeout(Some(Duration::from_secs(TCP_RW_TIMEOUT_SECS))) .ok(); let req_str = serde_json::to_string(&req)? + "\n"; stream.write_all(req_str.as_bytes())?; let mut line = String::new(); let mut reader = BufReader::new(&stream); reader.read_line(&mut line)?; let resp: Response = serde_json::from_str(&line) .context("解析 daemon 响应失败")?; if !resp.ok { bail!("{}", resp.error.as_deref().unwrap_or("未知错误")); } Ok(resp) } #[cfg(unix)] fn send_unix(req: Request) -> Result { use std::os::unix::net::UnixStream; let sock_path = config::sock_path(); let mut stream = UnixStream::connect(&sock_path) .context("连接 daemon socket 失败")?; stream.set_read_timeout(Some(Duration::from_secs(120))).ok(); stream.set_write_timeout(Some(Duration::from_secs(120))).ok(); let req_str = serde_json::to_string(&req)? + "\n"; stream.write_all(req_str.as_bytes())?; let mut line = String::new(); let mut reader = BufReader::new(&stream); reader.read_line(&mut line)?; let resp: Response = serde_json::from_str(&line) .context("解析 daemon 响应失败")?; if !resp.ok { bail!("{}", resp.error.as_deref().unwrap_or("未知错误")); } Ok(resp) } #[cfg(windows)] fn send_windows(req: Request) -> Result { use interprocess::local_socket::{prelude::*, GenericNamespaced, Stream}; let name = "wx-cli-daemon".to_ns_name::() .context("构造 pipe name 失败")?; let stream = Stream::connect(name) .context("连接 daemon named pipe 失败")?; // interprocess::Stream 同时实现 Read + Write,但需要拆分读写端 let mut reader = BufReader::new(stream); let req_str = serde_json::to_string(&req)? + "\n"; reader.get_mut().write_all(req_str.as_bytes())?; let mut line = String::new(); reader.read_line(&mut line)?; let resp: Response = serde_json::from_str(&line) .context("解析 daemon 响应失败")?; if !resp.ok { bail!("{}", resp.error.as_deref().unwrap_or("未知错误")); } Ok(resp) } #[cfg(test)] mod integration_tests { use super::*; use crate::ipc::{Request, Response}; use serde_json::json; use std::net::SocketAddr; use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; use tokio::net::TcpListener; /// Spawn a mock TCP server that responds to one request with the given JSON data. /// Returns the bound address (with the actual random port). async fn spawn_mock_server(response_body: serde_json::Value) -> SocketAddr { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); let (reader, mut writer) = stream.into_split(); // Read one line (the request) let mut buf_reader = tokio::io::BufReader::new(reader); let mut line = String::new(); buf_reader.read_line(&mut line).await.unwrap(); // Write response as a JSON line let resp = Response { ok: true, error: None, data: response_body, }; let resp_str = serde_json::to_string(&resp).unwrap() + "\n"; writer.write_all(resp_str.as_bytes()).await.unwrap(); writer.shutdown().await.unwrap(); }); addr } #[tokio::test(flavor = "multi_thread")] async fn test_send_tcp_round_trip() { let addr = spawn_mock_server(json!({ "sessions": [{"name": "test"}] })) .await; let resp = send_tcp(Request::Sessions { limit: 20 }, &addr.to_string()).unwrap(); assert!(resp.ok, "Response should be ok"); } #[tokio::test(flavor = "multi_thread")] async fn test_send_tcp_connection_refused() { // Port 59876 is very unlikely to have a listener let result = send_tcp(Request::Sessions { limit: 20 }, "127.0.0.1:59876"); assert!(result.is_err(), "Expected connection refused error"); } #[tokio::test(flavor = "multi_thread")] async fn test_is_alive_tcp_false() { // Port 59877 is very unlikely to have a listener let result = is_alive_tcp("127.0.0.1:59877"); assert!(!result, "Expected is_alive_tcp to return false for unused port"); } } /// Real TCP daemon integration tests — spawn the actual `wx` daemon binary, /// connect via TCP, and verify end-to-end request/response round-trip. /// /// These tests are `#[cfg(unix)]` only and require the `wx` binary to have /// been built with `cargo build --bin wx`. #[cfg(unix)] #[cfg(test)] mod tcp_integration_tests { use super::*; use crate::ipc::Request; use std::process::Command; /// Build the `wx` binary so the daemon subprocess is available. fn ensure_binary() -> std::path::PathBuf { let status = Command::new("cargo") .args(["build", "--bin", "wx"]) .output() .expect("cargo build failed to execute"); if !status.status.success() { panic!( "cargo build --bin wx failed:\n{}", String::from_utf8_lossy(&status.stderr) ); } // Binary path: target/debug/wx let mut p = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")); p.push("target/debug/wx"); assert!(p.exists(), "wx binary not found at {:?}", p); p } /// Wait for the daemon TCP endpoint to become ready. fn wait_for_tcp_ready(addr: &str) -> bool { let deadline = std::time::Instant::now() + std::time::Duration::from_secs(STARTUP_TIMEOUT_SECS); while std::time::Instant::now() < deadline { std::thread::sleep(std::time::Duration::from_millis(300)); if is_alive_tcp(addr) { return true; } } false } #[test] fn test_tcp_daemon_ping_round_trip() { let binary = ensure_binary(); // Pick a free ephemeral port let port = { let listener = std::net::TcpListener::bind("127.0.0.1:0") .expect("failed to bind ephemeral port"); listener.local_addr().unwrap().port() }; let addr = format!("127.0.0.1:{}", port); // Spawn the daemon subprocess in TCP-only mode let mut child = Command::new(&binary) .env("WX_DAEMON_MODE", "1") .env("WX_DAEMON_TCP_ADDR", &addr) .stdin(std::process::Stdio::null()) .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::piped()) .spawn() .expect("failed to spawn wx daemon"); let pid = child.id(); eprintln!("[test] spawned daemon PID {}", pid); // Wait for TCP readiness if !wait_for_tcp_ready(&addr) { let _ = child.kill(); let _ = child.wait(); panic!( "daemon did not become ready on {} within {}s (PID {})", addr, STARTUP_TIMEOUT_SECS, pid ); } eprintln!("[test] daemon ready on {}", addr); // Send Ping request and verify pong let resp = send_tcp(Request::Ping, &addr) .expect("send_tcp(Ping) should succeed"); assert!(resp.ok, "Response ok flag should be true"); let pong = resp.data.get("pong").and_then(|v| v.as_bool()); assert!( pong == Some(true), "Expected pong=true in response, got: {:?}", resp.data ); // Terminate daemon unsafe { libc::kill(pid as libc::pid_t, libc::SIGTERM) }; // Verify clean exit let exit_status = child.wait().expect("failed to wait on daemon"); assert!( exit_status.success(), "daemon should exit cleanly, got: {:?}", exit_status ); } #[test] fn test_tcp_daemon_connection_refused() { // Port 59889 is very unlikely to have a listener let addr = "127.0.0.1:59889"; let result = send_tcp(Request::Ping, addr); assert!( result.is_err(), "Expected connection refused error when no daemon is listening on {}", addr ); } /// Compare TCP and local transport responses for the same query. /// Marked `#[ignore]` because it requires WeChat data to be present on the machine. /// Run manually: `cargo test -- --ignored test_tcp_matches_local_sessions` #[test] #[ignore] fn test_tcp_matches_local_sessions() { let binary = ensure_binary(); // Pick a free ephemeral port for TCP daemon let port = { let listener = std::net::TcpListener::bind("127.0.0.1:0") .expect("failed to bind ephemeral port"); listener.local_addr().unwrap().port() }; let tcp_addr = format!("127.0.0.1:{}", port); // --- Phase 1: Query via TCP --- let mut tcp_child = Command::new(&binary) .env("WX_DAEMON_MODE", "1") .env("WX_DAEMON_TCP_ADDR", &tcp_addr) .stdin(std::process::Stdio::null()) .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::piped()) .spawn() .expect("failed to spawn wx daemon (TCP)"); let tcp_pid = tcp_child.id(); eprintln!("[test] TCP daemon PID {}", tcp_pid); if !wait_for_tcp_ready(&tcp_addr) { let _ = tcp_child.kill(); let _ = tcp_child.wait(); panic!("TCP daemon did not become ready within {}s", STARTUP_TIMEOUT_SECS); } eprintln!("[test] TCP daemon ready on {}", tcp_addr); let tcp_resp = send_tcp(Request::Sessions { limit: 20 }, &tcp_addr) .expect("send_tcp(Sessions) should succeed"); // Terminate TCP daemon unsafe { libc::kill(tcp_pid as libc::pid_t, libc::SIGTERM) }; let _ = tcp_child.wait(); eprintln!("[test] TCP daemon terminated"); // --- Phase 2: Query via local transport --- // send() with tcp_addr=None will auto-start a daemon on the Unix socket let local_resp = send(Request::Sessions { limit: 20 }, None) .expect("send(Sessions) via local transport should succeed"); // --- Phase 3: Deep-compare responses --- let tcp_data = serde_json::to_value(&tcp_resp.data) .expect("tcp_resp.data should be serializable"); let local_data = serde_json::to_value(&local_resp.data) .expect("local_resp.data should be serializable"); assert_eq!( tcp_data, local_data, "TCP and local transport responses differ!\nTCP: {}\nLocal: {}", serde_json::to_string_pretty(&tcp_data).unwrap(), serde_json::to_string_pretty(&local_data).unwrap(), ); eprintln!("[test] TCP and local responses match ✓"); } }