b4tman/socks5ws
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

windows service

Browse Source
This commit is contained in:
Dmitry Belyaev
2022-09-24 21:11:12 +03:00
parent 7f6ddb0562
commit f65d24b57c
4 changed files with 418 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
src/main.rs
View File

@@ -1,18 +1,55 @@
extern crate flexi_logger;
use tokio_util::sync::CancellationToken;
mod config;
mod server;
use crate::config::Config;
use crate::server::server_executor;
mod service;
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();
log::info!("cfg: {:#?}", cfg);
let token = CancellationToken::new();
let child_token = token.child_token();
let handle = std::thread::spawn(move || server_executor(cfg, child_token));
std::thread::sleep(std::time::Duration::from_secs(10));
token.cancel();
handle.join().unwrap();
match args.command {
Command::Install => handle_error!("install", service::install()),
Command::Uninstall => handle_error!("uninstall", service::uninstall()),
Command::Run => handle_error!("run", service::run()),
Command::Start => handle_error!("start", service::start()),
Command::Stop => handle_error!("stop", service::stop()),
}
}

183
src/service.rs Normal file
View File

@@ -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(())
}