From 4e37dc8090350d36fb8e37c0aa5fb106e082293b Mon Sep 17 00:00:00 2001
From: lejianwen <84855512@qq.com>
Date: Sun, 12 Jan 2025 20:07:38 +0800
Subject: [PATCH] feat: Add MUST_LOGIN Env

---
 src/main.rs              |  3 ++-
 src/rendezvous_server.rs | 51 +++++++++++++++++++++++++++++++++++++---
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 2cbd9c9..1ed6124 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -22,7 +22,8 @@ fn main() -> ResultType<()> {
         -r, --relay-servers=[HOST] 'Sets the default relay servers, separated by comma'
         -M, --rmem=[NUMBER(default={RMEM})] 'Sets UDP recv buffer size, set system rmem_max first, e.g., sudo sysctl -w net.core.rmem_max=52428800. vi /etc/sysctl.conf, net.core.rmem_max=52428800, sudo sysctl –p'
         , --mask=[MASK] 'Determine if the connection comes from LAN, e.g. 192.168.0.0/16'
-        -k, --key=[KEY] 'Only allow the client with the same key'",
+        -k, --key=[KEY] 'Only allow the client with the same key'
+        , --must-login=[Y|N] 'Only allow the client with login'",
     );
     init_args(&args, "hbbs", "RustDesk ID/Rendezvous Server");
     let port = get_arg_or("port", RENDEZVOUS_PORT.to_string()).parse::<i32>()?;
diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs
index b43f165..53ae586 100644
--- a/src/rendezvous_server.rs
+++ b/src/rendezvous_server.rs
@@ -71,6 +71,7 @@ static ROTATION_RELAY_SERVER: AtomicUsize = AtomicUsize::new(0);
 type RelayServers = Vec<String>;
 const CHECK_RELAY_TIMEOUT: u64 = 3_000;
 static ALWAYS_USE_RELAY: AtomicBool = AtomicBool::new(false);
+static MUST_LOGIN: AtomicBool = AtomicBool::new(false);
 
 #[derive(Clone)]
 struct Inner {
@@ -175,6 +176,25 @@ impl RendezvousServer {
                 "N"
             }
         );
+
+        let must_login = get_arg("must-login");
+        log::debug!("must_login={}", must_login);
+        if must_login.to_uppercase() == "Y" ||
+            (must_login == "" && std::env::var("MUST_LOGIN")
+            .unwrap_or_default()
+            .to_uppercase()
+            == "Y") {
+            MUST_LOGIN.store(true, Ordering::SeqCst);
+        }
+
+        log::info!(
+            "MUST_LOGIN={}",
+            if MUST_LOGIN.load(Ordering::SeqCst) {
+                "Y"
+            } else {
+                "N"
+            }
+        );
         if test_addr.to_lowercase() != "no" {
             let test_addr = if test_addr.is_empty() {
                 listener.local_addr()?
@@ -756,6 +776,15 @@ 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));
+        }
         let id = ph.id;
         // punch hole request from A, relay to B,
         // check if in same intranet first,
@@ -988,13 +1017,14 @@ impl RendezvousServer {
         match fds.next() {
             Some("h") => {
                 res = format!(
-                    "{}\n{}\n{}\n{}\n{}\n{}\n",
+                    "{}\n{}\n{}\n{}\n{}\n{}\n{}\n",
                     "relay-servers(rs) <separated by ,>",
                     "reload-geo(rg)",
                     "ip-blocker(ib) [<ip>|<number>] [-]",
                     "ip-changes(ic) [<id>|<number>] [-]",
-                    "always-use-relay(aur)",
-                    "test-geo(tg) <ip1> <ip2>"
+                    "always-use-relay(aur) [Y|N]",
+                    "test-geo(tg) <ip1> <ip2>",
+                    "must-login(ml) [Y|N]",
                 )
             }
             Some("relay-servers" | "rs") => {
@@ -1121,6 +1151,21 @@ impl RendezvousServer {
                     }
                 }
             }
+            Some("must-login" | "ml") => {
+                if let Some(rs) = fds.next() {
+                    if rs.to_uppercase() == "Y" {
+                        MUST_LOGIN.store(true, Ordering::SeqCst);
+                    } else {
+                        MUST_LOGIN.store(false, Ordering::SeqCst);
+                    }
+                } else {
+                    let _ = writeln!(
+                        res,
+                        "MUST_LOGIN: {:?}",
+                        MUST_LOGIN.load(Ordering::SeqCst)
+                    );
+                }
+            }
             _ => {}
         }
         res