feat: check token via jwt

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",

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"

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镜像地址
 

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<String> =
+    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<String, String> {
+    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<Claims, String> {
+    // 解码 JWT
+    let validation = Validation::new(Algorithm::HS256);
+
+    let decoded = decode::<Claims>(
+        &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()),
+    }
+}

src/lib.rs

@@ -4,3 +4,4 @@ pub mod common;
 mod database;
 mod peer;
 mod version;
+pub mod jwt;

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) {
+        // 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 first"),
+                    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,

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");
+}