use crate::utils::RustFSConfig; use dioxus::logger::tracing::{debug, error, info}; use lazy_static::lazy_static; use rust_embed::RustEmbed; use sha2::{Digest, Sha256}; use std::error::Error; use std::path::{Path, PathBuf}; use std::process::Command as StdCommand; use std::time::Duration; use tokio::fs; use tokio::fs::File; use tokio::io::AsyncWriteExt; use tokio::net::TcpStream; use tokio::sync::{mpsc, Mutex}; #[derive(RustEmbed)] #[folder = "$CARGO_MANIFEST_DIR/embedded-rustfs/"] struct Asset; // Use `lazy_static` to cache the checksum of embedded resources lazy_static! { static ref RUSTFS_HASH: Mutex = { let rustfs_file = if cfg!(windows) { "rustfs.exe" } else { "rustfs" }; let rustfs_data = Asset::get(rustfs_file).expect("RustFs binary not embedded"); let hash = hex::encode(Sha256::digest(&rustfs_data.data)); Mutex::new(hash) }; } /// Service command /// This enum represents the commands that can be sent to the service manager /// to start, stop, or restart the service /// The `Start` variant contains the configuration for the service /// The `Restart` variant contains the configuration for the service /// /// # Example /// ``` /// let config = RustFSConfig { /// address: "127.0.0.1:9000".to_string(), /// host: "127.0.0.1".to_string(), /// port: "9000".to_string(), /// access_key: "rustfsadmin".to_string(), /// secret_key: "rustfsadmin".to_string(), /// domain_name: "demo.rustfs.com".to_string(), /// volume_name: "data".to_string(), /// console_address: "127.0.0.1:9001".to_string(), /// }; /// /// let command = ServiceCommand::Start(config); /// println!("{:?}", command); /// /// assert_eq!(command, ServiceCommand::Start(config)); /// ``` pub enum ServiceCommand { Start(RustFSConfig), Stop, Restart(RustFSConfig), } /// Service operation result /// This struct represents the result of a service operation /// It contains information about the success of the operation, /// /// # Example /// ``` /// use chrono::Local; /// /// let result = ServiceOperationResult { /// success: true, /// start_time: chrono::Local::now(), /// end_time: chrono::Local::now(), /// message: "服务启动成功".to_string(), /// }; /// /// println!("{:?}", result); /// assert_eq!(result.success, true); /// ``` #[derive(Debug)] pub struct ServiceOperationResult { pub success: bool, pub start_time: chrono::DateTime, pub end_time: chrono::DateTime, pub message: String, } /// Service manager /// This struct represents a service manager that can be used to start, stop, or restart a service /// It contains a command sender that can be used to send commands to the service manager /// /// # Example /// ``` /// let service_manager = ServiceManager::new(); /// println!("{:?}", service_manager); /// ``` #[derive(Debug, Clone)] pub struct ServiceManager { command_tx: mpsc::Sender, // process: Arc>>, // pid: Arc>>, // Add PID storage // current_config: Arc>>, // Add configuration storage } impl ServiceManager { /// check if the service is running and return a pid /// This function is platform dependent /// On Unix systems, it uses the `ps` command to check for the service /// On Windows systems, it uses the `wmic` command to check for the service /// /// # Example /// ``` /// let pid = check_service_status().await; /// println!("{:?}", pid); /// ``` pub async fn check_service_status() -> Option { #[cfg(unix)] { // use the ps command on a unix system if let Ok(output) = StdCommand::new("ps").arg("-ef").output() { let output_str = String::from_utf8_lossy(&output.stdout); for line in output_str.lines() { // match contains `rustfs/bin/rustfs` of the line if line.contains("rustfs/bin/rustfs") && !line.contains("grep") { if let Some(pid_str) = line.split_whitespace().nth(1) { if let Ok(pid) = pid_str.parse::() { return Some(pid); } } } } } } #[cfg(windows)] { if let Ok(output) = StdCommand::new("wmic") .arg("process") .arg("where") .arg("caption='rustfs.exe'") .arg("get") .arg("processid") .output() { let output_str = String::from_utf8_lossy(&output.stdout); for line in output_str.lines() { if let Ok(pid) = line.trim().parse::() { return Some(pid); } } } } None } /// Prepare the service /// This function downloads the service executable if it doesn't exist /// It also creates the necessary directories for the service /// /// # Example /// ``` /// let executable_path = prepare_service().await; /// println!("{:?}", executable_path); /// ``` async fn prepare_service() -> Result> { // get the user directory let home_dir = dirs::home_dir().ok_or("无法获取用户目录")?; let rustfs_dir = home_dir.join("rustfs"); let bin_dir = rustfs_dir.join("bin"); let data_dir = rustfs_dir.join("data"); let logs_dir = rustfs_dir.join("logs"); // create the necessary directories for dir in [&bin_dir, &data_dir, &logs_dir] { if !dir.exists() { tokio::fs::create_dir_all(dir).await?; } } let rustfs_file = if cfg!(windows) { "rustfs.exe" } else { "rustfs" }; let executable_path = bin_dir.join(rustfs_file); let hash_path = bin_dir.join("embedded_rustfs.sha256"); if executable_path.exists() && hash_path.exists() { let cached_hash = fs::read_to_string(&hash_path).await?; let expected_hash = RUSTFS_HASH.lock().await; if cached_hash == *expected_hash { println!("Use cached rustfs: {:?}", executable_path); return Ok(executable_path); } } // Extract and write files let rustfs_data = Asset::get(rustfs_file).expect("RustFS binary not embedded"); let mut file = File::create(&executable_path).await?; file.write_all(&rustfs_data.data).await?; let expected_hash = hex::encode(Sha256::digest(&rustfs_data.data)); fs::write(&hash_path, expected_hash).await?; // set execution permissions on unix systems #[cfg(unix)] { use std::os::unix::fs::PermissionsExt; let mut perms = std::fs::metadata(&executable_path)?.permissions(); perms.set_mode(0o755); std::fs::set_permissions(&executable_path, perms)?; } Ok(executable_path) } /// Helper function: Extracts the port from the address string /// /// # Example /// ``` /// let address = "127.0.0.1:9000"; /// let port = extract_port(address); /// println!("{:?}", port); /// ``` fn extract_port(address: &str) -> Option { address.split(':').nth(1)?.parse().ok() } /// Create a new instance of the service manager /// /// # Example /// ``` /// let service_manager = ServiceManager::new(); /// println!("{:?}", service_manager); /// ``` pub(crate) fn new() -> Self { let (command_tx, mut command_rx) = mpsc::channel(10); // Start the control loop tokio::spawn(async move { while let Some(cmd) = command_rx.recv().await { match cmd { ServiceCommand::Start(config) => { if let Err(e) = Self::start_service(&config).await { Self::show_error(&format!("启动服务失败:{}", e)); } } ServiceCommand::Stop => { if let Err(e) = Self::stop_service().await { Self::show_error(&format!("停止服务失败:{}", e)); } } ServiceCommand::Restart(config) => { if Self::check_service_status().await.is_some() { if let Err(e) = Self::stop_service().await { Self::show_error(&format!("重启服务失败:{}", e)); continue; } } if let Err(e) = Self::start_service(&config).await { Self::show_error(&format!("重启服务失败:{}", e)); } } } } }); ServiceManager { command_tx } } /// Start the service /// This function starts the service with the given configuration /// /// # Example /// ``` /// let config = RustFSConfig { /// address: "127.0.0.1:9000".to_string(), /// host: "127.0.0.1".to_string(), /// port: "9000".to_string(), /// access_key: "rustfsadmin".to_string(), /// secret_key: "rustfsadmin".to_string(), /// domain_name: "demo.rustfs.com".to_string(), /// volume_name: "data".to_string(), /// console_address: "127.0.0.1:9001".to_string(), /// }; /// /// let result = start_service(&config).await; /// println!("{:?}", result); /// ``` async fn start_service(config: &RustFSConfig) -> Result<(), Box> { // Check if the service is already running if let Some(existing_pid) = Self::check_service_status().await { return Err(format!("服务已经在运行,PID: {}", existing_pid).into()); } // Prepare the service program let executable_path = Self::prepare_service().await?; // Check the data catalog let volume_name_path = Path::new(&config.volume_name); if !volume_name_path.exists() { tokio::fs::create_dir_all(&config.volume_name).await?; } // Extract the port from the configuration let main_port = Self::extract_port(&config.address).ok_or("无法解析主服务端口")?; let console_port = Self::extract_port(&config.console_address).ok_or("无法解析控制台端口")?; let host = config.address.split(':').next().ok_or("无法解析主机地址")?; // Check the port let ports = vec![main_port, console_port]; for port in ports { if Self::is_port_in_use(host, port).await { return Err(format!("端口 {} 已被占用", port).into()); } } // Start the service let mut child = tokio::process::Command::new(executable_path) .arg("--address") .arg(&config.address) .arg("--access-key") .arg(&config.access_key) .arg("--secret-key") .arg(&config.secret_key) .arg("--console-address") .arg(&config.console_address) .arg(config.volume_name.clone()) .spawn()?; let process_pid = child.id().unwrap(); // Wait for the service to start tokio::time::sleep(Duration::from_secs(2)).await; // Check if the service started successfully if Self::is_port_in_use(host, main_port).await { Self::show_info(&format!("服务启动成功!进程 ID: {}", process_pid)); Ok(()) } else { child.kill().await?; Err("服务启动失败".into()) } } /// Stop the service /// This function stops the service /// /// # Example /// ``` /// let result = stop_service().await; /// println!("{:?}", result); /// ``` async fn stop_service() -> Result<(), Box> { let existing_pid = Self::check_service_status().await; debug!("existing_pid: {:?}", existing_pid); if let Some(service_pid) = existing_pid { // An attempt was made to terminate the process #[cfg(unix)] { StdCommand::new("kill").arg("-9").arg(service_pid.to_string()).output()?; } #[cfg(windows)] { StdCommand::new("taskkill") .arg("/F") .arg("/PID") .arg(&service_pid.to_string()) .output()?; } // Verify that the service is indeed stopped tokio::time::sleep(Duration::from_secs(1)).await; if Self::check_service_status().await.is_some() { return Err("服务停止失败".into()); } Self::show_info("服务已成功停止"); Ok(()) } else { Err("服务未运行".into()) } } /// Check if the port is in use /// This function checks if the given port is in use on the given host /// /// # Example /// ``` /// let host = "127.0.0.1"; /// let port = 9000; /// let result = is_port_in_use(host, port).await; /// println!("{:?}", result); /// ``` async fn is_port_in_use(host: &str, port: u16) -> bool { TcpStream::connect(format!("{}:{}", host, port)).await.is_ok() } /// Show an error message /// This function shows an error message dialog /// /// # Example /// ``` /// show_error("This is an error message"); /// ``` pub(crate) fn show_error(message: &str) { rfd::MessageDialog::new() .set_title("错误") .set_description(message) .set_level(rfd::MessageLevel::Error) .show(); } /// Show an information message /// This function shows an information message dialog /// /// # Example /// ``` /// show_info("This is an information message"); /// ``` pub(crate) fn show_info(message: &str) { rfd::MessageDialog::new() .set_title("成功") .set_description(message) .set_level(rfd::MessageLevel::Info) .show(); } /// Start the service /// This function sends a `Start` command to the service manager /// /// # Example /// ``` /// let config = RustFSConfig { /// address: "127.0.0.1:9000".to_string(), /// host: "127.0.0.1".to_string(), /// port: "9000".to_string(), /// access_key: "rustfsadmin".to_string(), /// secret_key: "rustfsadmin".to_string(), /// domain_name: "demo.rustfs.com".to_string(), /// volume_name: "data".to_string(), /// console_address: "127.0.0.1:9001".to_string(), /// }; /// /// let service_manager = ServiceManager::new(); /// let result = service_manager.start(config).await; /// println!("{:?}", result); /// ``` /// /// # Errors /// This function returns an error if the service fails to start /// /// # Panics /// This function panics if the port number is invalid /// /// # Safety /// This function is not marked as unsafe /// /// # Performance /// This function is not optimized for performance /// /// # Design /// This function is designed to be simple and easy to use /// /// # Security /// This function does not have any security implications pub async fn start(&self, config: RustFSConfig) -> Result> { let start_time = chrono::Local::now(); self.command_tx.send(ServiceCommand::Start(config.clone())).await?; let host = &config.host; let port = config.port.parse::().expect("无效的端口号"); // wait for the service to actually start let mut retries = 0; while retries < 30 { // wait up to 30 seconds if Self::check_service_status().await.is_some() && Self::is_port_in_use(host, port).await { let end_time = chrono::Local::now(); return Ok(ServiceOperationResult { success: true, start_time, end_time, message: "服务启动成功".to_string(), }); } tokio::time::sleep(Duration::from_secs(1)).await; retries += 1; } Err("服务启动超时".into()) } /// Stop the service /// This function sends a `Stop` command to the service manager /// /// # Example /// ``` /// let service_manager = ServiceManager::new(); /// let result = service_manager.stop().await; /// println!("{:?}", result); /// ``` /// /// # Errors /// This function returns an error if the service fails to stop /// /// # Panics /// This function panics if the port number is invalid /// /// # Safety /// This function is not marked as unsafe /// /// # Performance /// This function is not optimized for performance /// /// # Design /// This function is designed to be simple and easy to use /// /// # Security /// This function does not have any security implications pub async fn stop(&self) -> Result> { let start_time = chrono::Local::now(); self.command_tx.send(ServiceCommand::Stop).await?; // Wait for the service to actually stop let mut retries = 0; while retries < 15 { // Wait up to 15 seconds if Self::check_service_status().await.is_none() { let end_time = chrono::Local::now(); return Ok(ServiceOperationResult { success: true, start_time, end_time, message: "服务停止成功".to_string(), }); } tokio::time::sleep(Duration::from_secs(1)).await; retries += 1; } Err("服务停止超时".into()) } /// Restart the service /// This function sends a `Restart` command to the service manager /// /// # Example /// ``` /// let config = RustFSConfig { /// address: "127.0.0.1:9000".to_string(), /// host: "127.0.0.1".to_string(), /// port: "9000".to_string(), /// access_key: "rustfsadmin".to_string(), /// secret_key: "rustfsadmin".to_string(), /// domain_name: "demo.rustfs.com".to_string(), /// volume_name: "data".to_string(), /// console_address: "127.0.0.1:9001".to_string(), /// }; /// /// let service_manager = ServiceManager::new(); /// let result = service_manager.restart(config).await; /// println!("{:?}", result); /// ``` /// /// # Errors /// This function returns an error if the service fails to restart /// /// # Panics /// This function panics if the port number is invalid /// /// # Safety /// This function is not marked as unsafe /// /// # Performance /// This function is not optimized for performance /// /// # Design /// This function is designed to be simple and easy to use /// /// # Security /// This function does not have any security implications pub async fn restart(&self, config: RustFSConfig) -> Result> { let start_time = chrono::Local::now(); self.command_tx.send(ServiceCommand::Restart(config.clone())).await?; let host = &config.host; let port = config.port.parse::().expect("无效的端口号"); // wait for the service to restart let mut retries = 0; while retries < 45 { // Longer waiting time is given as both the stop and start processes are involved if Self::check_service_status().await.is_some() && Self::is_port_in_use(host, port).await { match config.save() { Ok(_) => info!("save config success"), Err(e) => { error!("save config error: {}", e); self.command_tx.send(ServiceCommand::Stop).await?; Self::show_error("保存配置失败"); return Err("保存配置失败".into()); } } let end_time = chrono::Local::now(); return Ok(ServiceOperationResult { success: true, start_time, end_time, message: "服务重启成功".to_string(), }); } tokio::time::sleep(Duration::from_secs(1)).await; retries += 1; } Err("服务重启超时".into()) } }