diff --git a/vnt/src/channel/channel.rs b/vnt/src/channel/channel.rs index 22d0158..93b15ff 100644 --- a/vnt/src/channel/channel.rs +++ b/vnt/src/channel/channel.rs @@ -87,6 +87,9 @@ impl Context { pub fn main_local_port(&self) -> io::Result { self.inner.main_channel.local_addr().map(|k| k.port()) } + pub async fn send_main_udp(&self, buf: &[u8], addr: SocketAddr) -> io::Result { + self.inner.main_channel.send_to(buf, addr).await + } pub async fn send_main(&self, buf: &[u8], addr: SocketAddr) -> io::Result { if let Some(sender) = &self.inner.main_tcp_channel { let mut vec = vec![0; 4 + buf.len()]; diff --git a/vnt/src/channel/punch.rs b/vnt/src/channel/punch.rs index 3c57d9b..39592f2 100644 --- a/vnt/src/channel/punch.rs +++ b/vnt/src/channel/punch.rs @@ -24,12 +24,15 @@ pub enum NatType { } impl NatInfo { - pub fn new(public_ips: Vec, + pub fn new(mut public_ips: Vec, public_port: u16, public_port_range: u16, local_ip: Ipv4Addr, local_port: u16, nat_type: NatType, ) -> Self { + public_ips.retain(|ip| { + !ip.is_loopback() && !ip.is_private() + }); Self { public_ips, public_port, @@ -68,7 +71,7 @@ impl Punch { return Ok(()); } if !nat_info.local_ip.is_unspecified() || nat_info.local_port != 0 { - let _ = self.context.send_main(buf, SocketAddr::V4(SocketAddrV4::new(nat_info.local_ip, nat_info.local_port))).await; + let _ = self.context.send_main_udp(buf, SocketAddr::V4(SocketAddrV4::new(nat_info.local_ip, nat_info.local_port))).await; } match nat_info.nat_type { NatType::Symmetric => { @@ -122,7 +125,7 @@ impl Punch { for ip in nat_info.public_ips { let addr = SocketAddr::V4(SocketAddrV4::new(ip, nat_info.public_port)); if is_cone { - self.context.send_main(buf, addr).await?; + self.context.send_main_udp(buf, addr).await?; } else { //只有一方是对称，则对称方要使用全部端口发送数据，符合上述计算的概率 self.context.send_all(buf, addr).await?; @@ -143,7 +146,7 @@ impl Punch { return Ok(()); } let addr = SocketAddr::V4(SocketAddrV4::new(*pub_ip, *port)); - self.context.send_main(buf, addr).await?; + self.context.send_main_udp(buf, addr).await?; tokio::time::sleep(Duration::from_millis(2)).await; } } diff --git a/vnt/src/core/mod.rs b/vnt/src/core/mod.rs index ecaeef0..0beba1d 100644 --- a/vnt/src/core/mod.rs +++ b/vnt/src/core/mod.rs @@ -60,20 +60,25 @@ pub struct VntUtil { impl VntUtil { pub async fn new(config: Config) -> io::Result { let main_channel = UdpSocket::bind("0.0.0.0:0").await?; - let main_tcp_channel = if config.tcp { - Some(TcpStream::connect(config.server_address).await?) - } else { - None - }; Ok(VntUtil { config, main_channel, - main_tcp_channel, + main_tcp_channel: None, response: None, iface: None, }) } pub async fn connect(&mut self) -> Result { + if self.config.tcp { + match TcpStream::connect(self.config.server_address).await { + Ok(tcp) => { + let _ = self.main_tcp_channel.insert(tcp); + } + Err(e) => { + return Err(ReqEnum::Other(format!("connect error:{}", e))); + } + } + } match registration_handler::registration(&self.main_channel, self.main_tcp_channel.as_mut(), self.config.server_address, self.config.token.clone(), self.config.device_id.clone(), self.config.name.clone(), self.config.ip.unwrap_or(Ipv4Addr::UNSPECIFIED)).await { diff --git a/vnt/src/handle/heartbeat_handler.rs b/vnt/src/handle/heartbeat_handler.rs index 0b98a2f..e72d8dd 100644 --- a/vnt/src/handle/heartbeat_handler.rs +++ b/vnt/src/handle/heartbeat_handler.rs @@ -14,7 +14,7 @@ use crate::core::status::VntWorker; use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo}; use crate::protocol::control_packet::PingPacket; -use crate::protocol::{control_packet, NetPacket, Protocol, Version}; +use crate::protocol::{control_packet, MAX_TTL, NetPacket, Protocol, Version}; pub fn start_idle(mut worker: VntWorker, idle: Idle, sender: ChannelSender) { tokio::spawn(async move { @@ -91,7 +91,19 @@ async fn start_heartbeat_( return Ok(()); } let mut current_dev = current_device.load(); - if count % 6 == 0 { + if count % 10 == 0 { + let mut packet = NetPacket::new([0; 12])?; + packet.set_version(Version::V1); + packet.set_protocol(Protocol::Control); + packet.set_transport_protocol( + control_packet::Protocol::AddrRequest.into(), + ); + packet.first_set_ttl(MAX_TTL); + packet.set_source(current_dev.virtual_ip()); + packet.set_destination(current_dev.virtual_gateway); + let _ = sender.send_main_udp(packet.buffer(), current_dev.connect_server).await; + } + if count % 20 == 19 { if let Ok(mut addr) = server_address_str.to_socket_addrs() { if let Some(addr) = addr.next() { if addr != current_dev.connect_server { @@ -164,9 +176,9 @@ async fn start_heartbeat_( } } else { for (peer_ip, route_list) in sender.route_table().iter() { - set_now_time(&mut net_packet)?; net_packet.set_destination(*peer_ip); for route in route_list { + set_now_time(&mut net_packet)?; if let Err(e) = sender.send_by_key(net_packet.buffer(), &route.route_key()).await { log::warn!("peer_ip:{:?},route:{:?},e:{:?}", peer_ip, route, e); } diff --git a/vnt/src/handle/recv_handler.rs b/vnt/src/handle/recv_handler.rs index da4c0b4..08deb01 100644 --- a/vnt/src/handle/recv_handler.rs +++ b/vnt/src/handle/recv_handler.rs @@ -416,6 +416,36 @@ impl ChannelDataHandler { let route = Route::from(*route_key, metric, 99); context.add_route_if_absent(source, route); } + ControlPacket::AddrRequest => { + match route_key.addr.ip() { + std::net::IpAddr::V4(ipv4) => { + let mut packet = NetPacket::new([0;12+6])?; + packet.set_version(Version::V1); + packet.set_protocol(Protocol::Control); + packet.set_transport_protocol( + control_packet::Protocol::AddrResponse.into(), + ); + packet.first_set_ttl(MAX_TTL); + packet.set_source(current_device.virtual_ip()); + packet.set_destination(source); + let mut addr_packet = control_packet::AddrPacket::new(packet.payload_mut())?; + addr_packet.set_ipv4(ipv4); + addr_packet.set_port(route_key.addr.port()); + context.send_by_key(packet.buffer(), route_key).await?; + } + std::net::IpAddr::V6(_) => {} + } + } + ControlPacket::AddrResponse(addr_packet) => { + if addr_packet.port() != 0 + && !addr_packet.ipv4().is_multicast() + && !addr_packet.ipv4().is_broadcast() + && !addr_packet.ipv4().is_unspecified() + && !addr_packet.ipv4().is_loopback() + && !addr_packet.ipv4().is_private() { + self.nat_test.update_addr(addr_packet.ipv4(), addr_packet.port()) + } + } } Ok(()) } diff --git a/vnt/src/nat/mod.rs b/vnt/src/nat/mod.rs index ab58053..a4a2b73 100644 --- a/vnt/src/nat/mod.rs +++ b/vnt/src/nat/mod.rs @@ -70,6 +70,14 @@ impl NatTest { pub fn nat_info(&self) -> NatInfo { self.info.lock().clone() } + pub fn update_addr(&self, ip: Ipv4Addr, port: u16) { + let mut guard = self.info.lock(); + guard.public_port = port; + if !guard.public_ips.contains(&ip) { + guard.public_ips.push(ip); + } + println!("{:?}",guard); + } pub fn re_test( &self, public_ip: Ipv4Addr, diff --git a/vnt/src/protocol/control_packet.rs b/vnt/src/protocol/control_packet.rs index 333e3f5..e00fa3e 100644 --- a/vnt/src/protocol/control_packet.rs +++ b/vnt/src/protocol/control_packet.rs @@ -1,4 +1,5 @@ use std::{fmt, io}; +use std::net::Ipv4Addr; #[derive(Eq, PartialEq, Copy, Clone, Debug)] pub enum Protocol { @@ -17,6 +18,9 @@ pub enum Protocol { PunchRequest, /// 打洞响应 PunchResponse, + ///获取对端看到的地址 + AddrRequest, + AddrResponse, Unknown(u8), } @@ -27,6 +31,8 @@ impl From for Protocol { 2 => Protocol::Pong, 3 => Protocol::PunchRequest, 4 => Protocol::PunchResponse, + 5 => Protocol::AddrRequest, + 6 => Protocol::AddrResponse, val => Protocol::Unknown(val), } } @@ -39,15 +45,20 @@ impl Into for Protocol { Protocol::Pong => 2, Protocol::PunchRequest => 3, Protocol::PunchResponse => 4, + Protocol::AddrRequest => 5, + Protocol::AddrResponse => 6, Protocol::Unknown(val) => val, } } } + pub enum ControlPacket { PingPacket(PingPacket), PongPacket(PongPacket), PunchRequest, PunchResponse, + AddrRequest, + AddrResponse(AddrPacket), } impl> ControlPacket { @@ -57,6 +68,8 @@ impl> ControlPacket { Protocol::Pong => Ok(ControlPacket::PongPacket(PongPacket::new(buffer)?)), Protocol::PunchRequest => Ok(ControlPacket::PunchRequest), Protocol::PunchResponse => Ok(ControlPacket::PunchResponse), + Protocol::AddrRequest => Ok(ControlPacket::AddrRequest), + Protocol::AddrResponse => Ok(ControlPacket::AddrResponse(AddrPacket::new(buffer)?)), Protocol::Unknown(_) => Err(io::Error::new(io::ErrorKind::InvalidData, "Unsupported")), } } @@ -105,3 +118,42 @@ impl> fmt::Debug for PingPacket { .finish() } } + +pub struct AddrPacket { + buffer: B, +} + +impl> AddrPacket { + pub fn new(buffer: B) -> io::Result> { + let len = buffer.as_ref().len(); + if len != 6 { + return Err(io::Error::new(io::ErrorKind::InvalidData, "len != 6")); + } + Ok(AddrPacket { buffer }) + } + pub fn ipv4(&self) -> Ipv4Addr { + let buf = self.buffer.as_ref(); + Ipv4Addr::new(buf[0], buf[1], buf[2], buf[3]) + } + pub fn port(&self) -> u16 { + u16::from_be_bytes(self.buffer.as_ref()[4..6].try_into().unwrap()) + } +} + +impl + AsMut<[u8]>> AddrPacket { + pub fn set_ipv4(&mut self, ip: Ipv4Addr) { + self.buffer.as_mut()[..4].copy_from_slice(&ip.octets()) + } + pub fn set_port(&mut self, port: u16) { + self.buffer.as_mut()[4..6].copy_from_slice(&port.to_be_bytes()) + } +} + +impl> fmt::Debug for AddrPacket { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AddrPacket") + .field("ipv4", &self.ipv4()) + .field("port", &self.port()) + .finish() + } +}