open-trade 5 years ago
parent
commit
1acec93fdb
2 changed files with 24 additions and 1 deletions
  1. 3 0
      src/main.rs
  2. 21 1
      src/rendezvous_server.rs

+ 3 - 0
src/main.rs

@@ -57,11 +57,14 @@ async fn main() -> ResultType<()> {
57 57
         .collect();
58 58
     let addr = format!("0.0.0.0:{}", port);
59 59
     log::info!("Listening on {}", addr);
60
+    let addr2 = format!("0.0.0.0:{}", port.parse::<i32>().unwrap_or(0) - 1);
61
+    log::info!("Listening on {}, extra port for NAT test", addr2);
60 62
     log::info!("relay-server={}", relay_server);
61 63
     log::info!("serial={}", serial);
62 64
     log::info!("rendezvous-servers={:?}", rendezvous_servers);
63 65
     RendezvousServer::start(
64 66
         &addr,
67
+        &addr2,
65 68
         relay_server,
66 69
         serial,
67 70
         rendezvous_servers,

+ 21 - 1
src/rendezvous_server.rs

@@ -9,7 +9,7 @@ use hbb_common::{
9 9
     log,
10 10
     protobuf::Message as _,
11 11
     rendezvous_proto::*,
12
-    tcp::new_listener,
12
+    tcp::{new_listener, FramedStream},
13 13
     timeout,
14 14
     tokio::{self, net::TcpStream, sync::mpsc},
15 15
     tokio_util::codec::Framed,
@@ -135,6 +135,7 @@ pub struct RendezvousServer {
135 135
 impl RendezvousServer {
136 136
     pub async fn start(
137 137
         addr: &str,
138
+        addr2: &str,
138 139
         relay_server: String,
139 140
         serial: i32,
140 141
         rendezvous_servers: Vec<String>,
@@ -157,6 +158,7 @@ impl RendezvousServer {
157 158
             software_url,
158 159
         };
159 160
         let mut listener = new_listener(addr, false).await?;
161
+        let mut listener2 = new_listener(addr2, false).await?;
160 162
         loop {
161 163
             tokio::select! {
162 164
                 Some((msg, addr)) = rx.recv() => {
@@ -165,6 +167,24 @@ impl RendezvousServer {
165 167
                 Some(Ok((bytes, addr))) = socket.next() => {
166 168
                     allow_err!(rs.handle_msg(&bytes, addr, &mut socket).await);
167 169
                 }
170
+                Ok((stream, addr)) = listener2.accept() => {
171
+                    let stream = FramedStream::from(stream);
172
+                    tokio::spawn(async move {
173
+                        let mut stream = stream;
174
+                        if let Some(Ok(bytes)) = stream.next_timeout(30_000).await {
175
+                            if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
176
+                                if let Some(rendezvous_message::Union::test_nat_request(_)) = msg_in.union {
177
+                                    let mut msg_out = RendezvousMessage::new();
178
+                                    msg_out.set_test_nat_response(TestNatResponse {
179
+                                        port: addr.port() as _,
180
+                                        ..Default::default()
181
+                                    });
182
+                                    stream.send(&msg_out).await.ok();
183
+                                }
184
+                            }
185
+                        }
186
+                    });
187
+                }
168 188
                 Ok((stream, addr)) = listener.accept() => {
169 189
                     log::debug!("Tcp connection from {:?}", addr);
170 190
                     let (a, mut b) = Framed::new(stream, BytesCodec::new()).split();