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

peer_online_state: serve online state

Signed-off-by: fufesou <shuanglongchen@yeah.net>
fufesou лет назад: 3
Родитель
Сommit
a4940f4634
3 измененных файлов с 62 добавлено и 8 удалено
  1. 1 0
      .gitignore
  2. 12 0
      libs/hbb_common/protos/rendezvous.proto
  3. 49 8
      src/rendezvous_server.rs

+ 1 - 0
.gitignore

@@ -5,3 +5,4 @@ debian-build
5 5
 debian/.debhelper
6 6
 debian/debhelper-build-stamp
7 7
 .DS_Store
8
+.vscode

+ 12 - 0
libs/hbb_common/protos/rendezvous.proto

@@ -148,6 +148,16 @@ message PeerDiscovery {
148 148
   string misc = 7;
149 149
 }
150 150
 
151
+message OnlineRequest {
152
+  string id = 1;
153
+  repeated string peers = 2;
154
+}
155
+
156
+message OnlineResponse {
157
+  repeated string onlines = 1;
158
+  repeated string offlines = 2;
159
+}
160
+
151 161
 message RendezvousMessage {
152 162
   oneof union {
153 163
     RegisterPeer register_peer = 6;
@@ -167,5 +177,7 @@ message RendezvousMessage {
167 177
     TestNatRequest test_nat_request = 20;
168 178
     TestNatResponse test_nat_response = 21;
169 179
     PeerDiscovery peer_discovery = 22;
180
+    OnlineRequest online_request = 23;
181
+    OnlineResponse online_response = 24;
170 182
   }
171 183
 }

+ 49 - 8
src/rendezvous_server.rs

@@ -756,6 +756,40 @@ impl RendezvousServer {
756 756
         }
757 757
     }
758 758
 
759
+    #[inline]
760
+    async fn handle_online_request(
761
+        &mut self,
762
+        stream: &mut FramedStream,
763
+        peers: Vec<String>,
764
+    ) -> ResultType<()> {
765
+        let mut onlines = Vec::new();
766
+        for peer_id in &peers {
767
+            if let Some(peer) = self.pm.get_in_memory(&peer_id).await {
768
+                let (elapsed, _) = {
769
+                    let r = peer.read().await;
770
+                    (r.last_reg_time.elapsed().as_millis() as i32, r.socket_addr)
771
+                };
772
+                if elapsed < REG_TIMEOUT {
773
+                    onlines.push(peer_id.clone());
774
+                }
775
+            }
776
+        }
777
+        let offlines = peers
778
+            .into_iter()
779
+            .filter(|p| onlines.iter().position(|r| r == p).is_none())
780
+            .collect::<Vec<_>>();
781
+
782
+        let mut msg_out = RendezvousMessage::new();
783
+        msg_out.set_online_response(OnlineResponse {
784
+            onlines,
785
+            offlines,
786
+            ..Default::default()
787
+        });
788
+        stream.send(&msg_out).await?;
789
+
790
+        Ok(())
791
+    }
792
+
759 793
     #[inline]
760 794
     async fn send_to_tcp(&mut self, msg: RendezvousMessage, addr: SocketAddr) {
761 795
         let mut tcp = self.tcp_punch.lock().await.remove(&addr);
@@ -1014,8 +1048,8 @@ impl RendezvousServer {
1014 1048
     }
1015 1049
 
1016 1050
     async fn handle_listener2(&self, stream: TcpStream, addr: SocketAddr) {
1051
+        let mut rs = self.clone();
1017 1052
         if addr.ip().to_string() == "127.0.0.1" {
1018
-            let rs = self.clone();
1019 1053
             tokio::spawn(async move {
1020 1054
                 let mut stream = stream;
1021 1055
                 let mut buffer = [0; 64];
@@ -1033,13 +1067,20 @@ impl RendezvousServer {
1033 1067
             let mut stream = stream;
1034 1068
             if let Some(Ok(bytes)) = stream.next_timeout(30_000).await {
1035 1069
                 if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
1036
-                    if let Some(rendezvous_message::Union::TestNatRequest(_)) = msg_in.union {
1037
-                        let mut msg_out = RendezvousMessage::new();
1038
-                        msg_out.set_test_nat_response(TestNatResponse {
1039
-                            port: addr.port() as _,
1040
-                            ..Default::default()
1041
-                        });
1042
-                        stream.send(&msg_out).await.ok();
1070
+                    match msg_in.union {
1071
+                        Some(rendezvous_message::Union::TestNatRequest(_)) => {
1072
+                            let mut msg_out = RendezvousMessage::new();
1073
+                            msg_out.set_test_nat_response(TestNatResponse {
1074
+                                port: addr.port() as _,
1075
+                                ..Default::default()
1076
+                            });
1077
+                            stream.send(&msg_out).await.ok();
1078
+                        }
1079
+                        Some(rendezvous_message::Union::OnlineRequest(or)) => {
1080
+                            allow_err!(rs.handle_online_request(&mut stream, or.peers).await);
1081
+                        }
1082
+                        _ => {
1083
+                        }
1043 1084
                     }
1044 1085
                 }
1045 1086
             }