From 2f6d743931417cbb415ec44ab2a37089933b8edf Mon Sep 17 00:00:00 2001 From: lubeilin <1791778603@qq.com> Date: Thu, 29 Feb 2024 22:19:41 +0800 Subject: [PATCH] =?UTF-8?q?[mio]=20=E5=8A=A0=E5=85=A5=E7=AB=AF=E5=8F=A3?= =?UTF-8?q?=E7=BB=84=EF=BC=8C=E5=8E=BB=E9=99=A4tokio?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- vnt/src/nat/mod.rs | 77 +++++++++++-------------- vnt/src/nat/stun.rs | 137 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+), 42 deletions(-) create mode 100644 vnt/src/nat/stun.rs diff --git a/vnt/src/nat/mod.rs b/vnt/src/nat/mod.rs index 223440d..eeba527 100644 --- a/vnt/src/nat/mod.rs +++ b/vnt/src/nat/mod.rs @@ -1,6 +1,7 @@ use std::io; use std::net::UdpSocket; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; +use std::ops::Sub; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -10,7 +11,7 @@ use parking_lot::Mutex; use crate::channel::punch::{NatInfo, NatType}; use crate::proto::message::PunchNatType; -mod stun_test; +mod stun; pub fn local_ipv4_() -> io::Result { let socket = UdpSocket::bind("0.0.0.0:0")?; @@ -78,22 +79,20 @@ impl Into for PunchNatType { impl NatTest { pub fn new( mut stun_server: Vec, - public_ip: Ipv4Addr, - public_port: u16, local_ipv4: Option, ipv6: Option, - udp_port: u16, + udp_ports: Vec, tcp_port: u16, ) -> NatTest { let server = stun_server[0].clone(); stun_server.resize(3, server); let nat_info = NatInfo::new( - vec![public_ip], - public_port, + Vec::new(), + Vec::new(), 0, local_ipv4, ipv6, - udp_port, + udp_ports, tcp_port, NatType::Cone, ); @@ -101,7 +100,9 @@ impl NatTest { NatTest { stun_server, info, - time: Arc::new(AtomicCell::new(Instant::now())), + time: Arc::new(AtomicCell::new( + Instant::now().sub(Duration::from_secs(100)), + )), } } pub fn can_update(&self) -> bool { @@ -109,70 +110,62 @@ impl NatTest { last.elapsed() > Duration::from_secs(10) && self.time.compare_exchange(last, Instant::now()).is_ok() } + pub fn nat_info(&self) -> NatInfo { self.info.lock().clone() } - pub fn update_addr(&self, ip: Ipv4Addr, port: u16) { + pub fn update_addr(&self, index: usize, ip: Ipv4Addr, port: u16) { let mut guard = self.info.lock(); - guard.update_addr(ip, port) + guard.update_addr(index, ip, port) } - pub async fn re_test( + pub fn re_test( &self, - public_ip: Ipv4Addr, - public_port: u16, + public_ports: Vec, local_ipv4: Option, ipv6: Option, - udp_port: u16, + udp_ports: Vec, tcp_port: u16, ) -> NatInfo { let info = NatTest::re_test_( &self.stun_server, - public_ip, - public_port, + public_ports, local_ipv4, ipv6, - udp_port, + udp_ports, tcp_port, - ) - .await; + ); log::info!("探测nat类型={:?}", info); *self.info.lock() = info.clone(); info } - async fn re_test_( + fn re_test_( stun_server: &Vec, - public_ip: Ipv4Addr, - public_port: u16, + public_ports: Vec, local_ipv4: Option, ipv6: Option, - udp_port: u16, + udp_ports: Vec, tcp_port: u16, ) -> NatInfo { - return match stun_test::stun_test_nat(stun_server.clone()).await { - Ok((nat_type, mut public_ips, port_range)) => { - if !public_ips.contains(&public_ip) { - public_ips.push(public_ip) - } - NatInfo::new( - public_ips, - public_port, - port_range, - local_ipv4, - ipv6, - udp_port, - tcp_port, - nat_type, - ) - } + return match stun::stun_test_nat(stun_server.clone()) { + Ok((nat_type, public_ips, port_range)) => NatInfo::new( + public_ips, + public_ports, + port_range, + local_ipv4, + ipv6, + udp_ports, + tcp_port, + nat_type, + ), Err(e) => { log::warn!("{:?}", e); NatInfo::new( - vec![public_ip], - public_port, + Vec::new(), + public_ports, 0, local_ipv4, ipv6, - udp_port, + udp_ports, tcp_port, NatType::Cone, ) diff --git a/vnt/src/nat/stun.rs b/vnt/src/nat/stun.rs new file mode 100644 index 0000000..b1351ff --- /dev/null +++ b/vnt/src/nat/stun.rs @@ -0,0 +1,137 @@ +use std::collections::HashSet; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; +use std::time::Duration; +use std::{io, thread}; + +use crate::channel::punch::NatType; +use std::net::UdpSocket; +use stun_format::Attr; + +pub fn stun_test_nat(stun_servers: Vec) -> io::Result<(NatType, Vec, u16)> { + let mut h = Vec::new(); + for x in stun_servers { + let handle = thread::spawn(move || test_nat(x)); + h.push(handle); + } + let mut nat_type = NatType::Cone; + let mut port_range = 0; + let mut hash_set = HashSet::new(); + for x in h { + if let Ok(rs) = x.join() { + if let Ok((nat_type_t, ip_list_t, port_range_t)) = rs { + if nat_type_t == NatType::Symmetric { + nat_type = NatType::Symmetric; + } + for x in ip_list_t { + hash_set.insert(x); + } + if port_range < port_range_t { + port_range = port_range_t; + } + } + } + } + Ok((nat_type, hash_set.into_iter().collect(), port_range)) +} + +fn test_nat(stun_server: String) -> io::Result<(NatType, Vec, u16)> { + let udp = UdpSocket::bind("0.0.0.0:0")?; + udp.set_read_timeout(Some(Duration::from_millis(300)))?; + udp.connect(stun_server)?; + let mut port_range = 0; + let mut hash_set = HashSet::new(); + let mut nat_type = NatType::Cone; + match test_nat_(&udp, true, true) { + Ok((mapped_addr1, changed_addr1)) => { + match mapped_addr1.ip() { + IpAddr::V4(ip) => { + hash_set.insert(ip); + } + IpAddr::V6(_) => {} + } + if udp.connect(changed_addr1).is_ok() { + if let Ok((mapped_addr2, _)) = test_nat_(&udp, false, false) { + match mapped_addr2.ip() { + IpAddr::V4(ip) => { + hash_set.insert(ip); + if mapped_addr1 != mapped_addr2 { + nat_type = NatType::Symmetric; + } + } + IpAddr::V6(_) => {} + } + port_range = mapped_addr2.port().abs_diff(mapped_addr1.port()); + } + } + } + Err(_) => {} + } + Ok((nat_type, hash_set.into_iter().collect(), port_range)) +} + +fn test_nat_( + udp: &UdpSocket, + change_ip: bool, + change_port: bool, +) -> io::Result<(SocketAddr, SocketAddr)> { + for _ in 0..2 { + let mut buf = [0u8; 28]; + let mut msg = stun_format::MsgBuilder::from(buf.as_mut_slice()); + msg.typ(stun_format::MsgType::BindingRequest).unwrap(); + msg.tid(1).unwrap(); + msg.add_attr(Attr::ChangeRequest { + change_ip, + change_port, + }) + .unwrap(); + udp.send(msg.as_bytes())?; + let mut buf = [0; 10240]; + let (len, _addr) = match udp.recv_from(&mut buf) { + Ok(rs) => rs, + Err(_) => { + continue; + } + }; + let msg = stun_format::Msg::from(&buf[..len]); + let mut mapped_addr = None; + let mut changed_addr = None; + for x in msg.attrs_iter() { + match x { + Attr::MappedAddress(addr) => { + if mapped_addr.is_none() { + let _ = mapped_addr.insert(stun_addr(addr)); + } + } + Attr::ChangedAddress(addr) => { + if changed_addr.is_none() { + let _ = changed_addr.insert(stun_addr(addr)); + } + } + Attr::XorMappedAddress(addr) => { + if mapped_addr.is_none() { + let _ = mapped_addr.insert(stun_addr(addr)); + } + } + _ => {} + } + if changed_addr.is_some() && mapped_addr.is_some() { + return Ok((mapped_addr.unwrap(), changed_addr.unwrap())); + } + } + if let Some(addr) = mapped_addr { + return Ok((addr, changed_addr.unwrap_or(addr))); + } + } + Err(io::Error::new(io::ErrorKind::Other, "stun response err")) +} + +fn stun_addr(addr: stun_format::SocketAddr) -> SocketAddr { + match addr { + stun_format::SocketAddr::V4(ip, port) => { + SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from(ip), port)) + } + stun_format::SocketAddr::V6(ip, port) => { + SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::from(ip), port, 0, 0)) + } + } +}