rustdesk лет назад: 3
Родитель
Сommit
fc83fa0a04
2 измененных файлов с 24 добавлено и 15 удалено
  1. 17 10
      libs/hbb_common/src/lib.rs
  2. 7 5
      src/rendezvous_server.rs

+ 17 - 10
libs/hbb_common/src/lib.rs

@@ -95,19 +95,24 @@ pub type ResultType<F, E = anyhow::Error> = anyhow::Result<F, E>;
95
 
95
 
96
 pub struct AddrMangle();
96
 pub struct AddrMangle();
97
 
97
 
98
+#[inline]
99
+pub fn try_into_v4(addr: SocketAddr) -> SocketAddr {
100
+    match addr {
101
+        SocketAddr::V6(v6) if !addr.ip().is_loopback() => {
102
+            if let Some(v4) = v6.ip().to_ipv4() {
103
+                SocketAddr::new(IpAddr::V4(v4), addr.port())
104
+            } else {
105
+                addr
106
+            }
107
+        }
108
+        _ => addr,
109
+    }
110
+}
111
+
98
 impl AddrMangle {
112
 impl AddrMangle {
99
     pub fn encode(addr: SocketAddr) -> Vec<u8> {
113
     pub fn encode(addr: SocketAddr) -> Vec<u8> {
100
         // not work with [:1]:<port>
114
         // not work with [:1]:<port>
101
-        let addr = match addr {
102
-            SocketAddr::V6(v6) => {
103
-                if let Some(v4) = v6.ip().to_ipv4() {
104
-                    SocketAddr::new(IpAddr::V4(v4), addr.port())
105
-                } else {
106
-                    addr
107
-                }
108
-            }
109
-            _ => addr,
110
-        };
115
+        let addr = try_into_v4(addr);
111
         match addr {
116
         match addr {
112
             SocketAddr::V4(addr_v4) => {
117
             SocketAddr::V4(addr_v4) => {
113
                 let tm = (SystemTime::now()
118
                 let tm = (SystemTime::now()
@@ -353,5 +358,7 @@ mod test {
353
         assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v4)), addr_v4);
358
         assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v4)), addr_v4);
354
         let addr_v6 = "[ef::fe]:8080".parse().unwrap();
359
         let addr_v6 = "[ef::fe]:8080".parse().unwrap();
355
         assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v6)), addr_v6);
360
         assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v6)), addr_v6);
361
+        let addr_v6 = "[::1]:8080".parse().unwrap();
362
+        assert_eq!(AddrMangle::decode(&AddrMangle::encode(addr_v6)), addr_v6);
356
     }
363
     }
357
 }
364
 }

+ 7 - 5
src/rendezvous_server.rs

@@ -4,6 +4,7 @@ use hbb_common::{
4
     allow_err,
4
     allow_err,
5
     bytes::{Bytes, BytesMut},
5
     bytes::{Bytes, BytesMut},
6
     bytes_codec::BytesCodec,
6
     bytes_codec::BytesCodec,
7
+    config,
7
     futures::future::join_all,
8
     futures::future::join_all,
8
     futures_util::{
9
     futures_util::{
9
         sink::SinkExt,
10
         sink::SinkExt,
@@ -25,6 +26,7 @@ use hbb_common::{
25
         time::{interval, Duration},
26
         time::{interval, Duration},
26
     },
27
     },
27
     tokio_util::codec::Framed,
28
     tokio_util::codec::Framed,
29
+    try_into_v4,
28
     udp::FramedSocket,
30
     udp::FramedSocket,
29
     AddrMangle, ResultType,
31
     AddrMangle, ResultType,
30
 };
32
 };
@@ -475,7 +477,7 @@ impl RendezvousServer {
475
                 Some(rendezvous_message::Union::PunchHoleRequest(ph)) => {
477
                 Some(rendezvous_message::Union::PunchHoleRequest(ph)) => {
476
                     // there maybe several attempt, so sink can be none
478
                     // there maybe several attempt, so sink can be none
477
                     if let Some(sink) = sink.take() {
479
                     if let Some(sink) = sink.take() {
478
-                        self.tcp_punch.lock().await.insert(addr, sink);
480
+                        self.tcp_punch.lock().await.insert(try_into_v4(addr), sink);
479
                     }
481
                     }
480
                     allow_err!(self.handle_tcp_punch_hole_request(addr, ph, key, ws).await);
482
                     allow_err!(self.handle_tcp_punch_hole_request(addr, ph, key, ws).await);
481
                     return true;
483
                     return true;
@@ -483,7 +485,7 @@ impl RendezvousServer {
483
                 Some(rendezvous_message::Union::RequestRelay(mut rf)) => {
485
                 Some(rendezvous_message::Union::RequestRelay(mut rf)) => {
484
                     // there maybe several attempt, so sink can be none
486
                     // there maybe several attempt, so sink can be none
485
                     if let Some(sink) = sink.take() {
487
                     if let Some(sink) = sink.take() {
486
-                        self.tcp_punch.lock().await.insert(addr, sink);
488
+                        self.tcp_punch.lock().await.insert(try_into_v4(addr), sink);
487
                     }
489
                     }
488
                     if let Some(peer) = self.pm.get_in_memory(&rf.id).await {
490
                     if let Some(peer) = self.pm.get_in_memory(&rf.id).await {
489
                         let mut msg_out = RendezvousMessage::new();
491
                         let mut msg_out = RendezvousMessage::new();
@@ -1048,7 +1050,7 @@ impl RendezvousServer {
1048
 
1050
 
1049
     async fn handle_listener2(&self, stream: TcpStream, addr: SocketAddr) {
1051
     async fn handle_listener2(&self, stream: TcpStream, addr: SocketAddr) {
1050
         let mut rs = self.clone();
1052
         let mut rs = self.clone();
1051
-        if addr.ip().to_string() == "127.0.0.1" {
1053
+        if addr.ip().is_loopback() {
1052
             tokio::spawn(async move {
1054
             tokio::spawn(async move {
1053
                 let mut stream = stream;
1055
                 let mut stream = stream;
1054
                 let mut buffer = [0; 64];
1056
                 let mut buffer = [0; 64];
@@ -1203,7 +1205,7 @@ async fn check_relay_servers(rs0: Arc<RelayServers>, tx: Sender) {
1203
     for x in rs0.iter() {
1205
     for x in rs0.iter() {
1204
         let mut host = x.to_owned();
1206
         let mut host = x.to_owned();
1205
         if !host.contains(':') {
1207
         if !host.contains(':') {
1206
-            host = format!("{}:{}", host, hbb_common::config::RELAY_PORT);
1208
+            host = format!("{}:{}", host, config::RELAY_PORT);
1207
         }
1209
         }
1208
         let rs = rs.clone();
1210
         let rs = rs.clone();
1209
         let x = x.clone();
1211
         let x = x.clone();
@@ -1226,7 +1228,7 @@ async fn check_relay_servers(rs0: Arc<RelayServers>, tx: Sender) {
1226
 
1228
 
1227
 // temp solution to solve udp socket failure
1229
 // temp solution to solve udp socket failure
1228
 async fn test_hbbs(addr: SocketAddr) -> ResultType<()> {
1230
 async fn test_hbbs(addr: SocketAddr) -> ResultType<()> {
1229
-    let mut socket = FramedSocket::new("0.0.0.0:0").await?;
1231
+    let mut socket = FramedSocket::new(config::Config::get_any_listen_addr(addr.is_ipv4())).await?;
1230
     let mut msg_out = RendezvousMessage::new();
1232
     let mut msg_out = RendezvousMessage::new();
1231
     msg_out.set_register_peer(RegisterPeer {
1233
     msg_out.set_register_peer(RegisterPeer {
1232
         id: "(:test_hbbs:)".to_owned(),
1234
         id: "(:test_hbbs:)".to_owned(),