优化代码

This commit is contained in:
lubeilin
2023-03-19 20:11:51 +08:00
parent d7226d8522
commit 7840f30b64
14 changed files with 174 additions and 99 deletions
+3 -4
View File
@@ -156,10 +156,9 @@ impl ArgsConfig {
}
}
}
use fd_lock::RwLock;
pub fn lock_config() -> io::Result<RwLock<File>> {
let config_path = SWITCH_HOME_PATH.lock().clone().unwrap().join("config");
Ok(RwLock::new(File::open(config_path)?))
pub fn lock_file() -> io::Result<File> {
let path = SWITCH_HOME_PATH.lock().clone().unwrap().join(".lock");
Ok(File::create(path)?)
}
pub fn save_config(config: ArgsConfig) -> io::Result<()> {
+85 -18
View File
@@ -1,4 +1,4 @@
use console::style;
use console::{style, Style};
use crate::command::entity::{DeviceItem, RouteItem, Status};
@@ -16,51 +16,118 @@ pub fn console_status(status: Status) {
println!("Local ip: {}", style(status.local_ip).green());
}
pub fn console_route_table(list: Vec<RouteItem>) {
pub fn console_route_table(mut list: Vec<RouteItem>) {
if list.is_empty() {
println!("No route found");
return;
}
list.sort_by(|t1, t2| t1.destination.cmp(&t2.destination));
let mut out_list = Vec::with_capacity(list.len());
//表头
out_list.push(vec!["Destination".to_string(), "Next Hop".to_string(), "Metric".to_string(),
"Rt".to_string(), "Interface".to_string()]);
out_list.push(vec![("Destination".to_string(), Style::new()),
("Next Hop".to_string(), Style::new()),
("Metric".to_string(), Style::new()),
("Rt".to_string(), Style::new()),
("Interface".to_string(), Style::new()), ]);
for item in list {
out_list.push(vec![item.destination, item.next_hop, item.metric,
item.rt, item.interface]);
out_list.push(vec![(item.destination, Style::new().green()),
(item.next_hop, Style::new().green()),
(item.metric, Style::new().green()),
(item.rt, Style::new().green()),
(item.interface, Style::new().green())]);
}
table::println_table(out_list)
}
pub fn console_device_list(list: Vec<DeviceItem>) {
pub fn console_device_list(mut list: Vec<DeviceItem>) {
if list.is_empty() {
println!("No other devices found");
return;
}
list.sort_by(|t1, t2| t1.virtual_ip.cmp(&t2.virtual_ip));
list.sort_by(|t1, t2| t1.status.cmp(&t2.status));
let mut out_list = Vec::with_capacity(list.len());
//表头
out_list.push(vec!["Name".to_string(), "Virtual Ip".to_string(), "P2P/Relay".to_string(), "Rt".to_string(), "Status".to_string()]);
out_list.push(vec![("Name".to_string(), Style::new()),
("Virtual Ip".to_string(), Style::new()),
("Status".to_string(), Style::new()),
("P2P/Relay".to_string(), Style::new()),
("Rt".to_string(), Style::new())]);
for item in list {
out_list.push(vec![item.name, item.virtual_ip, item.nat_traversal_type,
item.rt, item.status]);
if &item.status == "Online" {
if &item.nat_traversal_type == "p2p" {
out_list.push(vec![(item.name, Style::new().green()),
(item.virtual_ip, Style::new().green()),
(item.status, Style::new().green()),
(item.nat_traversal_type, Style::new().green()),
(item.rt, Style::new().green())]);
} else {
out_list.push(vec![(item.name, Style::new().yellow()),
(item.virtual_ip, Style::new().yellow()),
(item.status, Style::new().yellow()),
(item.nat_traversal_type, Style::new().yellow()),
(item.rt, Style::new().yellow())]);
}
} else {
out_list.push(vec![(item.name, Style::new().color256(102)),
(item.virtual_ip, Style::new().color256(102)),
(item.status, Style::new().color256(102)),
("".to_string(), Style::new().color256(102)),
("".to_string(), Style::new().color256(102))]);
}
}
table::println_table(out_list)
}
pub fn console_device_list_all(list: Vec<DeviceItem>) {
pub fn console_device_list_all(mut list: Vec<DeviceItem>) {
if list.is_empty() {
println!("No other devices found");
return;
}
list.sort_by(|t1, t2| t1.virtual_ip.cmp(&t2.virtual_ip));
list.sort_by(|t1, t2| t1.status.cmp(&t2.status));
let mut out_list = Vec::with_capacity(list.len());
//表头
out_list.push(vec!["Name".to_string(), "Virtual Ip".to_string(), "NAT Type".to_string(),
"Public Ips".to_string(), "Local Ip".to_string(), "P2P/Relay".to_string(),
"Rt".to_string(), "Status".to_string()]);
out_list.push(vec![("Name".to_string(), Style::new()),
("Virtual Ip".to_string(), Style::new()),
("Status".to_string(), Style::new()),
("NAT Type".to_string(), Style::new()),
("Public Ips".to_string(), Style::new()),
("Local Ip".to_string(), Style::new()),
("P2P/Relay".to_string(), Style::new()),
("Rt".to_string(), Style::new())]);
for item in list {
out_list.push(vec![item.name, item.virtual_ip, item.nat_type,
item.public_ips, item.local_ip, item.nat_traversal_type,
item.rt, item.status]);
if &item.status == "Online" {
if &item.nat_traversal_type == "p2p" {
out_list.push(vec![(item.name, Style::new().green()),
(item.virtual_ip, Style::new().green()),
(item.status, Style::new().green()),
(item.nat_traversal_type, Style::new().green()),
(item.rt, Style::new().green()),
(item.nat_type, Style::new().green()),
(item.public_ips, Style::new().green()),
(item.local_ip, Style::new().green())]);
} else {
out_list.push(vec![(item.name, Style::new().yellow()),
(item.virtual_ip, Style::new().yellow()),
(item.status, Style::new().yellow()),
(item.nat_traversal_type, Style::new().yellow()),
(item.rt, Style::new().yellow()),
(item.nat_type, Style::new().yellow()),
(item.public_ips, Style::new().yellow()),
(item.local_ip, Style::new().yellow()), ]);
}
} else {
out_list.push(vec![(item.name, Style::new().color256(102)),
(item.virtual_ip, Style::new().color256(102)),
(item.status, Style::new().color256(102)),
("".to_string(), Style::new().color256(102)),
("".to_string(), Style::new().color256(102)),
("".to_string(), Style::new().color256(102)),
("".to_string(), Style::new().color256(102)),
("".to_string(), Style::new().color256(102)), ]);
}
}
table::println_table(out_list)
}
+6 -12
View File
@@ -1,29 +1,23 @@
use console::style;
use console::Style;
pub fn println_table(table: Vec<Vec<String>>) {
pub fn println_table(table: Vec<Vec<(String, Style)>>) {
if table.is_empty() {
return;
}
let mut width_list = vec![0; table[0].len()];
for in_list in table.iter() {
for (index, item) in in_list.iter().enumerate() {
for (index, (item, _)) in in_list.iter().enumerate() {
let width = console::measure_text_width(item) + 6;
if width_list[index] < width {
width_list[index] = width;
}
}
}
let mut head = true;
for in_list in table {
for (index, item) in in_list.iter().enumerate() {
if head {
print!("{item:width$}", item = item, width = width_list[index]);
} else {
let str = format!("{item:width$}", item = item, width = width_list[index]);
print!("{}", style(str).green());
}
for (col, (item, style)) in in_list.iter().enumerate() {
let str = format!("{:1$}", item, width_list[col]);
print!("{}", style.apply_to(str));
}
head = false;
println!()
}
}
+16 -12
View File
@@ -4,8 +4,6 @@ use console::style;
use switch::core::Switch;
use crate::config::log_config::log_init;
#[cfg(target_os = "windows")]
use crate::config::log_config::log_service_init;
mod command;
mod config;
@@ -86,13 +84,14 @@ pub struct StartArgs {
/// NAT detection service address. Use comma to separate
#[arg(long)]
nat_test_server: Option<String>,
/// 命令服务,开启后可以在其他进程使用 route、list等命令查看信息
/// 程序使用后台运行时需要增加此参数
/// Command service. After it is enabled, you can use route, list and other commands in other processes to view information.
/// This parameter needs to be added when the program is running in the background
/// 关闭命令服务,关闭后不能在其他进程直接使用route、list等命令查看信息
/// Turn off the command service. After turning off, you cannot directly use the route, list and other commands to view information in other processes
#[cfg(any(unix))]
#[arg(long)]
command_server: bool,
off_command_server: bool,
/// 记录日志,输出在 home/.switch 目录下,长时间使用时不建议开启
/// Output the log in the "home/.switch" directory
log: bool,
}
#[cfg(target_os = "windows")]
@@ -124,16 +123,17 @@ fn main() {
if args.len() == 3 && args[1] == windows::SERVICE_FLAG {
//以服务的方式启动
config::set_home(std::path::PathBuf::from(&args[2]));
let _ = log_service_init();
log::info!("config {:?}", std::path::PathBuf::from(&args[2]));
log::info!("config {:?}", config::read_config());
windows::service::start();
return;
} else {
let home = dirs::home_dir().unwrap().join(".switch");
config::set_home(home);
let _ = log_init();
let args = BaseArgs::parse();
if let Commands::Start(start_args) = &args.command {
if start_args.log {
let _ = log_init();
}
}
windows::main0(args);
}
}
@@ -142,8 +142,12 @@ fn main() {
fn main() {
let home = dirs::home_dir().unwrap().join(".switch");
config::set_home(home);
let _ = log_init();
let args = BaseArgs::parse();
if let Commands::Start(start_args) = &args.command {
if start_args.log {
let _ = log_init();
}
}
unix::main0(args);
}
+14 -16
View File
@@ -1,6 +1,7 @@
use std::sync::Arc;
use console::style;
use fs2::FileExt;
use switch::core::{Config, Switch};
@@ -11,7 +12,7 @@ use crate::command::{command, CommandEnum};
pub fn main0(base_args: BaseArgs) {
match base_args.command {
Commands::Start(args) => {
let open_command_server = args.command_server;
let off_command_server = args.off_command_server;
match config::default_config(args) {
Ok(start_config) => {
if sudo::RunningAs::Root != sudo::check() {
@@ -37,7 +38,7 @@ pub fn main0(base_args: BaseArgs) {
nat_test_server,
start_config.device_id.clone(),
);
let mut lock = match config::lock_config() {
let lock = match config::lock_file() {
Ok(lock) => {
lock
}
@@ -46,17 +47,13 @@ pub fn main0(base_args: BaseArgs) {
return;
}
};
let lock_guard = match lock.try_write() {
Ok(lock) => {
lock
}
Err(_) => {
println!("{}", style("文件被重复打开").red());
return;
}
};
if lock.try_lock_exclusive().is_err() {
println!("{}", style("文件被重复打开").red());
return;
}
if let Err(e) = config::save_config(args_config) {
log::error!("{:?}",e);
lock.unlock().unwrap();
return;
}
let switch = match Switch::start(config) {
@@ -65,12 +62,16 @@ pub fn main0(base_args: BaseArgs) {
}
Err(e) => {
log::error!("{:?}", e);
lock.unlock().unwrap();
return;
}
};
let switch = Arc::new(switch);
let command_server = crate::command::server::CommandServer::new();
if open_command_server {
if off_command_server {
crate::console_listen(&switch);
log::info!("前台任务结束");
} else {
if let Err(e) = config::update_pid(std::process::id()) {
log::error!("{:?}", e);
}
@@ -86,11 +87,8 @@ pub fn main0(base_args: BaseArgs) {
} else {
log::info!("后台任务结束");
}
} else {
crate::console_listen(&switch);
log::info!("前台任务结束");
}
drop(lock_guard)
lock.unlock().unwrap();
}
Err(e) => {
log::error!("{:?}", e);
+18 -15
View File
@@ -5,6 +5,7 @@ use std::path::PathBuf;
use std::time::Duration;
use console::style;
use fs2::FileExt;
use windows_service::Error;
use windows_service::service::{
ServiceAccess, ServiceErrorControl, ServiceInfo, ServiceStartType, ServiceState, ServiceType,
@@ -61,6 +62,7 @@ pub fn main0(base_args: BaseArgs) {
// 允许应用通过防火墙
let _udp = UdpSocket::bind("0.0.0.0:0").unwrap();
}
let out_log = args.log;
match config::default_config(args) {
Ok(start_config) => {
match service_state() {
@@ -76,7 +78,7 @@ pub fn main0(base_args: BaseArgs) {
log::error!("{:?}",e);
return;
}
match start() {
match start(out_log) {
Ok(_) => {
//需要检查启动状态
thread::sleep(Duration::from_secs(2));
@@ -107,22 +109,19 @@ pub fn main0(base_args: BaseArgs) {
start_config.server,
start_config.nat_test_server,
);
let mut lock = match config::lock_config() {
Ok(lock) => lock,
let lock = match config::lock_file() {
Ok(lock) => {
lock
}
Err(e) => {
log::error!("{:?}",e);
return;
}
};
let lock_guard = match lock.try_write() {
Ok(lock) => {
lock
}
Err(_) => {
println!("{}", style("程序文件被重复打开").red());
return;
}
};
if lock.try_lock_exclusive().is_err() {
println!("{}", style("文件被重复打开").red());
return;
}
match Switch::start(config) {
Ok(switch) => {
crate::console_listen(&switch);
@@ -131,7 +130,7 @@ pub fn main0(base_args: BaseArgs) {
log::error!("{:?}", e);
}
}
drop(lock_guard);
lock.unlock().unwrap();
return;
}
}
@@ -329,11 +328,15 @@ fn uninstall() -> Result<(), Error> {
Ok(())
}
fn start() -> Result<(), Error> {
fn start(out_log: bool) -> 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(&[""])
if out_log {
service.start(&["log"])
} else {
service.start(&[""])
}
}
fn service_state() -> Result<ServiceState, Error> {
+8 -1
View File
@@ -20,7 +20,14 @@ use crate::windows::config::read_config;
use crate::windows::SERVICE_NAME;
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>) {
if !arguments.is_empty() {
if let Some(str) = arguments[0].to_str() {
if str == "log" {
let _ = config::log_config::log_service_init();
}
}
}
thread::spawn(|| match service_main() {
Ok(_) => {}
Err(e) => {