连接服务端

Windows/CS/Framework4.0/Updater/src/main.rs

@@ -1,71 +1,174 @@
 use serde::{Deserialize, Serialize};
 use std::fs;
 use std::path::PathBuf;
+use std::process::Command;
 use tokio::signal;
 use tokio::time;
+use futures_util::sink::SinkExt;
+use futures_util::stream::StreamExt;
 
+/// 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,
-        }
+        Self { debug_mode: false }
     }
 }
 
-fn get_config_path() -> PathBuf {
+/// 获取 Updater 自身配置路径 AppData/Updater/config.json
+fn get_updater_config_path() -> PathBuf {
     let exe_path = std::env::current_exe().expect("Failed to get executable path");
-    let drive = if let Some(parent) = exe_path.parent() {
-        let drive = parent.as_os_str().to_str().unwrap_or("C:/").split('\\').next().unwrap_or("C:");
-        std::path::Path::new(&format!("{}/", drive)).to_path_buf()
-    } else {
-        std::path::Path::new("C:/").to_path_buf()
-    };
-    let appdata = drive.join("AppData").join("Updater");
-    fs::create_dir_all(&appdata).expect("Failed to create config directory");
-    appdata.join("config.json")
+    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 updater_dir = appdata.join("Updater");
+    let _ = fs::create_dir_all(&updater_dir);
+    updater_dir.join("config.json")
 }
 
-fn load_config() -> Config {
-    let config_path = get_config_path();
+/// 获取公共配置路径 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() {
-        match fs::read_to_string(&config_path) {
-            Ok(content) => {
-                match serde_json::from_str(&content) {
-                    Ok(config) => config,
-                    Err(_) => Config::default(),
-                }
+        if let Ok(content) = fs::read_to_string(&config_path) {
+            if let Ok(config) = serde_json::from_str::<Config>(&content) {
+                return config;
             }
-            Err(_) => Config::default(),
-        }
-    } else {
-        let default_config = Config::default();
-        match serde_json::to_string_pretty(&default_config) {
-            Ok(content) => {
-                match fs::write(&config_path, content) {
-                    Ok(_) => default_config,
-                    Err(_) => default_config,
-                }
-            }
-            Err(_) => default_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
 }
 
-async fn upgrade(debug_mode: bool) {
+/// 从公共 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::<serde_json::Value>(&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()
+}
+
+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(&current_pid) {
+            count += 1;
+        }
+    }
+
+    count > 0
+}
+
+async fn upgrade(server_url: &str, debug_mode: bool) {
     if debug_mode {
-        println!("开始升级");
+        println!("开始升级检查，连接服务端...");
+    }
+    connect_to_websocket(server_url, debug_mode).await;
+}
+
+async fn connect_to_websocket(server_url: &str, debug_mode: bool) {
+    use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+    use tokio_tungstenite::{connect_async, tungstenite::Message};
+
+    if debug_mode {
+        println!("Connecting to WebSocket server: {}", server_url);
+    }
+
+    if let Ok(request) = server_url.into_client_request() {
+        match connect_async(request).await {
+            Ok((mut ws_stream, _)) => {
+                if debug_mode {
+                    println!("Connected to WebSocket server");
+                }
+
+                // 注册身份
+                let app_info = serde_json::json!({
+                    "type": "Updater",
+                    "action": "register"
+                });
+
+                if let Err(e) = ws_stream.send(Message::Text(app_info.to_string())).await {
+                    eprintln!("Failed to send register info: {:?}", e);
+                } else {
+                    match ws_stream.next().await {
+                        Some(Ok(message)) => {
+                            if debug_mode {
+                                println!("Server response: {:?}", message);
+                            }
+                        }
+                        Some(Err(e)) => {
+                            eprintln!("WebSocket error: {:?}", e);
+                        }
+                        None => {
+                            eprintln!("WebSocket connection closed by server");
+                        }
+                    }
+                }
+            }
+            Err(e) => {
+                eprintln!("Failed to connect to WebSocket server: {:?}", e);
+            }
+        }
+    } else {
+        eprintln!("Invalid WebSocket URL: {}", server_url);
     }
 }
 
 #[tokio::main]
 async fn main() {
-    let config = load_config();
-    
+    // 检查是否已有 Updater 进程在运行
+    if is_process_running("Updater.exe") {
+        return;
+    }
+
+    // 加载 Updater 自身配置（debug_mode）
+    let config = load_updater_config();
+
+    // 从公共 config.json 解析 WebSocket 连接地址
+    let server_url = resolve_ws_url();
+
+    // 非 debug 模式下释放控制台，后台静默运行
     if !config.debug_mode {
         #[cfg(windows)]
         {
@@ -79,13 +182,16 @@ async fn main() {
             }
         }
     }
-    
+
     if config.debug_mode {
+        println!("Updater started in debug mode");
+        println!("Server URL: {}", server_url);
+
         let mut interval = time::interval(time::Duration::from_secs(300));
         loop {
             tokio::select! {
                 _ = interval.tick() => {
-                    upgrade(config.debug_mode).await;
+                    upgrade(&server_url, config.debug_mode).await;
                 }
                 _ = signal::ctrl_c() => {
                     println!("Received Ctrl+C, exiting...");
@@ -97,7 +203,7 @@ async fn main() {
         let mut interval = time::interval(time::Duration::from_secs(300));
         loop {
             interval.tick().await;
-            upgrade(config.debug_mode).await;
+            upgrade(&server_url, config.debug_mode).await;
         }
     }
 }