From 94e83d80a689cb1d34a50be709b6417c2f59bda5 Mon Sep 17 00:00:00 2001
From: lubeilin <1791778603@qq.com>
Date: Mon, 13 Feb 2023 22:49:01 +0800
Subject: [PATCH] =?UTF-8?q?=E5=88=9D=E5=A7=8B=E5=8C=96=E6=A8=A1=E5=9D=97?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 switch/Cargo.toml                         |  1 +
 switch/nat_traversal/Cargo.toml           | 11 +++
 switch/nat_traversal/src/cone/mod.rs      | 39 +++++++++++
 switch/nat_traversal/src/lib.rs           |  5 ++
 switch/nat_traversal/src/symmetric/mod.rs | 84 +++++++++++++++++++++++
 5 files changed, 140 insertions(+)
 create mode 100644 switch/nat_traversal/Cargo.toml
 create mode 100644 switch/nat_traversal/src/cone/mod.rs
 create mode 100644 switch/nat_traversal/src/lib.rs
 create mode 100644 switch/nat_traversal/src/symmetric/mod.rs

diff --git a/switch/Cargo.toml b/switch/Cargo.toml
index d399f73..27c9341 100644
--- a/switch/Cargo.toml
+++ b/switch/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021"
 
 [dependencies]
 packet = { path = "./packet" }
+nat_traversal = { path = "./nat_traversal" }
 bytes = "1.3.0"
 log = "0.4.17"
 libc = "0.2.137"
diff --git a/switch/nat_traversal/Cargo.toml b/switch/nat_traversal/Cargo.toml
new file mode 100644
index 0000000..776b151
--- /dev/null
+++ b/switch/nat_traversal/Cargo.toml
@@ -0,0 +1,11 @@
+[package]
+name = "nat_traversal"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+tokio = { version = "1.24.1", features = ["net"] }
+futures = "0.3"
+crossbeam-skiplist = "0.1"
diff --git a/switch/nat_traversal/src/cone/mod.rs b/switch/nat_traversal/src/cone/mod.rs
new file mode 100644
index 0000000..9f2bd8e
--- /dev/null
+++ b/switch/nat_traversal/src/cone/mod.rs
@@ -0,0 +1,39 @@
+use std::io;
+use std::net::SocketAddr;
+
+use tokio::net::UdpSocket;
+
+/// 锥形网络，使用一个端口
+pub struct Channel {
+    udp: UdpSocket,
+    server_address: SocketAddr,
+}
+
+impl Channel {
+    pub async fn new(server_address: SocketAddr) -> io::Result<Self> {
+        Ok(Self {
+            udp: UdpSocket::bind("0:0").await?,
+            server_address,
+        })
+    }
+}
+
+impl Channel {
+    #[inline]
+    pub async fn recv(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
+        self.udp.recv_from(buf).await
+    }
+}
+
+impl Channel {
+    #[inline]
+    pub async fn send_to(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
+        self.udp.send_to(buf, addr).await
+    }
+    #[inline]
+    pub async fn send_server(&self, buf: &[u8]) -> io::Result<usize> {
+        self.udp.send_to(buf, self.server_address).await
+    }
+}
+
+
diff --git a/switch/nat_traversal/src/lib.rs b/switch/nat_traversal/src/lib.rs
new file mode 100644
index 0000000..9b7699b
--- /dev/null
+++ b/switch/nat_traversal/src/lib.rs
@@ -0,0 +1,5 @@
+pub use cone::Channel as ConeChannel;
+pub use symmetric::Channel as SymmetricChannel;
+
+mod cone;
+mod symmetric;
\ No newline at end of file
diff --git a/switch/nat_traversal/src/symmetric/mod.rs b/switch/nat_traversal/src/symmetric/mod.rs
new file mode 100644
index 0000000..d0ba36e
--- /dev/null
+++ b/switch/nat_traversal/src/symmetric/mod.rs
@@ -0,0 +1,84 @@
+use std::io;
+use std::net::SocketAddr;
+use std::sync::Arc;
+
+use crossbeam_skiplist::SkipMap;
+use tokio::net::UdpSocket;
+
+/// 对称网络，绑定多个端口
+///
+///
+/// 假设一方是对称网络，一方是锥形网络
+/// 如果对称网络一方绑定n个端口，通过NAT对外映射出n个 公网ip:公网端口，随机尝试k次的情况下
+/// 猜中的概率 p = 1-((65535-n)/65535)*((65535-n-1)/(65535-1))*...*((65535-n-k+1)/(65535-k+1))
+/// n取76，k取600，猜中的概率就超过50%了
+pub struct Channel {
+    udp_list: Vec<Arc<UdpSocket>>,
+    addr_map: SkipMap<SocketAddr, usize>,
+    server_address: SocketAddr,
+}
+
+impl Channel {
+    pub async fn new(server_address: SocketAddr, num: usize) -> io::Result<Self> {
+        let mut udp_list = Vec::with_capacity(num);
+        for _ in 0..num {
+            udp_list.push(Arc::new(UdpSocket::bind("0:0").await?));
+        }
+        Ok(Self {
+            udp_list,
+            addr_map: SkipMap::new(),
+            server_address,
+        })
+    }
+}
+
+impl Channel {
+    #[inline]
+    pub async fn recv(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
+        let mut list = Vec::with_capacity(self.udp_list.len());
+        for udp in &self.udp_list {
+            let udp = udp.clone();
+            list.push(Box::pin(async move {
+                udp.readable().await
+            }));
+        }
+        let (rs, index, _) = futures::future::select_all(list.into_iter()).await;
+        let _ = rs?;
+        let (len, addr) = self.udp_list[index].try_recv_from(buf)?;
+        self.addr_map.insert(addr, index);
+        Ok((len, addr))
+    }
+}
+
+impl Channel {
+    /// 向一个已经穿透成功洞地址发数据
+    #[inline]
+    pub async fn send_to(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
+        if let Some(entry) = self.addr_map.get(&addr) {
+            self.udp_list[*entry.value()].send_to(buf, addr).await
+        } else {
+            Err(io::Error::from(io::ErrorKind::NotConnected))
+        }
+    }
+    /// 向所有渠道发数据，用于打洞
+    #[inline]
+    pub async fn send_all(&self, buf: &[u8], addr: SocketAddr) -> io::Result<()> {
+        for udp in &self.udp_list {
+            udp.send_to(buf, addr).await?;
+        }
+        Ok(())
+    }
+    /// 向服务器发送数据
+    #[inline]
+    pub async fn send_server(&self, buf: &[u8]) -> io::Result<usize> {
+        self.udp_list[0].send_to(buf, self.server_address).await
+    }
+}
+
+impl Channel {
+    pub fn remove_hole(&self, hole: &SocketAddr) {
+        self.addr_map.remove(hole);
+    }
+}
+
+