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

+ 7 - 10
src/rendezvous_server.rs

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