Browse Source

fix command line buffer and test addr

rustdesk 2 years ago
parent
commit
675bf3c1f5
2 changed files with 8 additions and 11 deletions
  1. 1 1
      src/relay_server.rs
  2. 7 10
      src/rendezvous_server.rs

+ 1 - 1
src/relay_server.rs

@@ -379,7 +379,7 @@ async fn handle_connection(
379 379
         let limiter = limiter.clone();
380 380
         tokio::spawn(async move {
381 381
             let mut stream = stream;
382
-            let mut buffer = [0; 64];
382
+            let mut buffer = [0; 1024];
383 383
             if let Ok(Ok(n)) = timeout(1000, stream.read(&mut buffer[..])).await {
384 384
                 if let Ok(data) = std::str::from_utf8(&buffer[..n]) {
385 385
                     let res = check_cmd(data, limiter).await;

+ 7 - 10
src/rendezvous_server.rs

@@ -38,7 +38,6 @@ use std::{
38 38
     sync::Arc,
39 39
     time::Instant,
40 40
 };
41
-const ADDR_127: IpAddr = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
42 41
 
43 42
 #[derive(Clone, Debug)]
44 43
 enum Data {
@@ -93,7 +92,6 @@ impl RendezvousServer {
93 92
     #[tokio::main(flavor = "multi_thread")]
94 93
     pub async fn start(port: i32, serial: i32, key: &str, rmem: usize) -> ResultType<()> {
95 94
         let (key, sk) = Self::get_server_sk(key);
96
-        let addr = format!("0.0.0.0:{port}");
97 95
         let nat_port = port - 1;
98 96
         let ws_port = port + 2;
99 97
         let pm = PeerMap::new().await?;
@@ -162,12 +160,11 @@ impl RendezvousServer {
162 160
             }
163 161
         );
164 162
         if test_addr.to_lowercase() != "no" {
165
-            let test_addr = (if test_addr.is_empty() {
166
-                addr.replace("0.0.0.0", "127.0.0.1")
163
+            let test_addr = if test_addr.is_empty() {
164
+                listener.local_addr()?
167 165
             } else {
168
-                test_addr
169
-            })
170
-            .parse::<SocketAddr>()?;
166
+                test_addr.parse()?
167
+            };
171 168
             tokio::spawn(async move {
172 169
                 allow_err!(test_hbbs(test_addr).await);
173 170
             });
@@ -427,7 +424,7 @@ impl RendezvousServer {
427 424
                     self.handle_local_addr(la, addr, Some(socket)).await?;
428 425
                 }
429 426
                 Some(rendezvous_message::Union::ConfigureUpdate(mut cu)) => {
430
-                    if addr.ip() == ADDR_127 && cu.serial > self.inner.serial {
427
+                    if addr.ip().is_loopback() && cu.serial > self.inner.serial {
431 428
                         let mut inner: Inner = (*self.inner).clone();
432 429
                         inner.serial = cu.serial;
433 430
                         self.inner = Arc::new(inner);
@@ -566,7 +563,7 @@ impl RendezvousServer {
566 563
                 ip != old.socket_addr.ip()
567 564
             } else {
568 565
                 ip.to_string() != old.info.ip
569
-            } && ip != ADDR_127;
566
+            } && !ip.is_loopback();
570 567
             let request_pk = old.pk.is_empty() || ip_change;
571 568
             if !request_pk {
572 569
                 old.socket_addr = socket_addr;
@@ -1053,7 +1050,7 @@ impl RendezvousServer {
1053 1050
         if addr.ip().is_loopback() {
1054 1051
             tokio::spawn(async move {
1055 1052
                 let mut stream = stream;
1056
-                let mut buffer = [0; 64];
1053
+                let mut buffer = [0; 1024];
1057 1054
                 if let Ok(Ok(n)) = timeout(1000, stream.read(&mut buffer[..])).await {
1058 1055
                     if let Ok(data) = std::str::from_utf8(&buffer[..n]) {
1059 1056
                         let res = rs.check_cmd(data).await;