From a3977ba52ad0eb181044e490afabfe926bdecc43 Mon Sep 17 00:00:00 2001 From: alex Date: Tue, 28 Apr 2026 17:19:20 +0800 Subject: [PATCH] fix access-token record error --- Cargo.lock | 94 ++++++++++++++++++++ Cargo.toml | 1 + src/bin/punchnet/api/mod.rs | 30 ++++++- src/bin/punchnet/main.rs | 167 ++++++++++++++++++++++++++---------- src/tcp/quic.rs | 13 +++ src/utils/command.rs | 29 +++++-- 6 files changed, 281 insertions(+), 53 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 88f2970..48ff953 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1874,6 +1874,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "ntapi" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3b335231dfd352ffb0f8017f3b6027a4917f7df785ea2143d8af2adc66980ae" +dependencies = [ + "winapi", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -1957,6 +1966,25 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.11.0", +] + +[[package]] +name = "objc2-io-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33fafba39597d6dc1fb709123dfa8289d39406734be322956a69f0931c73bb15" +dependencies = [ + "libc", + "objc2-core-foundation", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -2286,6 +2314,7 @@ dependencies = [ "simple-dns", "socket2", "structopt", + "sysinfo", "tokio", "tokio-util", "tracing", @@ -3300,6 +3329,20 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "sysinfo" +version = "0.38.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ab6a2f8bfe508deb3c6406578252e491d299cbbf3bc0529ecc3313aee4a52f" +dependencies = [ + "libc", + "memchr", + "ntapi", + "objc2-core-foundation", + "objc2-io-kit", + "windows 0.62.2", +] + [[package]] name = "system-configuration" version = "0.5.1" @@ -4069,6 +4112,27 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" +dependencies = [ + "windows-collections", + "windows-core 0.62.2", + "windows-future", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" +dependencies = [ + "windows-core 0.62.2", +] + [[package]] name = "windows-core" version = "0.52.0" @@ -4091,6 +4155,17 @@ dependencies = [ "windows-strings", ] +[[package]] +name = "windows-future" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" +dependencies = [ + "windows-core 0.62.2", + "windows-link", + "windows-threading", +] + [[package]] name = "windows-implement" version = "0.60.2" @@ -4119,6 +4194,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-numerics" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" +dependencies = [ + "windows-core 0.62.2", + "windows-link", +] + [[package]] name = "windows-result" version = "0.4.1" @@ -4254,6 +4339,15 @@ dependencies = [ "windows_x86_64_msvc 0.53.1", ] +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.42.2" diff --git a/Cargo.toml b/Cargo.toml index a116ddb..2b3103b 100755 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,6 +47,7 @@ simple-dns = "0.11.2" default-net = "0.22.0" socket2 = "0.6.3" hostname = "0.4.2" +sysinfo = "0.38.4" # rolling-file = { path = "../rolling-file" } [target.'cfg(unix)'.dependencies] diff --git a/src/bin/punchnet/api/mod.rs b/src/bin/punchnet/api/mod.rs index 768008a..fceebcc 100644 --- a/src/bin/punchnet/api/mod.rs +++ b/src/bin/punchnet/api/mod.rs @@ -1,5 +1,5 @@ use hmac::{Hmac, Mac as HamcMac}; -use punchnet::{ExitNodeConfiguration, TokenLogin, get_hostname}; +use punchnet::{CachedLoginInfo, ExitNodeConfiguration, TokenLogin, get_hostname, set_access_token}; use reqwest::Client; use sdlan_sn_rs::utils::{Mac, Result, SDLanError}; use serde::{Deserialize, Serialize}; @@ -78,6 +78,34 @@ impl HMacCalculator for UserPassLoginData<'_> { } } +impl TryInto for LoginResponse { + type Error = SDLanError; + fn try_into(self) -> std::result::Result { + if self.code != 0 { + return Err(SDLanError::IOError(format!("code error: {}", self.message))); + } + match self.data { + Some(data) => { + if let Err(_e) = set_access_token(&CachedLoginInfo { + access_token: data.access_token.clone(), + username: data.username.clone(), + user_type: data.user_type.clone(), + audit: data.audit, + network_id: data.network_id, + network_name: data.network_name.clone(), + network_domain: data.network_domain.clone(), + }) { + eprintln!("failed to save access_token"); + } + Ok(data) + }, + None => Err(SDLanError::IOError(format!("data is none: {}", self.message))), + } + + } +} + + #[derive(Debug, Deserialize)] pub struct LoginResponse { pub code: i32, diff --git a/src/bin/punchnet/main.rs b/src/bin/punchnet/main.rs index a25c6be..0711f99 100755 --- a/src/bin/punchnet/main.rs +++ b/src/bin/punchnet/main.rs @@ -3,11 +3,14 @@ mod local_udp_info; use std::collections::HashMap; use std::fs; +use std::path::Path; use std::process; +use std::time::Duration; use clap::Parser; use punchnet::ExitNodeConfiguration; use punchnet::store_configuration; +use sdlan_sn_rs::utils::SDLanError; use std::env; #[cfg(not(target_os = "windows"))] @@ -60,6 +63,8 @@ const APP_USER_ENV_NAME: &str = "PUNCH_USER"; const APP_PASS_ENV_NAME: &str = "PUNCH_PASS"; const APP_TOKEN_ENV_NAME: &str = "PUNCH_TOKEN"; +const PID_FILE: &str = "/tmp/punchnet.pid"; + fn parse_connect_result(res: Result) -> ConnectData { match res { Err(e) => { @@ -297,6 +302,38 @@ const SYSTEM: &'static str = "linux"; #[cfg(target_os = "linux")] const DEFAULT_BASE_DIR: &'static str = "/usr/local/punchnet"; +async fn login_with_token_or_user_pass( + client_id: &str, + mac: Mac, + system: &str, + version: &str, + token: &Option, + user: &Option, + pass: &Option +) -> Result{ + if let Some(ref tk) = token { + login_with_token(TEST_PREFIX, client_id, tk, mac, system, version) + .await?.try_into() + } else { + if let (Some(ref user), Some(ref pass)) = (&user, &pass) { + login_with_user_pass( + TEST_PREFIX, + &client_id, + &user, + &pass, + mac, + system, + version, + ) + .await?.try_into() + } else { + // eprintln!("invalid argument, use --help for help"); + Err(SDLanError::IOError("Invalid argument, use --help for help".to_string())) + // process::exit(-1); + } + } +} + fn main() { set_base_dir(DEFAULT_BASE_DIR); // let _guard = log::init_log(&format!("{}/.output", get_base_dir())); @@ -320,23 +357,20 @@ fn main() { Commands::Login(user) => { let rt = Runtime::new().unwrap(); rt.block_on(async move { - let _ = parse_login_result( - login_with_user_pass( - TEST_PREFIX, - &client_id, - &user.username, - &user.password, - mac, - system, - version, - ) - .await, - ); + if is_process_running() { + eprintln!("process is running, aborting..."); + std::process::exit(-1); + } + + if let Err(e) = login_with_token_or_user_pass(&client_id, mac, system, version, &user.token, &user.username, &user.password).await { + eprintln!("failed to login: {}", e.as_str()); + std::process::exit(-1); + } println!("login ok"); }); process::exit(0); - // TODO: do login with user } + /* Commands::TokenLogin(tk) => { let rt = Runtime::new().unwrap(); rt.block_on(async move { @@ -348,6 +382,7 @@ fn main() { }); process::exit(0); } + */ Commands::ExitNode(cmd) => { let rt = Runtime::new().unwrap(); @@ -369,7 +404,7 @@ fn main() { #[cfg(not(target_os = "windows"))] Commands::Stop => { - match fs::read_to_string("/tmp/punchnet.pid") { + match fs::read_to_string(PID_FILE) { Ok(content) => { let pid: i32 = match content.trim().parse() { Ok(value) => value, @@ -404,11 +439,19 @@ fn main() { process::exit(0); } Commands::AutoRun(runinfo) => { + if is_process_running() { + eprintln!("process is running, aborting..."); + std::process::exit(-1); + } if runinfo.foreground { should_daemonize = false; } } Commands::Start(startinfo) => { + if is_process_running() { + eprintln!("process is running, aborting..."); + std::process::exit(-1); + } if startinfo.foreground { should_daemonize = false; } @@ -418,7 +461,6 @@ fn main() { } } - #[cfg(not(target_os = "windows"))] if should_daemonize { use std::fs::OpenOptions; @@ -436,7 +478,7 @@ fn main() { .unwrap(); let daemonize = Daemonize::new() - .pid_file("/tmp/punchnet.pid") + .pid_file(PID_FILE) .chown_pid_file(true) .working_directory(get_base_dir()) .stdout(out) @@ -452,6 +494,8 @@ fn main() { } } } else { + let pid = std::process::id(); + record_pid_file(pid); run_it(cmd, client_id, allow_routing, mac, system, version); } @@ -488,6 +532,7 @@ fn run_it( let rt = Runtime::new().unwrap(); match &cmd.cmd { Commands::Start(rtinfo) => rt.block_on(async move { + let remembered_token = get_access_token(); if remembered_token.is_none() { eprintln!("not logged in, should login with user/pass or token first"); @@ -503,7 +548,7 @@ fn run_it( rtinfo.allow_routing || allow_routing, connect_info, remembered, - client_id, + client_id.clone(), rtinfo.route_file.clone(), rtinfo.route.clone(), mac, @@ -512,13 +557,18 @@ fn run_it( .await; }), Commands::AutoRun(tk) => rt.block_on(async move { - let mut remembered_token = get_access_token(); - if remembered_token.is_none() { - let data = parse_login_result( - login_with_token(TEST_PREFIX, &client_id, &tk.token, mac, system, version) - .await, - ); - remembered_token = Some(CachedLoginInfo { + loop { + let data = match login_with_token_or_user_pass(&client_id, mac, system, version, &tk.token, &tk.username, &tk.password).await { + Ok(data) => { + data + } + Err(e) => { + eprintln!("failed to login: {}, will try in 10 seconds", e.as_str()); + tokio::time::sleep(Duration::from_secs(10)).await; + continue; + } + }; + let remembered_token = Some(CachedLoginInfo { access_token: data.access_token, username: data.username, user_type: data.user_type, @@ -527,35 +577,66 @@ fn run_it( network_name: data.network_name, network_domain: data.network_domain, }); + + let remembered = remembered_token.unwrap(); + + let connect_info = parse_connect_result( + connect(TEST_PREFIX, &client_id, &remembered.access_token).await, + ); + + record_exit_node(&connect_info); + daemonize_me( + tk.allow_routing || allow_routing, + connect_info, + remembered, + client_id, + tk.route_file.clone(), + tk.route.clone(), + mac, + tk.take_dns, + ) + .await; + break; } - let remembered = remembered_token.unwrap(); - - let connect_info = parse_connect_result( - connect(TEST_PREFIX, &client_id, &remembered.access_token).await, - ); - - record_exit_node(&connect_info); - daemonize_me( - tk.allow_routing || allow_routing, - connect_info, - remembered, - client_id, - tk.route_file.clone(), - tk.route.clone(), - mac, - tk.take_dns, - ) - .await; }), - other => { + _other => { eprintln!("should not comes here"); process::exit(-1); } } } +pub fn record_pid_file(pid: u32) { + fs::write(PID_FILE, pid.to_string()); +} + pub fn delete_pid_file() { fs::remove_file("/tmp/punchnet.pid"); } + +fn read_pid_file() -> Option { + let content = fs::read_to_string(PID_FILE).ok()?; + content.trim().parse().ok() +} + +fn is_pid_running(pid: u32) -> bool { + let mut sys = sysinfo::System::new_all(); + sys.refresh_all(); + sys.process(sysinfo::Pid::from_u32(pid)).is_some() +} + +pub fn is_process_running() -> bool{ + if Path::new(PID_FILE).exists() { + if let Some(pid) = read_pid_file() { + if !is_pid_running(pid) { + fs::remove_file(PID_FILE); + return false; + } else { + return true; + } + } + } + false +} \ No newline at end of file diff --git a/src/tcp/quic.rs b/src/tcp/quic.rs index c3dd878..70b7a4c 100644 --- a/src/tcp/quic.rs +++ b/src/tcp/quic.rs @@ -664,6 +664,18 @@ impl ReadWriteActor { // stream.write("hello".as_bytes()).await; // let (reader, mut write) = stream.into_split(); + let tcp_pong_timeout = async move { + loop { + tokio::time::sleep(Duration::from_secs(10)).await; + let now = get_current_timestamp(); + let last = edge.tcp_pong.load(Ordering::Relaxed); + if now - last > 30 { + println!("tcp pong timedout"); + break; + } + } + }; + let read_from_tcp = async move { // let mut buffed_reader = BufReader::new(recv); loop { @@ -738,6 +750,7 @@ impl ReadWriteActor { _ = write_to_tcp => {}, _ = check_pong => {}, _ = check_stop => {}, + _ = tcp_pong_timeout => {}, } on_disconnected_callback().await; conn.close(0u32.into(), "close".as_bytes()); diff --git a/src/utils/command.rs b/src/utils/command.rs index fff6fe3..6ac033f 100755 --- a/src/utils/command.rs +++ b/src/utils/command.rs @@ -14,7 +14,7 @@ pub struct CommandLineInput2 { #[derive(Subcommand, Debug)] pub enum Commands { Login(UserLogin), - TokenLogin(TokenLogin), + // TokenLogin(TokenLogin), /// if logined in, just start, /// else, use the token to login, and start @@ -80,26 +80,37 @@ pub struct RouteCmdInfo { #[derive(Args, Debug)] pub struct UserLogin { - #[arg(short, long, env = APP_USER_ENV_NAME)] - pub username: String, + #[arg(short='u', long, env = APP_USER_ENV_NAME, help="specify username")] + pub username: Option, - #[arg(short, long, env = APP_PASS_ENV_NAME, required=false)] - pub password: String, + #[arg(short='p', long, env = APP_PASS_ENV_NAME, help="specify user's password")] + pub password: Option, + + #[arg(short='t', long, env=APP_TOKEN_ENV_NAME, help="login with a token")] + pub token: Option, } +/* #[derive(Args, Debug)] pub struct AutoRunTokenLogin { - #[arg(long, env=APP_TOKEN_ENV_NAME, required=false)] + #[arg(short='t', long, env=APP_TOKEN_ENV_NAME)] pub token: String, - #[arg(short, long, default_value="")] + #[arg(short='r', long, default_value="")] pub route: String, } +*/ #[derive(Args, Debug)] pub struct AutoRunArgument { - #[arg(long, env=APP_TOKEN_ENV_NAME, required=false)] - pub token: String, + #[arg(short='t', long, env=APP_TOKEN_ENV_NAME)] + pub token: Option, + + #[arg(short='u', long, env = APP_USER_ENV_NAME, help="specify username")] + pub username: Option, + + #[arg(short='p', long, env = APP_PASS_ENV_NAME, help="specify user's password")] + pub password: Option, #[arg(short, long, default_value_t=false)] pub allow_routing: bool,