1.更新对称NAT的打洞方式；2.支持windows服务；3.更新协议内容

switch-desktop/src/windows/mod.rs

@@ -1,7 +1,20 @@
+use std::{io, thread};
+use std::ffi::OsString;
+use std::path::PathBuf;
+use std::time::Duration;
+
 use clap::Parser;
 use console::style;
+use windows_service::Error;
+use windows_service::service::{
+    ServiceAccess, ServiceErrorControl, ServiceInfo, ServiceStartType, ServiceState, ServiceType,
+};
+use windows_service::service_manager::{ServiceManager, ServiceManagerAccess};
+
+use crate::config;
 
 pub mod service;
+mod windows_admin_check;
 
 #[derive(Parser, Debug)]
 #[command(
@@ -17,13 +30,13 @@ struct Args {
     /// Only devices with the same token can communicate with each other.
     /// It is recommended to use uuid to ensure uniqueness
     #[arg(long)]
-    token: String,
+    token: Option<String>,
     /// 给设备一个名称，为空时默认用系统版本信息
     #[arg(long)]
     name: Option<String>,
-    /// 安装服务，安装后可以后台运行
+    /// 安装服务，安装后可以后台运行，需要指定安装路径
     #[arg(long)]
-    install: bool,
+    install: Option<String>,
     /// 卸载服务
     #[arg(long)]
     uninstall: bool,
@@ -32,72 +45,165 @@ struct Args {
     #[arg(long)]
     start: bool,
     #[arg(long)]
-    /// 停止，安装服务后，使用--stop停止服务
+    /// 停止，安装服务后，使用 --stop停止服务
     stop: bool,
-
+    /// 启动服务后，使用 --list 查看设备列表
+    #[arg(long)]
+    list: bool,
+    /// 启动服务后，使用 --status 查看设备状态
+    #[arg(long)]
+    status: bool,
 }
 
-const SERVICE_FLAG: &'static str = "start_switch_service_";
-const SERVICE_NAME: &'static str = "switch-service";
+pub const SERVICE_FLAG: &'static str = "start_switch_service_";
+pub const SERVICE_NAME: &'static str = "switch-service";
+pub const SERVICE_TYPE: ServiceType = ServiceType::OWN_PROCESS;
 
 pub fn main0() {
-    let args: Vec<_> = std::env::args().collect();
-    if args.len() == 2 && args[1] == SERVICE_FLAG {
-        //以服务的方式启动
-        service::start();
-    }
     let args = Args::parse();
-    if args.install {
-        if let Err(e) = install() {
-            log::error!("{:?}",e);
-        }else{
-            println!("{}",style("安装成功").green())
+    if args.list || args.status {
+        match service_state() {
+            Ok(state) => {
+                if state == ServiceState::Running {
+                    let command_client = crate::command::client::CommandClient::new().unwrap();
+                    let out = if args.list {
+                        command_client.list().unwrap()
+                    } else if args.status {
+                        command_client.status().unwrap()
+                    } else {
+                        "".to_string()
+                    };
+                    println!("{}", out);
+                } else {
+                    println!("服务未启动")
+                }
+            }
+            Err(e) => {
+                println!("{:?}", e);
+            }
         }
-        pause();
         return;
     }
-    if args.uninstall {
+    if !windows_admin_check::is_app_elevated() {
+        println!("{}", style("请使用管理员权限运行").red());
+        return;
+    }
+    if let Some(path) = args.install {
+        let path: PathBuf = path.into();
+        if !path.exists() {
+            std::fs::create_dir_all(&path).unwrap();
+        }
+        if !path.is_dir() {
+            println!("参数必须为文件目录");
+        } else {
+            if let Err(e) = install(path) {
+                log::error!("{:?}", e);
+            } else {
+                println!("{}", style("安装成功").green())
+            }
+        }
+    } else if args.uninstall {
         if let Err(e) = uninstall() {
-            log::error!("{:?}",e);
-        }else{
-            println!("{}",style("卸载成功").green())
+            log::error!("{:?}", e);
+        } else {
+            println!("{}", style("卸载成功").green())
         }
-        pause();
-        return;
-    }
-    if args.start {
-        if let Err(e) = start() {
-            log::error!("{:?}",e);
-            // 在当前进程启动
+    } else if args.start {
+        match service_state() {
+            Ok(state) => {
+                if state == ServiceState::Stopped {
+                    if args.token.is_none() {
+                        println!("{}", style("需要参数 --token").red());
+                    } else {
+                        let token = args.token.clone().unwrap();
+                        config::save_config(config::ArgsConfig::new(
+                            token.clone(),
+                            args.name.clone(),
+                        ))
+                            .unwrap();
+                        match start() {
+                            Ok(_) => {
+                                //需要检查启动状态
+                                println!("{}", style("启动成功").green())
+                            }
+                            Err(e) => {
+                                match e {
+                                    Error::Winapi(ref e) => {
+                                        if let Some(code) = e.raw_os_error() {
+                                            if code == 1060 {
+                                                //指定的服务未安装。
+                                                println!(
+                                                    "{}",
+                                                    style("服务未安装，在当前进程启动").red()
+                                                );
+                                                crate::start(token, args.name);
+                                                return;
+                                            }
+                                        }
+                                    }
+                                    _ => {}
+                                }
+                                log::error!("{:?}", e);
+                            }
+                        }
+                    }
+                } else {
+                    println!("服务未停止");
+                }
+            }
+            Err(e) => {
+                println!("{:?}", e);
+            }
         }
-        pause();
+    } else if args.stop {
+        match stop() {
+            Ok(_) => {
+                println!("{}", style("停止成功").green())
+            }
+            Err(e) => {
+                log::error!("{:?}", e);
+            }
+        }
+    } else {
+        println!("使用参数 -h 查看帮助")
     }
+    pause();
 }
 
 fn pause() {
-    println!("按任意键退出...");
-    std::io::stdin().read_u8().unwrap();
+    println!("{}", style("按任意键退出...").green());
+    use console::Term;
+    let term = Term::stdout();
+    let _ = term.read_char().unwrap();
 }
 
-
-fn install() -> Result<(), windows_service::Error> {
-    use std::ffi::OsString;
-    use windows_service::{
-        service::{ServiceAccess, ServiceErrorControl, ServiceInfo, ServiceStartType, ServiceType},
-        service_manager::{ServiceManager, ServiceManagerAccess},
-    };
-
+fn install(path: PathBuf) -> Result<(), Error> {
     let manager_access = ServiceManagerAccess::CONNECT | ServiceManagerAccess::CREATE_SERVICE;
     let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?;
-    let service_binary_path = std::env::current_exe().unwrap();
+    let current_exe_path = std::env::current_exe().unwrap();
+    let service_path = path.join("switch-service.exe");
+    std::fs::copy(current_exe_path, service_path.as_path()).unwrap();
+    if let Err(e) = std::fs::copy("wintun.dll", path.join("wintun.dll").as_path()) {
+        if e.kind() == io::ErrorKind::NotFound {
+            println!("Not fount 'wintun.dll'. Please put 'wintun.dll' in the current directory");
+            std::process::exit(0);
+        } else {
+            panic!("{:?}", e)
+        }
+    }
+    let mut launch_arguments = Vec::new();
+    launch_arguments.push(OsString::from(SERVICE_FLAG));
+    launch_arguments.push(OsString::from(
+        dirs::home_dir().unwrap().join(".switch").to_str().unwrap(),
+    ));
     let service_info = ServiceInfo {
         name: OsString::from(SERVICE_NAME),
         display_name: OsString::from("switch service"),
         service_type: ServiceType::OWN_PROCESS,
         start_type: ServiceStartType::OnDemand,
         error_control: ServiceErrorControl::Normal,
-        executable_path: service_binary_path.into(),
-        launch_arguments: vec![OsString::from(SERVICE_FLAG); 1],
+        executable_path: service_path.into(),
+        launch_arguments,
         dependencies: vec![],
         account_name: None, // run as System
         account_password: None,
@@ -107,13 +213,7 @@ fn install() -> Result<(), windows_service::Error> {
     Ok(())
 }
 
-fn uninstall() -> Result<(), windows_service::Error> {
-    use std::{thread, time::Duration};
-    use windows_service::{
-        service::{ServiceAccess, ServiceState},
-        service_manager::{ServiceManager, ServiceManagerAccess},
-    };
-
+fn uninstall() -> Result<(), Error> {
     let manager_access = ServiceManagerAccess::CONNECT;
     let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?;
 
@@ -131,14 +231,27 @@ fn uninstall() -> Result<(), windows_service::Error> {
     Ok(())
 }
 
-fn start() -> Result<(), windows_service::Error> {
-    use std::env;
-    use windows_service::{
-        service::ServiceAccess,
-        service_manager::{ServiceManager, ServiceManagerAccess},
-    };
+fn start() -> Result<(), Error> {
     let manager_access = ServiceManagerAccess::CONNECT;
     let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?;
     let service = service_manager.open_service(SERVICE_NAME, ServiceAccess::START)?;
-    service.start(&[])
-}
+    service.start(&[""])
+}
+
+fn service_state() -> Result<ServiceState, Error> {
+    let manager_access = ServiceManagerAccess::CONNECT;
+    let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?;
+
+    let service_access = ServiceAccess::QUERY_STATUS;
+    let service = service_manager.open_service(SERVICE_NAME, service_access)?;
+    let service_status = service.query_status()?;
+    return Ok(service_status.current_state);
+}
+
+fn stop() -> Result<(), Error> {
+    let manager_access = ServiceManagerAccess::CONNECT;
+    let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?;
+    let service = service_manager.open_service(SERVICE_NAME, ServiceAccess::STOP)?;
+    service.stop()?;
+    Ok(())
+}

switch-desktop/src/windows/service.rs

@@ -2,12 +2,110 @@
 // extern crate windows_service;
 
 use std::ffi::OsString;
-use windows_service::{define_windows_service, service_dispatcher};
+use std::sync::Arc;
+use std::thread;
+use std::time::Duration;
+
+
+use windows_service::service::{
+    ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus,
+};
+use windows_service::service_control_handler::ServiceControlHandlerResult;
+use windows_service::{define_windows_service, service_control_handler, service_dispatcher};
+
+use switch::{Config, Switch};
+
+use crate::windows::config::read_config;
 
 define_windows_service!(ffi_service_main, switch_service_main);
-pub fn switch_service_main(arguments: Vec<OsString>) {
+pub fn switch_service_main(_arguments: Vec<OsString>) {
+    thread::spawn(|| match service_main() {
+        Ok(_) => {}
+        Err(e) => {
+            log::warn!("{:?}", e);
+        }
+    });
+}
 
+fn service_main() -> windows_service::Result<()> {
+    let parker = crossbeam::sync::Parker::new();
+    let un_parker = parker.unparker().clone();
+    let event_handler = move |control_event| -> ServiceControlHandlerResult {
+        match control_event {
+            // Notifies a service to report its current status information to the service
+            // control manager. Always return NoError even if not implemented.
+            ServiceControl::Interrogate => ServiceControlHandlerResult::NoError,
+
+            // Handle stop
+            ServiceControl::Stop => {
+                log::info!("handler 服务停止");
+                un_parker.unpark();
+                ServiceControlHandlerResult::NoError
+            }
+            _ => ServiceControlHandlerResult::NotImplemented,
+        }
+    };
+
+    // Register system service event handler.
+    // The returned status handle should be used to report service status changes to the system.
+    let status_handle =
+        service_control_handler::register(crate::windows::SERVICE_NAME, event_handler)?;
+
+    // Tell the system that service is running
+    status_handle.set_service_status(ServiceStatus {
+        service_type: crate::windows::SERVICE_TYPE,
+        current_state: ServiceState::Running,
+        controls_accepted: ServiceControlAccept::STOP,
+        exit_code: ServiceExitCode::Win32(0),
+        checkpoint: 0,
+        wait_hint: Duration::default(),
+        process_id: None,
+    })?;
+    if let Some(config) = read_config() {
+        let mac_address = mac_address::get_mac_address().unwrap().unwrap().to_string();
+        let un_parker = parker.unparker().clone();
+        match Config::new(config.token, mac_address, config.name, move || {
+            un_parker.unpark();
+        }) {
+            Ok(config) => match Switch::start(config) {
+                Ok(switch) => {
+                    log::info!("switch-service服务启动");
+                    let switch = Arc::new(switch);
+                    let command_server = crate::command::server::CommandServer::new();
+                    let switch1 = switch.clone();
+                    thread::spawn(move || {
+                        if let Err(e) = command_server.start(switch1) {
+                            log::warn!("{:?}", e);
+                        }
+                    });
+                    parker.park();
+                    switch.stop_async();
+                    thread::sleep(Duration::from_secs(1));
+                    log::info!("switch-service服务停止");
+                }
+                Err(e) => {
+                    log::error!("{:?}", e);
+                }
+            },
+            Err(e) => {
+                log::error!("{:?}", e);
+            }
+        };
+    } else {
+        log::info!("配置文件为空");
+    }
+    status_handle.set_service_status(ServiceStatus {
+        service_type: crate::windows::SERVICE_TYPE,
+        current_state: ServiceState::Stopped,
+        controls_accepted: ServiceControlAccept::empty(),
+        exit_code: ServiceExitCode::Win32(0),
+        checkpoint: 0,
+        wait_hint: Duration::default(),
+        process_id: None,
+    })
+}
+
+pub fn start() {
+    log::info!("以服务的方式启动");
+    service_dispatcher::start("switch-service", ffi_service_main).unwrap();
 }
-pub fn start(){
-   service_dispatcher::start("switch-service",ffi_service_main).unwrap();
-}

switch-desktop/src/windows/windows_admin_check.rs

@@ -0,0 +1,76 @@
+/// 使用 https://github.com/spa5k/is_sudo/blob/main/src/window.rs
+use std::io::Error;
+use std::ptr;
+
+use winapi::um::handleapi::CloseHandle;
+use winapi::um::processthreadsapi::{GetCurrentProcess, OpenProcessToken};
+use winapi::um::securitybaseapi::GetTokenInformation;
+use winapi::um::winnt::{TokenElevation, HANDLE, TOKEN_ELEVATION, TOKEN_QUERY};
+
+// Use std::io::Error::last_os_error for errors.
+// NOTE: For this example I'm simple passing on the OS error.
+// However, customising the error could provide more context
+
+/// Returns true if the current process has admin rights, otherwise false.
+pub fn is_app_elevated() -> bool {
+    _is_app_elevated().unwrap_or(false)
+}
+
+/// On success returns a bool indicating if the current process has admin rights.
+/// Otherwise returns an OS error.
+///
+/// This is unlikely to fail but if it does it's even more unlikely that you have admin permissions anyway.
+/// Therefore the public function above simply eats the error and returns a bool.
+fn _is_app_elevated() -> Result<bool, Error> {
+    let token = QueryAccessToken::from_current_process()?;
+    token.is_elevated()
+}
+
+/// A safe wrapper around querying Windows access tokens.
+pub struct QueryAccessToken(HANDLE);
+
+impl QueryAccessToken {
+    pub fn from_current_process() -> Result<Self, Error> {
+        unsafe {
+            let mut handle: HANDLE = ptr::null_mut();
+            let result = OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut handle);
+
+            if result != 0 {
+                Ok(Self(handle))
+            } else {
+                Err(Error::last_os_error())
+            }
+        }
+    }
+
+    /// On success returns a bool indicating if the access token has elevated privilidges.
+    /// Otherwise returns an OS error.
+    pub fn is_elevated(&self) -> Result<bool, Error> {
+        unsafe {
+            let mut elevation = TOKEN_ELEVATION::default();
+            let size = std::mem::size_of::<TOKEN_ELEVATION>() as u32;
+            let mut ret_size = size;
+            // The weird looking repetition of `as *mut _` is casting the reference to a c_void pointer.
+            if GetTokenInformation(
+                self.0,
+                TokenElevation,
+                &mut elevation as *mut _ as *mut _,
+                size,
+                &mut ret_size,
+            ) != 0
+            {
+                Ok(elevation.TokenIsElevated != 0)
+            } else {
+                Err(Error::last_os_error())
+            }
+        }
+    }
+}
+
+impl Drop for QueryAccessToken {
+    fn drop(&mut self) {
+        if !self.0.is_null() {
+            unsafe { CloseHandle(self.0) };
+        }
+    }
+}