windows service

Cargo.lock

@@ -61,6 +61,45 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
+[[package]]
+name = "clap"
+version = "3.2.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "86447ad904c7fb335a790c9d7fe3d0d971dc523b8ccd1561a520de9a85302750"
+dependencies = [
+ "atty",
+ "bitflags",
+ "clap_derive",
+ "clap_lex",
+ "indexmap",
+ "once_cell",
+ "strsim",
+ "termcolor",
+ "textwrap",
+]
+
+[[package]]
+name = "clap_derive"
+version = "3.2.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ea0c8bce528c4be4da13ea6fead8965e95b6073585a2f05204bd8f4119f82a65"
+dependencies = [
+ "heck",
+ "proc-macro-error",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "clap_lex"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2850f2f5a82cbf437dd5af4d49848fbdfc27c157c3d010345776f952765261c5"
+dependencies = [
+ "os_str_bytes",
+]
+
 [[package]]
 name = "crossbeam-channel"
 version = "0.5.6"
@@ -91,6 +130,20 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "err-derive"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c34a887c8df3ed90498c1c437ce21f211c8e27672921a8ffa293cb8d6d4caa9e"
+dependencies = [
+ "proc-macro-error",
+ "proc-macro2",
+ "quote",
+ "rustversion",
+ "syn",
+ "synstructure",
+]
+
 [[package]]
 name = "fast-socks5"
 version = "0.8.1"
@@ -144,6 +197,18 @@ version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9b919933a397b79c37e33b77bb2aa3dc8eb6e165ad809e58ff75bc7db2e34574"
 
+[[package]]
+name = "hashbrown"
+version = "0.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
+
+[[package]]
+name = "heck"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2540771e65fc8cb83cd6e8a237f70c319bd5c29f78ed1084ba5d50eeac86f7f9"
+
 [[package]]
 name = "hermit-abi"
 version = "0.1.19"
@@ -153,6 +218,16 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "indexmap"
+version = "1.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10a35a97730320ffe8e2d410b5d3b69279b98d2c14bdb8b70ea89ecf7888d41e"
+dependencies = [
+ "autocfg",
+ "hashbrown",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.3"
@@ -233,6 +308,12 @@ version = "1.15.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e82dad04139b71a90c080c8463fe0dc7902db5192d939bd0950f074d014339e1"
 
+[[package]]
+name = "os_str_bytes"
+version = "6.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ff7415e9ae3fff1225851df9e0d9e4e5479f947619774677a63572e55e80eff"
+
 [[package]]
 name = "parking_lot"
 version = "0.12.1"
@@ -262,6 +343,30 @@ version = "0.2.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e0a7ae3ac2f1173085d398531c705756c94a4c56843785df85a60c1a0afac116"
 
+[[package]]
+name = "proc-macro-error"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c"
+dependencies = [
+ "proc-macro-error-attr",
+ "proc-macro2",
+ "quote",
+ "syn",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro-error-attr"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "version_check",
+]
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.43"
@@ -367,6 +472,7 @@ dependencies = [
 name = "socks5ws"
 version = "0.1.0"
 dependencies = [
+ "clap",
  "fast-socks5",
  "flexi_logger",
  "log",
@@ -376,8 +482,15 @@ dependencies = [
  "tokio-stream",
  "tokio-util",
  "toml",
+ "windows-service",
 ]
 
+[[package]]
+name = "strsim"
+version = "0.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623"
+
 [[package]]
 name = "syn"
 version = "1.0.100"
@@ -389,6 +502,33 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "synstructure"
+version = "0.12.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f36bdaa60a83aca3921b5259d5400cbf5e90fc51931376a9bd4a0eb79aa7210f"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+ "unicode-xid",
+]
+
+[[package]]
+name = "termcolor"
+version = "1.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bab24d30b911b2376f3a13cc2cd443142f0c81dda04c118693e35b3835757755"
+dependencies = [
+ "winapi-util",
+]
+
+[[package]]
+name = "textwrap"
+version = "0.15.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "949517c0cf1bf4ee812e2e07e08ab448e3ae0d23472aee8a06c985f0c8815b16"
+
 [[package]]
 name = "thiserror"
 version = "1.0.35"
@@ -498,12 +638,30 @@ version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dcc811dc4066ac62f84f11307873c4850cb653bfa9b1719cee2bd2204a4bc5dd"
 
+[[package]]
+name = "unicode-xid"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f962df74c8c05a667b5ee8bcf162993134c104e96440b663c8daa176dc772d8c"
+
+[[package]]
+name = "version_check"
+version = "0.9.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
+
 [[package]]
 name = "wasi"
 version = "0.11.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
 
+[[package]]
+name = "widestring"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "653f141f39ec16bba3c5abe400a0c60da7468261cc2cbf36805022876bc721a8"
+
 [[package]]
 name = "winapi"
 version = "0.3.9"
@@ -520,12 +678,33 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
 
+[[package]]
+name = "winapi-util"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178"
+dependencies = [
+ "winapi",
+]
+
 [[package]]
 name = "winapi-x86_64-pc-windows-gnu"
 version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
 
+[[package]]
+name = "windows-service"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "917fdb865e7ff03af9dd86609f8767bc88fefba89e8efd569de8e208af8724b3"
+dependencies = [
+ "bitflags",
+ "err-derive",
+ "widestring",
+ "windows-sys",
+]
+
 [[package]]
 name = "windows-sys"
 version = "0.36.1"

Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+clap = { version = "3.2.22", features = ["derive"] }
 fast-socks5 = "0.8.1"
 flexi_logger = { version = "0.23.3", features = ["specfile_without_notification", "async"] }
 log = "0.4.17"
@@ -15,3 +16,4 @@ tokio = { version = "1.21.1", features = ["full", "winapi", "mio"] }
 tokio-stream = "0.1.10"
 tokio-util = "0.7.4"
 toml = "0.5.9"
+windows-service = "0.5.0"

src/main.rs

@@ -1,18 +1,55 @@
 extern crate flexi_logger;
 
-use tokio_util::sync::CancellationToken;
-
 mod config;
 mod server;
-use crate::config::Config;
+mod service;
-use crate::server::server_executor;
 
 use flexi_logger::{AdaptiveFormat, Age, Cleanup, Criterion, Duplicate, FileSpec, Logger, Naming};
 
+use clap::{Parser, Subcommand};
+
+#[derive(Subcommand, Debug)]
+enum Command {
+    /// install service
+    Install,
+    /// uninstall service
+    Uninstall,
+    /// start service
+    Start,
+    /// stop service
+    Stop,
+    /// run service (by Windows)
+    Run,
+}
+
+/// SOCKS5 proxy windows service
+#[derive(Parser, Debug)]
+#[clap(author, version, about, long_about = None)]
+#[clap(propagate_version = true)]
+struct Cli {
+    #[clap(subcommand)]
+    command: Command,
+}
+
+macro_rules! handle_error {
+    ($name:expr, $res:expr) => {
+        match $res {
+            Err(e) => {
+                log::error!("{} error: {:#?}", $name, e)
+            }
+            _ => (),
+        }
+    };
+}
+
 fn main() {
+    let args = Cli::parse();
+
     Logger::try_with_str("info")
         .unwrap()
-        .log_to_file(FileSpec::default())
+        .log_to_file(
+            FileSpec::default().directory(std::env::current_exe().unwrap().parent().unwrap()),
+        )
         .rotate(
             Criterion::Age(Age::Day),
             Naming::Timestamps,
@@ -21,18 +58,18 @@ fn main() {
         .adaptive_format_for_stderr(AdaptiveFormat::Detailed)
         .print_message()
         .duplicate_to_stderr(Duplicate::Warn)
-        .start_with_specfile("logspec.toml")
+        .start_with_specfile(
+            std::env::current_exe()
+                .unwrap()
+                .with_file_name("logspec.toml"),
+        )
         .unwrap();
 
-    let cfg = Config::get();
+    match args.command {
-    log::info!("cfg: {:#?}", cfg);
+        Command::Install => handle_error!("install", service::install()),
-
+        Command::Uninstall => handle_error!("uninstall", service::uninstall()),
-    let token = CancellationToken::new();
+        Command::Run => handle_error!("run", service::run()),
-    let child_token = token.child_token();
+        Command::Start => handle_error!("start", service::start()),
-    let handle = std::thread::spawn(move || server_executor(cfg, child_token));
+        Command::Stop => handle_error!("stop", service::stop()),
-
+    }
-    std::thread::sleep(std::time::Duration::from_secs(10));
-    token.cancel();
-
-    handle.join().unwrap();
 }

src/service.rs

@@ -0,0 +1,183 @@
+use tokio_util::sync::CancellationToken;
+
+use std::{ffi::OsString, sync::mpsc, thread, time::Duration};
+use windows_service::{
+    define_windows_service,
+    service::{
+        ServiceAccess, ServiceControl, ServiceControlAccept, ServiceErrorControl, ServiceExitCode,
+        ServiceInfo, ServiceStartType, ServiceState, ServiceStatus, ServiceType,
+    },
+    service_control_handler::{self, ServiceControlHandlerResult},
+    service_dispatcher,
+    service_manager::{ServiceManager, ServiceManagerAccess},
+    Result,
+};
+
+use crate::config::Config;
+use crate::server::server_executor;
+
+const SERVICE_NAME: &str = "socks5ws_srv";
+const SERVICE_TYPE: ServiceType = ServiceType::OWN_PROCESS;
+const SERVICE_DISPLAY: &str = "socks5ws proxy";
+const SERVICE_DESCRIPTION: &str = "SOCKS5 proxy windows service";
+
+pub fn install() -> windows_service::Result<()> {
+    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 service_info = ServiceInfo {
+        name: SERVICE_NAME.into(),
+        display_name: OsString::from(SERVICE_DISPLAY),
+        service_type: ServiceType::OWN_PROCESS,
+        start_type: ServiceStartType::OnDemand,
+        error_control: ServiceErrorControl::Normal,
+        executable_path: service_binary_path,
+        launch_arguments: vec!["run".into()],
+        dependencies: vec![],
+        account_name: Some(OsString::from(r#"NT AUTHORITY\NetworkService"#)),
+        account_password: None,
+    };
+    let service = service_manager.create_service(&service_info, ServiceAccess::CHANGE_CONFIG)?;
+    service.set_description(SERVICE_DESCRIPTION)?;
+    log::info!("service installed");
+    Ok(())
+}
+
+pub fn uninstall() -> windows_service::Result<()> {
+    let manager_access = ServiceManagerAccess::CONNECT;
+    let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?;
+
+    let service_access = ServiceAccess::QUERY_STATUS | ServiceAccess::STOP | ServiceAccess::DELETE;
+    let service = service_manager.open_service(SERVICE_NAME, service_access)?;
+
+    let service_status = service.query_status()?;
+    if service_status.current_state != ServiceState::Stopped {
+        log::warn!("stopping service");
+        service.stop()?;
+        // Wait for service to stop
+        thread::sleep(Duration::from_secs(5));
+    }
+
+    service.delete()?;
+    log::warn!("service deleted");
+    Ok(())
+}
+
+pub fn stop() -> windows_service::Result<()> {
+    let manager_access = ServiceManagerAccess::CONNECT;
+    let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?;
+
+    let service_access = ServiceAccess::QUERY_STATUS | ServiceAccess::STOP;
+    let service = service_manager.open_service(SERVICE_NAME, service_access)?;
+
+    let service_status = service.query_status()?;
+    if service_status.current_state != ServiceState::Stopped {
+        log::info!("stopping service");
+        service.stop()?;
+    }
+    Ok(())
+}
+
+pub fn start() -> windows_service::Result<()> {
+    let manager_access = ServiceManagerAccess::CONNECT;
+    let service_manager = ServiceManager::local_computer(None::<&str>, manager_access)?;
+
+    let service_access = ServiceAccess::QUERY_STATUS | ServiceAccess::START;
+    let service = service_manager.open_service(SERVICE_NAME, service_access)?;
+
+    let service_status = service.query_status()?;
+    if service_status.current_state != ServiceState::Running {
+        log::info!("start service");
+        service.start(Vec::<&str>::new().as_slice())?;
+    }
+    Ok(())
+}
+
+pub fn run() -> Result<()> {
+    // Register generated `ffi_service_main` with the system and start the service, blocking
+    // this thread until the service is stopped.
+    log::info!("service run");
+    service_dispatcher::start(SERVICE_NAME, ffi_service_main)
+}
+
+// Generate the windows service boilerplate.
+// The boilerplate contains the low-level service entry function (ffi_service_main) that parses
+// incoming service arguments into Vec<OsString> and passes them to user defined service
+// entry (my_service_main).
+define_windows_service!(ffi_service_main, my_service_main);
+
+// Service entry function which is called on background thread by the system with service
+// parameters. There is no stdout or stderr at this point so make sure to configure the log
+// output to file if needed.
+pub fn my_service_main(_arguments: Vec<OsString>) {
+    if let Err(e) = run_service() {
+        log::error!("{:#?}", e);
+    }
+}
+
+pub fn run_service() -> Result<()> {
+    // Create a channel to be able to poll a stop event from the service worker loop.
+    let (shutdown_tx, shutdown_rx) = mpsc::channel();
+
+    // Define system service event handler that will be receiving service events.
+    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::debug!("Stop signal from system");
+                shutdown_tx.send(()).unwrap();
+                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(SERVICE_NAME, event_handler)?;
+
+    // Tell the system that service is running
+    status_handle.set_service_status(ServiceStatus {
+        service_type: SERVICE_TYPE,
+        current_state: ServiceState::Running,
+        controls_accepted: ServiceControlAccept::STOP,
+        exit_code: ServiceExitCode::Win32(0),
+        checkpoint: 0,
+        wait_hint: Duration::default(),
+        process_id: None,
+    })?;
+
+    let cfg = Config::get();
+    log::info!("start with config: {:#?}", cfg);
+
+    let token = CancellationToken::new();
+    let child_token = token.child_token();
+    let server_handle = std::thread::spawn(move || server_executor(cfg, child_token));
+
+    shutdown_rx.recv().unwrap(); // wait for shutdown signal
+    log::info!("service stop");
+
+    // stop server
+    token.cancel();
+    server_handle.join().unwrap();
+
+    // Tell the system that service has stopped.
+    status_handle.set_service_status(ServiceStatus {
+        service_type: SERVICE_TYPE,
+        current_state: ServiceState::Stopped,
+        controls_accepted: ServiceControlAccept::empty(),
+        exit_code: ServiceExitCode::Win32(0),
+        checkpoint: 0,
+        wait_hint: Duration::default(),
+        process_id: None,
+    })?;
+
+    Ok(())
+}