From bc980a3b84c6d30b8d415c27fed6d0155ba3472c Mon Sep 17 00:00:00 2001 From: lejianwen <84855512@qq.com> Date: Wed, 15 Jan 2025 19:19:18 +0800 Subject: [PATCH] feat: check token via jwt --- Cargo.lock | 5 ++-- Cargo.toml | 2 +- README.md | 1 + src/jwt.rs | 54 ++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/rendezvous_server.rs | 31 +++++++++++++++++------ tests/jwt_tests.rs | 24 ++++++++++++++++++ 7 files changed, 107 insertions(+), 11 deletions(-) create mode 100644 src/jwt.rs create mode 100644 tests/jwt_tests.rs diff --git a/Cargo.lock b/Cargo.lock index eb4d034..20eb27a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "ahash" @@ -225,6 +225,7 @@ dependencies = [ "libc", "num-integer", "num-traits", + "serde", "time 0.1.43", "winapi", ] @@ -779,7 +780,7 @@ dependencies = [ [[package]] name = "hbbs" -version = "1.1.12" +version = "1.1.13" dependencies = [ "async-speed-limit", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 792b378..d2caf81 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,7 +36,7 @@ async-trait = "0.1" async-speed-limit = { git = "https://github.com/open-trade/async-speed-limit" } uuid = { version = "1.0", features = ["v4"] } bcrypt = "0.13" -chrono = "0.4" +chrono = { version = "0.4", features = ["serde"] } jsonwebtoken = "8" headers = "0.3" once_cell = "1.8" diff --git a/README.md b/README.md index d71641c..05d36c6 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ - 解决当客户端登录了`Api`账号时链接超时的问题 - s6镜像添加了`Api`支持，`Api`开源地址 https://github.com/lejianwen/rustdesk-api - 是否必须登录才能链接， `MUST_LOGIN` 默认为 `N`，设置为 `Y` 则必须登录才能链接 +- `RUSTDESK_API_JWT_KEY`，设置后会通过`JWT`校验token的合法性 ## docker镜像地址 diff --git a/src/jwt.rs b/src/jwt.rs new file mode 100644 index 0000000..c11cabb --- /dev/null +++ b/src/jwt.rs @@ -0,0 +1,54 @@ +use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation}; +use once_cell::sync::Lazy; +use serde::{Deserialize, Serialize}; +use std::env; +pub static SECRET: Lazy = + Lazy::new(|| env::var("RUSTDESK_API_JWT_KEY").unwrap_or_else(|_| "".to_string())); + +// 定义一个结构体来表示 JWT 的 payload +#[derive(Debug, Serialize, Deserialize)] +pub struct Claims { + user_id: u32, + exp: usize, +} + +pub fn generate_token(user_id: u32, exp: i64) -> Result { + println!("secret: {:}", SECRET.to_string()); + let claims = Claims { + user_id, + exp: (chrono::Utc::now() + chrono::Duration::seconds(exp)).timestamp() as usize, + }; + + let token = encode( + &Header::default(), + &claims, + &EncodingKey::from_secret(SECRET.as_ref()), + ); + + match token { + Ok(t) => Ok(t), + Err(e) => Err(e.to_string()), + } +} +// 验证 JWT 的函数 +pub fn verify_token(token: &str) -> Result { + // 解码 JWT + let validation = Validation::new(Algorithm::HS256); + + let decoded = decode::( + &token, + &DecodingKey::from_secret(SECRET.as_ref()), + &validation, + ); + match decoded { + Ok(token_data) => { + let now = chrono::Utc::now().timestamp() as usize; + if token_data.claims.exp > now { + Ok(token_data.claims) + } else { + Err("Token status invalid or expired".to_string()) + } + } + Err(_) => Err("Invalid token".to_string()), + } +} diff --git a/src/lib.rs b/src/lib.rs index 8da29a2..78aa01f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,3 +4,4 @@ pub mod common; mod database; mod peer; mod version; +pub mod jwt; \ No newline at end of file diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs index 53ae586..3b59a01 100644 --- a/src/rendezvous_server.rs +++ b/src/rendezvous_server.rs @@ -46,6 +46,7 @@ use std::{ sync::Arc, time::Instant, }; +use crate::jwt; #[derive(Clone, Debug)] enum Data { @@ -776,14 +777,28 @@ impl RendezvousServer { }); return Ok((msg_out, None)); } - // Todo check token by jwt - if ph.token.is_empty() && MUST_LOGIN.load(Ordering::SeqCst) { - let mut msg_out = RendezvousMessage::new(); - msg_out.set_punch_hole_response(PunchHoleResponse { - other_failure: String::from("Connection failed, please login first"), - ..Default::default() - }); - return Ok((msg_out, None)); + // if secret is not empty check token by jwt + if MUST_LOGIN.load(Ordering::SeqCst) { + if ph.token.is_empty() { + let mut msg_out = RendezvousMessage::new(); + msg_out.set_punch_hole_response(PunchHoleResponse { + other_failure: String::from("Connection failed, please login!"), + ..Default::default() + }); + return Ok((msg_out, None)); + } else if !jwt::SECRET.is_empty() { + let token = ph.token; + let token = jwt::verify_token(token.as_str()); + if token.is_err() { + let mut msg_out = RendezvousMessage::new(); + msg_out.set_punch_hole_response(PunchHoleResponse { + //提示重新登录 + other_failure: String::from("Token error, please log out and log back in!"), + ..Default::default() + }); + return Ok((msg_out, None)); + } + } } let id = ph.id; // punch hole request from A, relay to B, diff --git a/tests/jwt_tests.rs b/tests/jwt_tests.rs new file mode 100644 index 0000000..cce097f --- /dev/null +++ b/tests/jwt_tests.rs @@ -0,0 +1,24 @@ +use hbb_common::tokio; +use hbbs::jwt; + +#[test] +fn test_generate_token() { + std::env::set_var("RUSTDESK_API_JWT_KEY", "testjwt"); + let token = jwt::generate_token(1, 3600).unwrap(); + println!("Generated Token: {}", token); + assert!(!token.is_empty(), "Generated token should not be empty"); +} + +#[tokio::test] +async fn test_verify_token() { + std::env::set_var("RUSTDESK_API_JWT_KEY", "testjwt"); + let token = jwt::generate_token(1, 2).unwrap(); + // let token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoxLCJleHAiOjE3MzY4NzA4NjF9.u5pxmwNMrYUwtkspF1FuZj-R5ANAR9WT9_dMHuQhV0Y"; + println!("Token : {:?}, now: {:?}", token, chrono::Utc::now().timestamp()); + + // hbb_common::sleep(3f32).await; + // println!("Token : {:?}, now: {:?}", token, chrono::Utc::now().timestamp()); + let result = jwt::verify_token(&token); + println!("Token Verification Result: {:?}", result); + assert!(result.is_ok(), "Token should be valid"); +} \ No newline at end of file