From 3029c55ae570924ec81f4e5a4e06548241aa21af Mon Sep 17 00:00:00 2001 From: zhangjia Date: Thu, 10 Sep 2026 00:37:34 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20Windows=20CodeBuddy=20CN=20=E4=BC=98?= =?UTF-8?q?=E9=9B=85=E5=85=B3=E9=97=AD=E4=B8=8E=E5=86=99=E5=BA=93=E4=BF=9D?= =?UTF-8?q?=E6=8A=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - process: 新增向指定 PID 可见顶层窗口发送 WM_CLOSE 的辅助（user32/Add-Type） - codebuddy_cn_ide: Windows 关闭改为 WM_CLOSE 优雅退出，残留仍持窗口时提示用户而非直接强杀，避免未保存内容丢失；启动后新增存活校验，避免单例锁秒退仍提示切换成功 - vscode_cn_inject: 写库前等待 DB 释放独占句柄、用 VACUUM INTO 生成一致快照备份、写后解密回读校验并在失败时回滚备份 --- .../src/modules/codebuddy_cn_ide.rs | 84 ++++++++- crates/wb-switch-core/src/modules/process.rs | 119 +++++++++++++ .../src/modules/vscode_cn_inject.rs | 168 +++++++++++++----- 3 files changed, 321 insertions(+), 50 deletions(-) diff --git a/crates/wb-switch-core/src/modules/codebuddy_cn_ide.rs b/crates/wb-switch-core/src/modules/codebuddy_cn_ide.rs index 248fb02c..a779e681 100644 --- a/crates/wb-switch-core/src/modules/codebuddy_cn_ide.rs +++ b/crates/wb-switch-core/src/modules/codebuddy_cn_ide.rs @@ -740,19 +740,48 @@ fn close_codebuddy_cn_windows(timeout_secs: i64) -> Result<(), String> { return Ok(()); } let pids: Vec = rows.iter().map(|r| r.pid).collect(); - for pid in &pids { - let pid_s = pid.to_string(); - let _ = run_cmd("taskkill", &["/PID", &pid_s, "/T"], 10); - } - let started = Instant::now(); let timeout = Duration::from_secs(timeout_secs.max(1) as u64); - let graceful_budget = Duration::from_secs(8).min(timeout); - let remaining = process::wait_windows_pids_gone(&pids, graceful_budget); + + // 1) 优雅阶段：向拥有可见顶层窗口的进程（即 Electron GUI 主进程）发送 + // WM_CLOSE，让 CodeBuddy CN 自行走窗口关闭/保存/退出流程。只影响主进程， + // 渲染子进程会随主进程一并退出，无需逐个处理。 + let notified = process::windows_send_close_to_pids(&pids); + if notified == 0 { + // 未发现任何可见顶层窗口（如仅驻留托盘/后台）：退化为对所有进程执行 + // 不带 /F 的 taskkill——对 GUI 进程而言这等价于发送关闭请求。 + eprintln!("[codebuddy-cn-ide] no visible window found, fallback graceful taskkill…"); + for pid in &pids { + let pid_s = pid.to_string(); + let _ = run_cmd("taskkill", &["/PID", &pid_s, "/T"], 10); + } + } else { + eprintln!("[codebuddy-cn-ide] sent WM_CLOSE to {notified} visible window(s)"); + } + + // 2) 优雅窗口：等待进程全部消失（预算 8s 或剩余时间）。 + let graceful = Duration::from_secs(8).min(timeout); + let remaining = process::wait_windows_pids_gone(&pids, graceful); if remaining.is_empty() { return Ok(()); } + // 3) 保护用户数据：若残留进程仍持有可见窗口（典型场景：IDE 弹了 + // “未保存文件”确认框并阻塞退出），此时不自动强杀，避免丢失编辑内容， + // 改为提示用户保存/关闭后重试。这是与原实现“8s 后无条件 /F”的关键差异。 + let with_window = process::windows_visible_window_pids(&remaining); + if !with_window.is_empty() { + return Err(format!( + "CodeBuddy CN 仍有可见窗口未关闭（PID {}），可能包含未保存内容。请先在 CodeBuddy CN 中保存并关闭窗口后，再重试切换。", + with_window + .iter() + .map(u32::to_string) + .collect::>() + .join(", ") + )); + } + + // 4) 无窗口残留（后台进程）→ 可安全强制结束进程树。 for pid in &remaining { let pid_s = pid.to_string(); let _ = run_cmd("taskkill", &["/PID", &pid_s, "/T", "/F"], 10); @@ -890,6 +919,42 @@ fn launch_codebuddy_cn_macos() -> Result<(), String> { validate_macos_cn_startup(&app) } +/// Windows 启动存活校验：spawn 后轮询 CodeBuddy CN 进程出现并持续存活。 +/// +/// - 从未出现 → 超时 Err； +/// - 出现后持续存活 ≥ `sustain` → Ok； +/// - 出现过但随后消失（单例锁导致秒退的典型特征）→ 立即 Err。 +/// 与 macOS 分支 `validate_macos_cn_startup` 语义对齐，避免“spawn 成功即假成功”。 +#[cfg(target_os = "windows")] +fn validate_windows_cn_startup(exe: &Path) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(30); + let sustain = Duration::from_secs(10); + let mut seen_at: Option = None; + while Instant::now() < deadline { + let alive = !windows_cn_process_rows().is_empty(); + if alive { + match seen_at { + None => seen_at = Some(Instant::now()), + Some(start) => { + if start.elapsed() >= sustain { + return Ok(()); + } + } + } + } else if seen_at.is_some() { + return Err(format!( + "CodeBuddy CN 启动后立即退出（疑似残留单例锁）。请先手动打开一次 CodeBuddy CN（路径: {}）。", + exe.display() + )); + } + std::thread::sleep(Duration::from_millis(500)); + } + Err(format!( + "启动 CodeBuddy CN 超时，未能确认运行（路径: {}）。请先手动打开一次 CodeBuddy CN。", + exe.display() + )) +} + pub fn launch_codebuddy_cn() -> Result<(), String> { #[cfg(target_os = "macos")] { @@ -907,11 +972,14 @@ pub fn launch_codebuddy_cn() -> Result<(), String> { )); } persist_cn_app_cache(&exe); - process::cmd_builder(&exe) + let spawned = process::cmd_builder(&exe) .stdout(std::process::Stdio::null()) .stderr(std::process::Stdio::null()) .spawn() .map_err(|e| format!("启动 CodeBuddy CN 失败: {e}（路径: {}）", exe.display()))?; + // 启动后存活校验：避免单例锁导致秒退却仍提示“切换成功”。 + validate_windows_cn_startup(&exe)?; + drop(spawned); Ok(()) } #[cfg(not(any(target_os = "macos", target_os = "windows")))] diff --git a/crates/wb-switch-core/src/modules/process.rs b/crates/wb-switch-core/src/modules/process.rs index 3306f986..79168071 100644 --- a/crates/wb-switch-core/src/modules/process.rs +++ b/crates/wb-switch-core/src/modules/process.rs @@ -438,6 +438,125 @@ pub(crate) fn existing_windows_drives() -> Vec { .collect() } +/// 生成内联 C#（user32 P/Invoke）PowerShell 脚本的公共常量模板。 +/// +/// 占位符：`__PIDS__`（PowerShell 数组字面量）、`__MODE__`（close/query）。 +/// 脚本单次 Add-Type 编译后调用，仅用于低频的优雅关闭流程。 +#[cfg(target_os = "windows")] +const WM_CLOSE_TEMPLATE: &str = r#" +$ErrorActionPreference = 'SilentlyContinue' +Add-Type -TypeDefinition @" +using System; +using System.Runtime.InteropServices; +using System.Collections.Generic; +public static class WbSwitchWin { + public delegate bool EnumWindowsProc(IntPtr hWnd, IntPtr lParam); + [DllImport("user32.dll")] public static extern bool EnumWindows(EnumWindowsProc cb, IntPtr lParam); + [DllImport("user32.dll")] public static extern uint GetWindowThreadProcessId(IntPtr hWnd, out uint pid); + [DllImport("user32.dll")] public static extern bool IsWindowVisible(IntPtr hWnd); + [DllImport("user32.dll")] public static extern bool PostMessageW(IntPtr hWnd, uint msg, IntPtr wParam, IntPtr lParam); + public static List VisiblePids(HashSet targets) { + var res = new List(); + EnumWindows(delegate(IntPtr h, IntPtr l) { + uint pid; GetWindowThreadProcessId(h, out pid); + if (targets.Contains(pid) && IsWindowVisible(h)) { + if (!res.Contains(pid)) res.Add(pid); + } + return true; + }, IntPtr.Zero); + return res; + } + public static int SendClose(HashSet targets) { + int count = 0; + EnumWindows(delegate(IntPtr h, IntPtr l) { + uint pid; GetWindowThreadProcessId(h, out pid); + if (targets.Contains(pid) && IsWindowVisible(h)) { + PostMessageW(h, 0x0010, IntPtr.Zero, IntPtr.Zero); + count++; + } + return true; + }, IntPtr.Zero); + return count; + } +} +"@ +$targets = New-Object 'System.Collections.Generic.HashSet[uint32]' +__PIDS__ | ForEach-Object { [void]$targets.Add($_) } +if ('__MODE__' -eq 'close') { + [WbSwitchWin]::SendClose($targets) +} else { + ([WbSwitchWin]::VisiblePids($targets) -join ',') +} +"#; + +/// 按脚本模式组装内联 PowerShell 脚本。 +#[cfg(target_os = "windows")] +fn windows_wm_close_script(script_mode: &str, pid_csv: &str) -> String { + let pids = if pid_csv.is_empty() { + "@()".to_string() + } else { + format!( + "@({})", + pid_csv + .split(',') + .filter(|s| !s.trim().is_empty()) + .map(|s| format!("[uint32]{s}")) + .collect::>() + .join(",") + ) + }; + WM_CLOSE_TEMPLATE + .replace("__PIDS__", &pids) + .replace("__MODE__", script_mode) +} + +/// 向给定 PID 的所有可见顶层窗口发送 `WM_CLOSE`，返回实际送达的窗口数。 +/// 用于 GUI 优雅退出：Electron 主进程收到后可自行走保存/退出流程。 +#[cfg(target_os = "windows")] +pub(crate) fn windows_send_close_to_pids(pids: &[u32]) -> usize { + if pids.is_empty() { + return 0; + } + let csv = pids + .iter() + .map(u32::to_string) + .collect::>() + .join(","); + let script = windows_wm_close_script("close", &csv); + match run_cmd_timeout("powershell", &["-NoProfile", "-NonInteractive", "-Command", &script], 30) + { + Some(out) => String::from_utf8_lossy(&out.stdout) + .trim() + .parse::() + .unwrap_or(0), + None => 0, + } +} + +/// 返回给定 PID 中仍拥有可见顶层窗口的 PID 子集。 +/// 供优雅关闭后判断：若目标进程仍有可见窗口（如未保存确认框）， +/// 应提示用户手动处理，而不是直接强杀。 +#[cfg(target_os = "windows")] +pub(crate) fn windows_visible_window_pids(pids: &[u32]) -> Vec { + if pids.is_empty() { + return Vec::new(); + } + let csv = pids + .iter() + .map(u32::to_string) + .collect::>() + .join(","); + let script = windows_wm_close_script("query", &csv); + let stdout = match ps_output(&script, 30) { + Some(out) => out, + None => return Vec::new(), + }; + stdout + .split(',') + .filter_map(|s| s.trim().parse::().ok()) + .collect() +} + #[cfg(target_os = "windows")] fn windows_running_workbuddy_exe() -> Option { let script = "Get-Process -Name WorkBuddy,CodeBuddy -ErrorAction SilentlyContinue | \ diff --git a/crates/wb-switch-core/src/modules/vscode_cn_inject.rs b/crates/wb-switch-core/src/modules/vscode_cn_inject.rs index 0679f811..fefe5f9f 100644 --- a/crates/wb-switch-core/src/modules/vscode_cn_inject.rs +++ b/crates/wb-switch-core/src/modules/vscode_cn_inject.rs @@ -9,6 +9,7 @@ //! - Linux: secret-tool / peanuts 固定密钥 → AES-128-CBC `v11`/`v10` use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant}; #[cfg(not(target_os = "windows"))] use aes::Aes128; @@ -479,6 +480,76 @@ pub fn read_codebuddy_cn_secret(user_data_dir: Option<&Path>) -> Result Result<(), String> { + if !db_path.exists() { + return Ok(()); + } + let deadline = Instant::now() + timeout; + loop { + match std::fs::OpenOptions::new() + .read(true) + .write(true) + .open(db_path) + { + Ok(_) => return Ok(()), + Err(_) if Instant::now() < deadline => { + std::thread::sleep(Duration::from_millis(200)); + } + Err(e) => { + return Err(format!( + "等待 state.vscdb 可访问超时（可能仍被 CodeBuddy CN 占用，请先完全退出 IDE）: {e}" + )); + } + } + } +} + +/// 用 SQLite `VACUUM INTO` 生成一致快照备份（比文件复制更安全，天然包含 WAL 内容）。 +/// 返回备份文件路径。 +fn snapshot_backup(conn: &Connection) -> Result { + let dir = crate::modules::config::backup_dir(); + std::fs::create_dir_all(&dir).map_err(|e| format!("创建备份目录失败: {e}"))?; + let backup = dir.join(format!( + "state.vscdb.{}.bak", + crate::modules::config::now_ms() + )); + let _ = std::fs::remove_file(&backup); + let escaped = backup.to_string_lossy().replace('\'', "''"); + conn.execute_batch(&format!("VACUUM INTO '{escaped}'")) + .map_err(|e| format!("生成 state.vscdb 备份失败: {e}"))?; + if !backup.exists() { + return Err("生成 state.vscdb 备份失败：备份文件未生成".to_string()); + } + Ok(backup) +} + +/// 校验：解密回读刚写入的 secret，必须与目标明文一致。 +fn verify_written_secret( + conn: &Connection, + db_key: &str, + data_root: &Path, + expected: &str, +) -> Result<(), String> { + let written: String = conn + .query_row( + "SELECT value FROM ItemTable WHERE key = ?", + [db_key], + |row| row.get(0), + ) + .map_err(|e| format!("写后回读失败: {e}"))?; + let plain = decode_secret_storage_value(&written, data_root) + .map_err(|e| format!("写后解密回读失败: {e}"))?; + if plain != expected { + return Err("写后校验失败：解密回读内容与目标账号不一致".to_string()); + } + Ok(()) +} + /// 加密并写入 CodeBuddy CN secret。 pub fn inject_codebuddy_cn_secret( plaintext: &str, @@ -490,56 +561,69 @@ pub fn inject_codebuddy_cn_secret( std::fs::create_dir_all(parent) .map_err(|e| format!("创建 state.vscdb 父目录失败: {e}"))?; } + + // 1) 句柄握手：等待 DB 不被其它进程独占（Windows 强杀后句柄释放有延迟）。 + wait_for_db_released(&db_path, Duration::from_secs(5))?; + let conn = Connection::open(&db_path).map_err(|e| format!("打开 state.vscdb 失败: {e}"))?; - conn.execute( - "CREATE TABLE IF NOT EXISTS ItemTable (key TEXT PRIMARY KEY, value TEXT)", - [], - ) - .map_err(|e| format!("初始化 ItemTable 失败: {e}"))?; + conn.busy_timeout(Duration::from_secs(5)) + .map_err(|e| format!("设置 busy_timeout 失败: {e}"))?; - let db_key = secret_storage_item_key(); - let existing_prefix: Option = match conn.query_row( - "SELECT value FROM ItemTable WHERE key = ?", - [db_key.as_str()], - |row| row.get::<_, String>(0), - ) { - Ok(val) => { - if let Ok(parsed) = serde_json::from_str::(&val) { - if let Ok(bytes) = decode_buffer_data(&parsed) { - detect_prefix(&bytes).map(|s| s.to_string()) + // 2) 写前生成一致快照备份；失败回滚用。 + let backup = snapshot_backup(&conn)?; + + let db_path_for_result = db_path.clone(); + let result: Result = (|| { + conn.execute( + "CREATE TABLE IF NOT EXISTS ItemTable (key TEXT PRIMARY KEY, value TEXT)", + [], + ) + .map_err(|e| format!("初始化 ItemTable 失败: {e}"))?; + + let db_key = secret_storage_item_key(); + let existing_prefix: Option = match conn.query_row( + "SELECT value FROM ItemTable WHERE key = ?", + [db_key.as_str()], + |row| row.get::<_, String>(0), + ) { + Ok(val) => { + if let Ok(parsed) = serde_json::from_str::(&val) { + if let Ok(bytes) = decode_buffer_data(&parsed) { + detect_prefix(&bytes).map(|s| s.to_string()) + } else { + None + } } else { None } - } else { - None } - } - Err(_) => None, - }; - - let encrypted = - encrypt_secret_payload(plaintext.as_bytes(), existing_prefix.as_deref(), &data_root)?; - let buffer_str = encode_secret_buffer(encrypted)?; - conn.execute( - "INSERT OR REPLACE INTO ItemTable (key, value) VALUES (?, ?)", - rusqlite::params![db_key, buffer_str], - ) - .map_err(|e| format!("写入 state.vscdb 失败: {e}"))?; + Err(_) => None, + }; - // 写后校验：行存在且为 Buffer JSON - let written: String = conn - .query_row( - "SELECT value FROM ItemTable WHERE key = ?", - [db_key.as_str()], - |row| row.get(0), + let encrypted = + encrypt_secret_payload(plaintext.as_bytes(), existing_prefix.as_deref(), &data_root)?; + let buffer_str = encode_secret_buffer(encrypted)?; + conn.execute( + "INSERT OR REPLACE INTO ItemTable (key, value) VALUES (?, ?)", + rusqlite::params![db_key, buffer_str], ) - .map_err(|e| format!("写后校验失败: {e}"))?; - let parsed: serde_json::Value = - serde_json::from_str(&written).map_err(|e| format!("写后校验 JSON 失败: {e}"))?; - if parsed.get("type").and_then(|v| v.as_str()) != Some("Buffer") { - return Err("写后校验失败：value 不是 Buffer".to_string()); - } - Ok(db_path) + .map_err(|e| format!("写入 state.vscdb 失败: {e}"))?; + + // 3) 写后解密回读，确认写入内容与目标账号一致。 + verify_written_secret(&conn, &db_key, &data_root, plaintext) + .map_err(|e| { + let _ = std::fs::copy(&backup, &db_path_for_result); + e + })?; + Ok(db_path_for_result) + })(); + + if result.is_err() { + // 兜底：确保连接已关闭后再尝试恢复备份。 + drop(conn); + let _ = std::fs::copy(&backup, &db_path); + } + result } #[cfg(test)]