use serde::{Deserialize, Serialize}; use std::io::Read; use std::fs::{self, File}; use std::io::Write as IoWrite; use std::path::PathBuf; use std::pin::Pin; use std::process::Command; use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; use cube_lib::websocket::{WebSocketClient, WebSocketConfig}; // ===================== 文件下载状态 ===================== struct DownloadState { filename: String, offset: u64, temp_path: Option, file: Option, } impl Default for DownloadState { fn default() -> Self { Self { filename: String::new(), offset: 0, temp_path: None, file: None, } } } static DOWNLOAD_STATE: std::sync::Mutex = std::sync::Mutex::new(DownloadState { filename: String::new(), offset: 0, temp_path: None, file: None, }); // ===================== 版本比较与下载 ===================== /// 获取 Updater 数据目录(X:\AppData\,存放 BootLoader.exe 等) fn get_updater_data_dir() -> PathBuf { let exe_path = std::env::current_exe().expect("Failed to get executable path"); let drive = exe_path .parent() .and_then(|p| p.as_os_str().to_str()) .and_then(|s| s.split('\\').next()) .unwrap_or("C:"); let appdata = PathBuf::from(format!("{}/AppData", drive)); let _ = fs::create_dir_all(&appdata); appdata } /// 获取本地文件版本号(使用 PowerShell 获取 PE 文件版本信息) fn get_local_file_version(filename: &str) -> String { let file_path = get_updater_data_dir().join(filename); if !file_path.exists() { return "0.0.0".to_string(); } #[cfg(windows)] { let path_str = file_path.to_string_lossy().to_string(); let ps_script = format!( "(Get-Item -LiteralPath '{}' -ErrorAction SilentlyContinue).VersionInfo.FileVersion", path_str.replace("'", "''") ); let output = Command::new("powershell") .args(["-NoProfile", "-NonInteractive", "-Command", &ps_script]) .output(); if let Ok(output) = output { let stdout = String::from_utf8_lossy(&output.stdout); let version = stdout.trim(); if !version.is_empty() && version != "0" && !version.contains("没有版本") { return version.to_string(); } } // 备用:使用文件修改时间作为"版本" if let Ok(metadata) = fs::metadata(&file_path) { use std::time::UNIX_EPOCH; if let Ok(modified) = metadata.modified() { let duration = modified.duration_since(UNIX_EPOCH).unwrap_or_default(); return format!("{}.{}", duration.as_secs(), duration.subsec_nanos() / 1_000_000); } } "0.0.0".to_string() } #[cfg(not(windows))] { "0.0.0".to_string() } } /// 比较两个版本号(a < b 返回 true) fn version_less_than(a: &str, b: &str) -> bool { let parse_v = |v: &str| -> Vec { v.split(|c: char| !c.is_ascii_digit() && c != '.') .filter(|s| !s.is_empty()) .filter_map(|s| s.parse().ok()) .collect() }; let a_parts: Vec = parse_v(a); let b_parts: Vec = parse_v(b); let max_len = a_parts.len().max(b_parts.len()); for i in 0..max_len { let av = if i < a_parts.len() { a_parts[i] } else { 0 }; let bv = if i < b_parts.len() { b_parts[i] } else { 0 }; if av < bv { return true; } if av > bv { return false; } } false } /// 计算本地文件前 N 字节的 MD5 hash(字节数为 0 表示全部) fn compute_file_hash(filename: &str, bytes: u64, _debug: bool) -> Option { let file_path = get_updater_data_dir().join(filename); if !file_path.exists() { return None; } let mut file = File::open(&file_path).ok()?; let file_size = file.metadata().ok()?.len(); // 如果 bytes=0 或超过文件大小,则计算整个文件 let read_bytes = if bytes == 0 || bytes > file_size { file_size } else { bytes }; // 读取前 read_bytes 字节 let mut buffer = vec![0u8; read_bytes as usize]; file.read_exact(&mut buffer).ok()?; let hash = md5::compute(&buffer); Some(format!("{:x}", hash)) } /// 获取临时文件的当前大小(字节数) fn get_tmp_file_size(filename: &str) -> u64 { let tmp_path = get_updater_data_dir().join(format!("{}.tmp", filename)); tmp_path.metadata().map(|m| m.len()).unwrap_or(0) } /// 发送 GetFileMd5 请求 fn request_file_md5(sender: &cube_lib::websocket::MessageSender, filename: &str, bytes: u64, debug: bool) { let msg_str = format!( r#"{{"Type":"GetFileMd5","Data":{{"filename":"{}","bytes":{}}}}}"#, filename, bytes ); if debug { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} 发送消息:{}", ts, msg_str); } sender.send(msg_str); } /// 发送文件下载请求(断点续传) fn request_download(sender: &cube_lib::websocket::MessageSender, filename: &str, offset: u64, debug: bool) { let msg_str = format!( r#"{{"Type":"DownloadFile","Data":{{"filename":"{}","offset":{}}}}}"#, filename, offset ); if debug { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} 发送消息:{}", ts, msg_str); } sender.send(msg_str); } /// 处理收到的文件块 fn handle_file_chunk(data: &serde_json::Map, debug: bool) { let filename = data.get("filename").and_then(|v| v.as_str()).unwrap_or(""); let offset = data.get("offset").and_then(|v| v.as_u64()).unwrap_or(0); let chunk_data = data.get("data").and_then(|v| v.as_str()).unwrap_or(""); let is_last = data.get("is_last").and_then(|v| v.as_bool()).unwrap_or(false); let mut state = DOWNLOAD_STATE.lock().unwrap(); // 如果是新文件或 offset=0,重置状态 if state.filename != filename || offset == 0 { // 关闭旧文件 if let Some(f) = state.file.take() { drop(f); } // 删除旧临时文件 if let Some(ref tp) = state.temp_path { let _ = fs::remove_file(tp); } // 使用 AppData/Updater/ 作为数据目录 let data_dir = get_updater_data_dir(); let temp_path = data_dir.join(format!("{}.tmp", filename)); // 创建临时文件(截断) match File::create(&temp_path) { Ok(file) => { state.filename = filename.to_string(); state.offset = 0; state.temp_path = Some(temp_path); state.file = Some(file); } Err(_) => { if debug { eprintln!("[下载] 无法创建临时文件: {}", filename); } return; } } } // 解码并追加数据 let current_offset = state.offset; if offset == current_offset { if let Some(ref mut file) = state.file { match BASE64.decode(chunk_data) { Ok(decoded) => { match file.write_all(&decoded) { Ok(_) => { let _ = file.flush(); state.offset += decoded.len() as u64; } Err(e) => { if debug { eprintln!("[下载] 写入失败: {}", e); } } } } Err(e) => { if debug { eprintln!("[下载] Base64 解码失败: {}", e); } } } } } else { if debug { eprintln!("[下载] 偏移不匹配: 期望 {}, 收到 {}", current_offset, offset); } } // 最后一块:重命名临时文件为正式文件 if is_last { let temp_path = state.temp_path.clone(); let final_offset = state.offset; // 关闭文件句柄 if let Some(f) = state.file.take() { drop(f); } if let Some(ref temp) = temp_path { let data_dir = get_updater_data_dir(); let final_path = data_dir.join(filename); // 原子重命名 if let Err(e) = fs::rename(temp, &final_path) { if debug { eprintln!("[下载] 重命名失败: {}", e); } // 尝试直接覆盖写入 let _ = fs::copy(temp, &final_path); let _ = fs::remove_file(temp); } if debug { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [下载] {} 下载完成,共 {} 字节", ts, filename, final_offset); } } // 重置状态 state.filename = String::new(); state.offset = 0; state.temp_path = None; state.file = None; } } /// 处理下载完成消息 fn handle_download_complete(data: &serde_json::Map, debug: bool) { let filename = data.get("filename").and_then(|v| v.as_str()).unwrap_or(""); let size = data.get("size").and_then(|v| v.as_u64()).unwrap_or(0); if debug { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [下载] {} 下载完成,最终大小: {} 字节", ts, filename, size); } let mut state = DOWNLOAD_STATE.lock().unwrap(); state.filename = String::new(); state.offset = 0; state.temp_path = None; state.file = None; } /// Updater 自身配置(AppData/Updater/config.json) /// 只负责 Updater 自己的行为参数,连接地址从公共 config.json 加载 #[derive(Debug, Serialize, Deserialize)] struct Config { /// 调试模式:true 时保留控制台窗口并输出日志 debug_mode: bool, } impl Default for Config { fn default() -> Self { Self { debug_mode: false } } } /// 获取 Updater 自身配置路径 AppData/Updater/config.json fn get_updater_config_path() -> PathBuf { get_updater_data_dir().join("Updater").join("config.json") } /// 获取公共配置路径 AppData/config.json(与 BootLoader 同级) fn get_public_config_path() -> PathBuf { let exe_path = std::env::current_exe().expect("Failed to get executable path"); let drive = exe_path .parent() .and_then(|p| p.as_os_str().to_str()) .and_then(|s| s.split('\\').next()) .unwrap_or("C:"); PathBuf::from(format!("{}/AppData/config.json", drive)) } /// 加载 Updater 自身配置;若文件不存在则写入默认值 fn load_updater_config() -> Config { let config_path = get_updater_config_path(); if config_path.exists() { if let Ok(content) = fs::read_to_string(&config_path) { if let Ok(config) = serde_json::from_str::(&content) { return config; } } } // 文件不存在或解析失败 → 写入默认值 let default_config = Config::default(); if let Ok(content) = serde_json::to_string_pretty(&default_config) { let _ = fs::write(&config_path, content); } default_config } /// 从公共 config.json 读取 ServerUrl 字段 fn resolve_ws_url() -> String { let config_path = get_public_config_path(); if let Ok(content) = fs::read_to_string(&config_path) { if let Ok(json) = serde_json::from_str::(&content) { if let Some(url) = json.get("ServerUrl").and_then(|v| v.as_str()) { return url.to_string(); } } } // 读取失败 → 降级到默认值 "ws://127.0.0.1:8087/ws".to_string() } /// 从公共 config.json 读取 DeviceNumber 字段 fn resolve_device_number() -> String { let config_path = get_public_config_path(); if let Ok(content) = fs::read_to_string(&config_path) { if let Ok(json) = serde_json::from_str::(&content) { // 尝试多个可能的字段名 if let Some(id) = json.get("DeviceNumber").and_then(|v| v.as_str()) { if !id.is_empty() { return id.to_string(); } } if let Some(id) = json.get("StationId").and_then(|v| v.as_str()) { if !id.is_empty() { return id.to_string(); } } if let Some(id) = json.get("Station").and_then(|v| v.as_str()) { if !id.is_empty() { return id.to_string(); } } } } // 读取失败 → 降级到默认值 "UNKNOWN".to_string() } fn is_process_running(process_name: &str) -> bool { use std::process::id; let current_pid = id().to_string(); let output = Command::new("tasklist") .args(["/FI", &format!("IMAGENAME eq {}", process_name), "/FO", "CSV"]) .output() .expect("Failed to execute tasklist"); let output_str = String::from_utf8_lossy(&output.stdout); let lines: Vec<&str> = output_str.lines().collect(); let mut count = 0; for line in lines { if line.contains(&format!("\"{}\"", process_name)) && !line.contains(¤t_pid) { count += 1; } } count > 0 } /// 运行 Updater(使用 CubeLib 内置的自动重连) async fn run_updater(debug_mode: bool) { // 加载初始 URL let server_url = resolve_ws_url(); if debug_mode { println!("========================================"); println!("Updater 启动 (调试模式)"); println!("服务器地址: {}", server_url); println!("自动重连: 启用 (指数退避: 1s - 30s)"); println!("========================================"); } // 创建 WebSocket 配置(启用自动重连) let config = WebSocketConfig::new(&server_url) .with_client_type("Updater") .with_debug(debug_mode) .with_reconnect(true) // 启用自动重连 .with_reconnect_delay(1000) // 初始延迟 1s .with_max_reconnect_delay(30000); // 最大延迟 30s // 创建 WebSocket 客户端 let mut client = WebSocketClient::new(config); // 设置连接成功回调 let debug_connected = debug_mode; client.on_connected(move |url| { if debug_connected { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} 收到消息:{}", ts, url); } }); // 设置消息接收回调 let debug_msg = debug_mode; let device_number = resolve_device_number(); client.on_message(move |msg_type, data, sender| { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); if debug_msg { // 收到消息日志:data 是完整 parsed 对象,取其中的 Data 字段打印 let actual_data = data.get("Data").unwrap_or(&serde_json::Value::Null); let data_str = serde_json::to_string(actual_data).unwrap_or_else(|_| "{}".to_string()); println!("{} 收到消息:{{\"Type\":{},\"Data\":{}}}", ts, serde_json::to_string(&msg_type).unwrap_or_default(), data_str ); } // 收到 welcome 后,发送 GetFileVer if msg_type == "welcome" { if !device_number.is_empty() && device_number != "UNKNOWN" { let msg_str = format!( r#"{{"Type":"GetFileVer","Data":{{"DeviceNumber":"{}","file_list":["BootLoader.exe"]}}}}"#, device_number ); if debug_msg { let ts2 = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} 发送消息:{}", ts2, msg_str); } sender.send(msg_str); } } // 处理 FileVer 响应 if msg_type == "FileVer" { if let Some(file_versions) = data.get("Data").and_then(|d| d.get("file_versions")).and_then(|v| v.as_object()) { for (filename, server_ver) in file_versions { let server_version = server_ver.as_str().unwrap_or("0.0.0"); let local_version = get_local_file_version(filename); let tmp_size = get_tmp_file_size(filename); if debug_msg { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [版本] {}: 服务端={}, 本地={}, tmp大小={}", ts, filename, server_version, local_version, tmp_size); } // 比较版本:如果本地不存在或比服务端旧,则下载 let need_update = local_version == "0.0.0" || version_less_than(&local_version, server_version); if need_update { if tmp_size > 0 { // 有临时文件,请求 hash 校验 if debug_msg { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [续传] {} 发现未完成下载,请求 hash 校验...", ts, filename); } request_file_md5(&sender, filename, tmp_size, debug_msg); } else { // 无临时文件,从头下载 if debug_msg { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [升级] {} 需要更新,开始下载...", ts, filename); } request_download(&sender, filename, 0, debug_msg); } } } } } // 处理 Md5 响应 if msg_type == "Md5" { if let Some(md5_data) = data.get("Data").and_then(|v| v.as_object()) { let filename = md5_data.get("filename").and_then(|v| v.as_str()).unwrap_or(""); let server_md5 = md5_data.get("md5").and_then(|v| v.as_str()).unwrap_or(""); let bytes = md5_data.get("bytes").and_then(|v| v.as_u64()).unwrap_or(0); let tmp_filename = format!("{}.tmp", filename); // 计算本地 tmp 文件前 bytes 字节的 md5 let local_md5 = compute_file_hash(&tmp_filename, bytes, debug_msg); if debug_msg { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); if let Some(ref lm) = local_md5 { println!("{} [续传] {} md5对比: 本地={}, 服务端={}", ts, filename, lm, server_md5); } else { println!("{} [续传] {} 无法计算本地md5,重新下载", ts, filename); } } // 比较 md5 if local_md5.as_deref() == Some(server_md5) { // md5 相同,从 offset bytes 处续传 if debug_msg { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [续传] {} md5匹配,从 {} 字节处续传", ts, filename, bytes); } request_download(&sender, filename, bytes, debug_msg); } else { // md5 不同,删除临时文件,重新下载 if debug_msg { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [续传] {} md5不匹配,重新下载", ts, filename); } let tmp_path = get_updater_data_dir().join(&tmp_filename); let _ = fs::remove_file(&tmp_path); request_download(&sender, filename, 0, debug_msg); } } } // 处理文件块 if msg_type == "FileChunk" { if let Some(data_obj) = data.get("Data").and_then(|v| v.as_object()) { handle_file_chunk(data_obj, debug_msg); } } // 处理下载完成 if msg_type == "DownloadComplete" { if let Some(data_obj) = data.get("Data").and_then(|v| v.as_object()) { handle_download_complete(data_obj, debug_msg); } } }); // 设置断开连接回调 let debug_disconnect = debug_mode; client.on_disconnected(move || { if debug_disconnect { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [断开] 连接已断开", ts); } }); // 设置错误回调 client.on_error(|error| { eprintln!("[错误] WebSocket: {}", error); }); // 设置首次连接回调(GetFileVer 改到收到 welcome 后发送) let debug_first = debug_mode; let device_number_first = resolve_device_number(); client.on_first_connect(move |_url, _sender| { let device_number = device_number_first.clone(); Box::pin(async move { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); if device_number.is_empty() || device_number == "UNKNOWN" { if debug_first { println!("{} [连接] 已连接,未配置设备号,仅维持心跳", ts); } } else { if debug_first { println!("{} [连接] 已连接,等待服务器欢迎消息...", ts); } } }) as Pin + Send + Sync>> }); // 设置重连回调 let debug_reconnect = debug_mode; client.on_reconnecting(move |attempt, url_arc| { Box::pin(async move { if debug_reconnect { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [重连] 第 {} 次重连中...", ts, attempt); } let new_url = resolve_ws_url(); *url_arc.lock().await = new_url.clone(); if debug_reconnect { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); println!("{} [重连] 配置已更新,服务器: {}", ts, new_url); } }) as Pin + Send + Sync>> }); // 设置重连成功回调 let debug_reconnected = debug_mode; let device_number_reconn = resolve_device_number(); client.on_reconnected(move |_url, _sender| { let device_number = device_number_reconn.clone(); Box::pin(async move { let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f"); if device_number.is_empty() || device_number == "UNKNOWN" { if debug_reconnected { println!("{} [重连] 已重连,未配置设备号", ts); } } else { if debug_reconnected { println!("{} [重连] 已重连,等待服务器欢迎消息...", ts); } } }) as Pin + Send + Sync>> }); // 连接(CubeLib 会自动处理重连) if debug_mode { println!("[启动] 开始连接..."); } client.connect().await; if debug_mode { println!("[启动] 连接已结束"); println!("Updater 已停止"); } } #[tokio::main] async fn main() { // 检查是否已有 Updater 进程在运行 if is_process_running("Updater.exe") { return; } // 加载 Updater 自身配置(debug_mode) let config = load_updater_config(); // 非 debug 模式下释放控制台,后台静默运行 if !config.debug_mode { #[cfg(windows)] { use windows::Win32::System::Console; use windows::Win32::Foundation::HWND; unsafe { let console = Console::GetConsoleWindow(); if console != HWND::default() { let _ = Console::FreeConsole(); } } } } // 运行 Updater run_updater(config.debug_mode).await; if config.debug_mode { println!("Updater 已退出"); } }