yeyinlaoying/vnt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

修复tcp模式下的p2p问题

Browse Source
This commit is contained in:
lubeilin
2023-07-24 21:37:19 +08:00
parent 9b42c5d092
commit aa6d3a6843
7 changed files with 126 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vnt/src/channel/channel.rs
+3
View File
@@ -87,6 +87,9 @@ impl Context {
pub fn main_local_port(&self) -> io::Result<u16> {
self.inner.main_channel.local_addr().map(|k| k.port())
}
pub async fn send_main_udp(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
self.inner.main_channel.send_to(buf, addr).await
}
pub async fn send_main(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
if let Some(sender) = &self.inner.main_tcp_channel {
let mut vec = vec![0; 4 + buf.len()];
vnt/src/channel/punch.rs
+7 -4
View File
@@ -24,12 +24,15 @@ pub enum NatType {
}
impl NatInfo {
pub fn new(public_ips: Vec<Ipv4Addr>,
pub fn new(mut public_ips: Vec<Ipv4Addr>,
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;
}
}
vnt/src/core/mod.rs
+11 -6
View File
@@ -60,20 +60,25 @@ pub struct VntUtil {
impl VntUtil {
pub async fn new(config: Config) -> io::Result<VntUtil> {
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<RegResponse, ReqEnum> {
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 {
vnt/src/handle/heartbeat_handler.rs
+15 -3
View File
@@ -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);
}
vnt/src/handle/recv_handler.rs
+30
View File
@@ -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(())
}
vnt/src/nat/mod.rs
+8
View File
@@ -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,
vnt/src/protocol/control_packet.rs
+52
View File
@@ -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<u8> 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<u8> for Protocol {
Protocol::Pong => 2,
Protocol::PunchRequest => 3,
Protocol::PunchResponse => 4,
Protocol::AddrRequest => 5,
Protocol::AddrResponse => 6,
Protocol::Unknown(val) => val,
}
}
}
pub enum ControlPacket<B> {
PingPacket(PingPacket<B>),
PongPacket(PongPacket<B>),
PunchRequest,
PunchResponse,
AddrRequest,
AddrResponse(AddrPacket<B>),
}
impl<B: AsRef<[u8]>> ControlPacket<B> {
@@ -57,6 +68,8 @@ impl<B: AsRef<[u8]>> ControlPacket<B> {
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<B: AsRef<[u8]>> fmt::Debug for PingPacket<B> {
.finish()
}
}
pub struct AddrPacket<B> {
buffer: B,
}
impl<B: AsRef<[u8]>> AddrPacket<B> {
pub fn new(buffer: B) -> io::Result<AddrPacket<B>> {
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<B: AsRef<[u8]> + AsMut<[u8]>> AddrPacket<B> {
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<B: AsRef<[u8]>> fmt::Debug for AddrPacket<B> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("AddrPacket")
.field("ipv4", &self.ipv4())
.field("port", &self.port())
.finish()
}
}