Просмотр исходного кода

refacor relay_response to support initiate relay from server side

open-trade лет назад: 5
Родитель
Сommit
77b1072c7d
2 измененных файлов с 26 добавлено и 25 удалено
  1. 1 1
      libs/hbb_common
  2. 25 24
      src/rendezvous_server.rs

+ 1 - 1
libs/hbb_common

@@ -1 +1 @@
1
-Subproject commit ef0bf8e43ad9017d5b7851ad0f261b6503554e65
1
+Subproject commit 661e7303cc4ce14cc2a8e812200e2464a4da1c32

+ 25 - 24
src/rendezvous_server.rs

@@ -181,7 +181,7 @@ impl RendezvousServer {
181
                                         } else {
181
                                         } else {
182
                                             break;
182
                                             break;
183
                                         }
183
                                         }
184
-                                        allow_err!(rs.handle_tcp_punch_hole_request(addr, ph.id).await);
184
+                                        allow_err!(rs.handle_tcp_punch_hole_request(addr, ph).await);
185
                                     }
185
                                     }
186
                                     Some(rendezvous_message::Union::request_relay(mut rf)) => {
186
                                     Some(rendezvous_message::Union::request_relay(mut rf)) => {
187
                                         if let Some(sender) = sender.take() {
187
                                         if let Some(sender) = sender.take() {
@@ -196,11 +196,17 @@ impl RendezvousServer {
196
                                             rs.tx.send((msg_out, peer.socket_addr)).ok();
196
                                             rs.tx.send((msg_out, peer.socket_addr)).ok();
197
                                         }
197
                                         }
198
                                     }
198
                                     }
199
-                                    Some(rendezvous_message::Union::request_relay_response(mut rfr)) => {
200
-                                        let addr_b = AddrMangle::decode(&rfr.socket_addr);
201
-                                        rfr.socket_addr = Default::default();
199
+                                    Some(rendezvous_message::Union::relay_response(mut rr)) => {
200
+                                        let addr_b = AddrMangle::decode(&rr.socket_addr);
201
+                                        rr.socket_addr = Default::default();
202
+                                        let id = rr.get_id();
203
+                                        if !id.is_empty() {
204
+                                            if let Some(peer) = rs.pm.get(&id).await {
205
+                                                rr.set_pk(peer.pk.clone());
206
+                                            }
207
+                                        }
202
                                         let mut msg_out = RendezvousMessage::new();
208
                                         let mut msg_out = RendezvousMessage::new();
203
-                                        msg_out.set_request_relay_response(rfr);
209
+                                        msg_out.set_relay_response(rr);
204
                                         allow_err!(rs.send_to_tcp_sync(&msg_out, addr_b).await);
210
                                         allow_err!(rs.send_to_tcp_sync(&msg_out, addr_b).await);
205
                                         break;
211
                                         break;
206
                                     }
212
                                     }
@@ -297,14 +303,13 @@ impl RendezvousServer {
297
                     socket.send(&msg_out, addr).await?
303
                     socket.send(&msg_out, addr).await?
298
                 }
304
                 }
299
                 Some(rendezvous_message::Union::punch_hole_request(ph)) => {
305
                 Some(rendezvous_message::Union::punch_hole_request(ph)) => {
300
-                    let id = ph.id;
301
-                    if self.pm.is_in_memory(&id) {
302
-                        self.handle_udp_punch_hole_request(addr, id).await?;
306
+                    if self.pm.is_in_memory(&ph.id) {
307
+                        self.handle_udp_punch_hole_request(addr, ph).await?;
303
                     } else {
308
                     } else {
304
                         // not in memory, fetch from db with spawn in case blocking me
309
                         // not in memory, fetch from db with spawn in case blocking me
305
                         let mut me = self.clone();
310
                         let mut me = self.clone();
306
                         tokio::spawn(async move {
311
                         tokio::spawn(async move {
307
-                            allow_err!(me.handle_udp_punch_hole_request(addr, id).await);
312
+                            allow_err!(me.handle_udp_punch_hole_request(addr, ph).await);
308
                         });
313
                         });
309
                     }
314
                     }
310
                 }
315
                 }
@@ -420,14 +425,10 @@ impl RendezvousServer {
420
             Some(peer) => peer.pk,
425
             Some(peer) => peer.pk,
421
             _ => Vec::new(),
426
             _ => Vec::new(),
422
         };
427
         };
423
-        let mut relay_server = phs.relay_server;
424
-        if relay_server.is_empty() {
425
-            relay_server = self.relay_server.clone();
426
-        }
427
         let mut p = PunchHoleResponse {
428
         let mut p = PunchHoleResponse {
428
             socket_addr: AddrMangle::encode(addr),
429
             socket_addr: AddrMangle::encode(addr),
429
             pk,
430
             pk,
430
-            relay_server,
431
+            relay_server: phs.relay_server.clone(),
431
             ..Default::default()
432
             ..Default::default()
432
         };
433
         };
433
         if let Ok(t) = phs.nat_type.enum_value() {
434
         if let Ok(t) = phs.nat_type.enum_value() {
@@ -458,13 +459,9 @@ impl RendezvousServer {
458
             &addr
459
             &addr
459
         );
460
         );
460
         let mut msg_out = RendezvousMessage::new();
461
         let mut msg_out = RendezvousMessage::new();
461
-        let mut relay_server = la.relay_server;
462
-        if relay_server.is_empty() {
463
-            relay_server = self.relay_server.clone();
464
-        }
465
         let mut p = PunchHoleResponse {
462
         let mut p = PunchHoleResponse {
466
             socket_addr: la.local_addr.clone(),
463
             socket_addr: la.local_addr.clone(),
467
-            relay_server,
464
+            relay_server: la.relay_server,
468
             ..Default::default()
465
             ..Default::default()
469
         };
466
         };
470
         p.set_is_local(true);
467
         p.set_is_local(true);
@@ -481,8 +478,9 @@ impl RendezvousServer {
481
     async fn handle_punch_hole_request(
478
     async fn handle_punch_hole_request(
482
         &mut self,
479
         &mut self,
483
         addr: SocketAddr,
480
         addr: SocketAddr,
484
-        id: String,
481
+        ph: PunchHoleRequest,
485
     ) -> ResultType<(RendezvousMessage, Option<SocketAddr>)> {
482
     ) -> ResultType<(RendezvousMessage, Option<SocketAddr>)> {
483
+        let id = ph.id;
486
         // punch hole request from A, relay to B,
484
         // punch hole request from A, relay to B,
487
         // check if in same intranet first,
485
         // check if in same intranet first,
488
         // fetch local addrs if in same intranet.
486
         // fetch local addrs if in same intranet.
@@ -518,6 +516,7 @@ impl RendezvousServer {
518
                 );
516
                 );
519
                 msg_out.set_fetch_local_addr(FetchLocalAddr {
517
                 msg_out.set_fetch_local_addr(FetchLocalAddr {
520
                     socket_addr,
518
                     socket_addr,
519
+                    relay_server: self.relay_server.clone(),
521
                     ..Default::default()
520
                     ..Default::default()
522
                 });
521
                 });
523
             } else {
522
             } else {
@@ -529,6 +528,8 @@ impl RendezvousServer {
529
                 );
528
                 );
530
                 msg_out.set_punch_hole(PunchHole {
529
                 msg_out.set_punch_hole(PunchHole {
531
                     socket_addr,
530
                     socket_addr,
531
+                    nat_type: ph.nat_type,
532
+                    relay_server: self.relay_server.clone(),
532
                     ..Default::default()
533
                     ..Default::default()
533
                 });
534
                 });
534
             }
535
             }
@@ -574,9 +575,9 @@ impl RendezvousServer {
574
     async fn handle_tcp_punch_hole_request(
575
     async fn handle_tcp_punch_hole_request(
575
         &mut self,
576
         &mut self,
576
         addr: SocketAddr,
577
         addr: SocketAddr,
577
-        id: String,
578
+        ph: PunchHoleRequest,
578
     ) -> ResultType<()> {
579
     ) -> ResultType<()> {
579
-        let (msg, to_addr) = self.handle_punch_hole_request(addr, id).await?;
580
+        let (msg, to_addr) = self.handle_punch_hole_request(addr, ph).await?;
580
         if let Some(addr) = to_addr {
581
         if let Some(addr) = to_addr {
581
             self.tx.send((msg, addr))?;
582
             self.tx.send((msg, addr))?;
582
         } else {
583
         } else {
@@ -589,9 +590,9 @@ impl RendezvousServer {
589
     async fn handle_udp_punch_hole_request(
590
     async fn handle_udp_punch_hole_request(
590
         &mut self,
591
         &mut self,
591
         addr: SocketAddr,
592
         addr: SocketAddr,
592
-        id: String,
593
+        ph: PunchHoleRequest,
593
     ) -> ResultType<()> {
594
     ) -> ResultType<()> {
594
-        let (msg, to_addr) = self.handle_punch_hole_request(addr, id).await?;
595
+        let (msg, to_addr) = self.handle_punch_hole_request(addr, ph).await?;
595
         self.tx.send((
596
         self.tx.send((
596
             msg,
597
             msg,
597
             match to_addr {
598
             match to_addr {