From 1acec93fdbd43d54a6faa20e2b4e5ec761eb2d1c Mon Sep 17 00:00:00 2001
From: open-trade <info@opentradesolutions.com>
Date: Sun, 27 Sep 2020 22:20:42 +0800
Subject: [PATCH] port2

---
 src/main.rs              |  3 +++
 src/rendezvous_server.rs | 22 +++++++++++++++++++++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/main.rs b/src/main.rs
index 7982a1c..1c81632 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -57,11 +57,14 @@ async fn main() -> ResultType<()> {
         .collect();
     let addr = format!("0.0.0.0:{}", port);
     log::info!("Listening on {}", addr);
+    let addr2 = format!("0.0.0.0:{}", port.parse::<i32>().unwrap_or(0) - 1);
+    log::info!("Listening on {}, extra port for NAT test", addr2);
     log::info!("relay-server={}", relay_server);
     log::info!("serial={}", serial);
     log::info!("rendezvous-servers={:?}", rendezvous_servers);
     RendezvousServer::start(
         &addr,
+        &addr2,
         relay_server,
         serial,
         rendezvous_servers,
diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs
index ad387c1..bce2f80 100644
--- a/src/rendezvous_server.rs
+++ b/src/rendezvous_server.rs
@@ -9,7 +9,7 @@ use hbb_common::{
     log,
     protobuf::Message as _,
     rendezvous_proto::*,
-    tcp::new_listener,
+    tcp::{new_listener, FramedStream},
     timeout,
     tokio::{self, net::TcpStream, sync::mpsc},
     tokio_util::codec::Framed,
@@ -135,6 +135,7 @@ pub struct RendezvousServer {
 impl RendezvousServer {
     pub async fn start(
         addr: &str,
+        addr2: &str,
         relay_server: String,
         serial: i32,
         rendezvous_servers: Vec<String>,
@@ -157,6 +158,7 @@ impl RendezvousServer {
             software_url,
         };
         let mut listener = new_listener(addr, false).await?;
+        let mut listener2 = new_listener(addr2, false).await?;
         loop {
             tokio::select! {
                 Some((msg, addr)) = rx.recv() => {
@@ -165,6 +167,24 @@ impl RendezvousServer {
                 Some(Ok((bytes, addr))) = socket.next() => {
                     allow_err!(rs.handle_msg(&bytes, addr, &mut socket).await);
                 }
+                Ok((stream, addr)) = listener2.accept() => {
+                    let stream = FramedStream::from(stream);
+                    tokio::spawn(async move {
+                        let mut stream = stream;
+                        if let Some(Ok(bytes)) = stream.next_timeout(30_000).await {
+                            if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
+                                if let Some(rendezvous_message::Union::test_nat_request(_)) = msg_in.union {
+                                    let mut msg_out = RendezvousMessage::new();
+                                    msg_out.set_test_nat_response(TestNatResponse {
+                                        port: addr.port() as _,
+                                        ..Default::default()
+                                    });
+                                    stream.send(&msg_out).await.ok();
+                                }
+                            }
+                        }
+                    });
+                }
                 Ok((stream, addr)) = listener.accept() => {
                     log::debug!("Tcp connection from {:?}", addr);
                     let (a, mut b) = Framed::new(stream, BytesCodec::new()).split();