mirror of
https://mirror.suhoan.cn/https://github.com/EasyTier/EasyTier.git
synced 2025-12-14 05:37:23 +08:00
@@ -9,11 +9,11 @@ use std::{
|
||||
use futures::{stream::FuturesUnordered, Future, Sink, Stream};
|
||||
use network_interface::NetworkInterfaceConfig as _;
|
||||
use pin_project_lite::pin_project;
|
||||
use tokio::io::{AsyncRead, AsyncWrite};
|
||||
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
||||
|
||||
use bytes::{Buf, Bytes, BytesMut};
|
||||
use bytes::{Buf, BufMut, Bytes, BytesMut};
|
||||
use tokio_stream::StreamExt;
|
||||
use tokio_util::io::{poll_read_buf, poll_write_buf};
|
||||
use tokio_util::io::poll_write_buf;
|
||||
use zerocopy::FromBytes as _;
|
||||
|
||||
use crate::{
|
||||
@@ -149,13 +149,18 @@ where
|
||||
*self_mut.max_packet_size * 64,
|
||||
);
|
||||
|
||||
match ready!(poll_read_buf(
|
||||
self_mut.reader.as_mut(),
|
||||
cx,
|
||||
&mut self_mut.buf
|
||||
)) {
|
||||
Ok(size) => {
|
||||
if size == 0 {
|
||||
let cap = self_mut.buf.capacity() - self_mut.buf.len();
|
||||
let buf = self_mut.buf.chunk_mut().as_mut_ptr();
|
||||
let buf = unsafe { std::slice::from_raw_parts_mut(buf, cap) };
|
||||
let mut buf = ReadBuf::new(buf);
|
||||
|
||||
let ret = ready!(self_mut.reader.as_mut().poll_read(cx, &mut buf));
|
||||
let len = buf.filled().len();
|
||||
unsafe { self_mut.buf.advance_mut(len) };
|
||||
|
||||
match ret {
|
||||
Ok(_) => {
|
||||
if len == 0 {
|
||||
return Poll::Ready(None);
|
||||
}
|
||||
}
|
||||
@@ -173,14 +178,16 @@ pub trait ZCPacketToBytes {
|
||||
|
||||
pub struct TcpZCPacketToBytes;
|
||||
impl ZCPacketToBytes for TcpZCPacketToBytes {
|
||||
fn into_bytes(&self, mut item: ZCPacket) -> Result<Bytes, TunnelError> {
|
||||
fn into_bytes(&self, item: ZCPacket) -> Result<Bytes, TunnelError> {
|
||||
let mut item = item.convert_type(ZCPacketType::TCP);
|
||||
|
||||
let tcp_len = PEER_MANAGER_HEADER_SIZE + item.payload_len();
|
||||
let Some(header) = item.mut_tcp_tunnel_header() else {
|
||||
return Err(TunnelError::InvalidPacket("packet too short".to_string()));
|
||||
};
|
||||
header.len.set(tcp_len.try_into().unwrap());
|
||||
|
||||
Ok(item.into_bytes(ZCPacketType::TCP))
|
||||
Ok(item.into_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -96,23 +96,57 @@ const PAYLOAD_OFFSET_FOR_NIC_PACKET: usize = max(
|
||||
WG_TUNNEL_HEADER_SIZE,
|
||||
) + PEER_MANAGER_HEADER_SIZE;
|
||||
|
||||
const INVALID_OFFSET: usize = usize::MAX;
|
||||
|
||||
const fn get_converted_offset(old_hdr_size: usize, new_hdr_size: usize) -> usize {
|
||||
if old_hdr_size < new_hdr_size {
|
||||
INVALID_OFFSET
|
||||
} else {
|
||||
old_hdr_size - new_hdr_size
|
||||
}
|
||||
}
|
||||
|
||||
impl ZCPacketType {
|
||||
pub fn get_packet_offsets(&self) -> ZCPacketOffsets {
|
||||
match self {
|
||||
ZCPacketType::TCP => ZCPacketOffsets {
|
||||
payload_offset: TCP_TUNNEL_HEADER_SIZE + PEER_MANAGER_HEADER_SIZE,
|
||||
peer_manager_header_offset: TCP_TUNNEL_HEADER_SIZE,
|
||||
..Default::default()
|
||||
tcp_tunnel_header_offset: 0,
|
||||
udp_tunnel_header_offset: get_converted_offset(
|
||||
TCP_TUNNEL_HEADER_SIZE,
|
||||
UDP_TUNNEL_HEADER_SIZE,
|
||||
),
|
||||
wg_tunnel_header_offset: get_converted_offset(
|
||||
TCP_TUNNEL_HEADER_SIZE,
|
||||
WG_TUNNEL_HEADER_SIZE,
|
||||
),
|
||||
},
|
||||
ZCPacketType::UDP => ZCPacketOffsets {
|
||||
payload_offset: UDP_TUNNEL_HEADER_SIZE + PEER_MANAGER_HEADER_SIZE,
|
||||
peer_manager_header_offset: UDP_TUNNEL_HEADER_SIZE,
|
||||
..Default::default()
|
||||
tcp_tunnel_header_offset: get_converted_offset(
|
||||
UDP_TUNNEL_HEADER_SIZE,
|
||||
TCP_TUNNEL_HEADER_SIZE,
|
||||
),
|
||||
udp_tunnel_header_offset: 0,
|
||||
wg_tunnel_header_offset: get_converted_offset(
|
||||
UDP_TUNNEL_HEADER_SIZE,
|
||||
WG_TUNNEL_HEADER_SIZE,
|
||||
),
|
||||
},
|
||||
ZCPacketType::WG => ZCPacketOffsets {
|
||||
payload_offset: WG_TUNNEL_HEADER_SIZE + PEER_MANAGER_HEADER_SIZE,
|
||||
peer_manager_header_offset: WG_TUNNEL_HEADER_SIZE,
|
||||
..Default::default()
|
||||
tcp_tunnel_header_offset: get_converted_offset(
|
||||
WG_TUNNEL_HEADER_SIZE,
|
||||
TCP_TUNNEL_HEADER_SIZE,
|
||||
),
|
||||
udp_tunnel_header_offset: get_converted_offset(
|
||||
WG_TUNNEL_HEADER_SIZE,
|
||||
UDP_TUNNEL_HEADER_SIZE,
|
||||
),
|
||||
wg_tunnel_header_offset: 0,
|
||||
},
|
||||
ZCPacketType::NIC => ZCPacketOffsets {
|
||||
payload_offset: PAYLOAD_OFFSET_FOR_NIC_PACKET,
|
||||
@@ -155,8 +189,10 @@ impl ZCPacket {
|
||||
|
||||
pub fn new_with_payload(payload: &[u8]) -> Self {
|
||||
let mut ret = Self::new_nic_packet();
|
||||
let total_len = ret.packet_type.get_packet_offsets().payload_offset + payload.len();
|
||||
ret.inner.resize(total_len, 0);
|
||||
let payload_off = ret.packet_type.get_packet_offsets().payload_offset;
|
||||
let total_len = payload_off + payload.len();
|
||||
ret.inner.reserve(total_len);
|
||||
unsafe { ret.inner.set_len(total_len) };
|
||||
ret.mut_payload()[..payload.len()].copy_from_slice(&payload);
|
||||
ret
|
||||
}
|
||||
@@ -165,7 +201,7 @@ impl ZCPacket {
|
||||
let mut ret = Self::new_nic_packet();
|
||||
ret.inner.reserve(cap);
|
||||
let total_len = ret.packet_type.get_packet_offsets().payload_offset - packet_info_len;
|
||||
ret.inner.resize(total_len, 0);
|
||||
unsafe { ret.inner.set_len(total_len) };
|
||||
ret
|
||||
}
|
||||
|
||||
@@ -275,45 +311,56 @@ impl ZCPacket {
|
||||
hdr.len.set(payload_len as u32);
|
||||
}
|
||||
|
||||
pub fn into_bytes(mut self, target_packet_type: ZCPacketType) -> Bytes {
|
||||
fn tunnel_payload(&self) -> &[u8] {
|
||||
&self.inner[self
|
||||
.packet_type
|
||||
.get_packet_offsets()
|
||||
.peer_manager_header_offset..]
|
||||
}
|
||||
|
||||
pub fn convert_type(mut self, target_packet_type: ZCPacketType) -> Self {
|
||||
if target_packet_type == self.packet_type {
|
||||
return self.inner.freeze();
|
||||
} else {
|
||||
assert_eq!(
|
||||
self.packet_type,
|
||||
ZCPacketType::NIC,
|
||||
"only support NIC, got {:?}",
|
||||
self
|
||||
);
|
||||
return self;
|
||||
}
|
||||
|
||||
match target_packet_type {
|
||||
ZCPacketType::TCP => self
|
||||
.inner
|
||||
.split_off(
|
||||
self.packet_type
|
||||
.get_packet_offsets()
|
||||
.tcp_tunnel_header_offset,
|
||||
)
|
||||
.freeze(),
|
||||
ZCPacketType::UDP => self
|
||||
.inner
|
||||
.split_off(
|
||||
self.packet_type
|
||||
.get_packet_offsets()
|
||||
.udp_tunnel_header_offset,
|
||||
)
|
||||
.freeze(),
|
||||
ZCPacketType::WG => self
|
||||
.inner
|
||||
.split_off(
|
||||
self.packet_type
|
||||
.get_packet_offsets()
|
||||
.wg_tunnel_header_offset,
|
||||
)
|
||||
.freeze(),
|
||||
let new_offset = match target_packet_type {
|
||||
ZCPacketType::TCP => {
|
||||
self.packet_type
|
||||
.get_packet_offsets()
|
||||
.tcp_tunnel_header_offset
|
||||
}
|
||||
ZCPacketType::UDP => {
|
||||
self.packet_type
|
||||
.get_packet_offsets()
|
||||
.udp_tunnel_header_offset
|
||||
}
|
||||
ZCPacketType::WG => {
|
||||
self.packet_type
|
||||
.get_packet_offsets()
|
||||
.wg_tunnel_header_offset
|
||||
}
|
||||
ZCPacketType::NIC => unreachable!(),
|
||||
};
|
||||
|
||||
tracing::debug!(?self.packet_type, ?target_packet_type, ?new_offset, "convert zc packet type");
|
||||
|
||||
if new_offset == INVALID_OFFSET {
|
||||
// copy peer manager header and payload to new buffer
|
||||
let tunnel_payload = self.tunnel_payload();
|
||||
let new_pm_offset = target_packet_type
|
||||
.get_packet_offsets()
|
||||
.peer_manager_header_offset;
|
||||
let mut buf = BytesMut::with_capacity(new_pm_offset + tunnel_payload.len());
|
||||
unsafe { buf.set_len(new_pm_offset) };
|
||||
buf.extend_from_slice(tunnel_payload);
|
||||
return Self::new_from_buf(buf, target_packet_type);
|
||||
}
|
||||
|
||||
return Self::new_from_buf(self.inner.split_off(new_offset), target_packet_type);
|
||||
}
|
||||
|
||||
pub fn into_bytes(self) -> Bytes {
|
||||
self.inner.freeze()
|
||||
}
|
||||
|
||||
pub fn inner(self) -> BytesMut {
|
||||
@@ -349,7 +396,7 @@ mod tests {
|
||||
assert_eq!(packet.payload_len(), 11);
|
||||
println!("{:?}", packet.inner);
|
||||
|
||||
let tcp_packet = packet.into_bytes(ZCPacketType::TCP);
|
||||
let tcp_packet = packet.convert_type(ZCPacketType::TCP).into_bytes();
|
||||
assert_eq!(&tcp_packet[..1], b"\x0b");
|
||||
println!("{:?}", tcp_packet);
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering::Relaxed};
|
||||
use std::sync::atomic::{AtomicU32, Ordering::Relaxed};
|
||||
|
||||
pub struct WindowLatency {
|
||||
latency_us_window: Vec<AtomicU32>,
|
||||
@@ -48,48 +48,49 @@ impl WindowLatency {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct Throughput {
|
||||
tx_bytes: AtomicU64,
|
||||
rx_bytes: AtomicU64,
|
||||
tx_bytes: u64,
|
||||
rx_bytes: u64,
|
||||
|
||||
tx_packets: AtomicU64,
|
||||
rx_packets: AtomicU64,
|
||||
tx_packets: u64,
|
||||
rx_packets: u64,
|
||||
}
|
||||
|
||||
impl Throughput {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tx_bytes: AtomicU64::new(0),
|
||||
rx_bytes: AtomicU64::new(0),
|
||||
|
||||
tx_packets: AtomicU64::new(0),
|
||||
rx_packets: AtomicU64::new(0),
|
||||
}
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn tx_bytes(&self) -> u64 {
|
||||
self.tx_bytes.load(Relaxed)
|
||||
self.tx_bytes
|
||||
}
|
||||
|
||||
pub fn rx_bytes(&self) -> u64 {
|
||||
self.rx_bytes.load(Relaxed)
|
||||
self.rx_bytes
|
||||
}
|
||||
|
||||
pub fn tx_packets(&self) -> u64 {
|
||||
self.tx_packets.load(Relaxed)
|
||||
self.tx_packets
|
||||
}
|
||||
|
||||
pub fn rx_packets(&self) -> u64 {
|
||||
self.rx_packets.load(Relaxed)
|
||||
self.rx_packets
|
||||
}
|
||||
|
||||
pub fn record_tx_bytes(&self, bytes: u64) {
|
||||
self.tx_bytes.fetch_add(bytes, Relaxed);
|
||||
self.tx_packets.fetch_add(1, Relaxed);
|
||||
#[allow(invalid_reference_casting)]
|
||||
unsafe {
|
||||
*(&self.tx_bytes as *const u64 as *mut u64) += bytes;
|
||||
*(&self.tx_packets as *const u64 as *mut u64) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn record_rx_bytes(&self, bytes: u64) {
|
||||
self.rx_bytes.fetch_add(bytes, Relaxed);
|
||||
self.rx_packets.fetch_add(1, Relaxed);
|
||||
#[allow(invalid_reference_casting)]
|
||||
unsafe {
|
||||
*(&self.rx_bytes as *const u64 as *mut u64) += bytes;
|
||||
*(&self.rx_packets as *const u64 as *mut u64) += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -125,20 +125,21 @@ async fn forward_from_ring_to_udp(
|
||||
let Some(buf) = ring_recv.next().await else {
|
||||
return None;
|
||||
};
|
||||
let mut packet = match buf {
|
||||
let packet = match buf {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
return Some(e);
|
||||
}
|
||||
};
|
||||
|
||||
let mut packet = packet.convert_type(ZCPacketType::UDP);
|
||||
let udp_payload_len = packet.udp_payload().len();
|
||||
let header = packet.mut_udp_tunnel_header().unwrap();
|
||||
header.conn_id.set(conn_id);
|
||||
header.len.set(udp_payload_len as u16);
|
||||
header.msg_type = UdpPacketType::Data as u8;
|
||||
|
||||
let buf = packet.into_bytes(ZCPacketType::UDP);
|
||||
let buf = packet.into_bytes();
|
||||
tracing::trace!(?udp_payload_len, ?buf, "udp forward from ring to udp");
|
||||
let ret = socket.send_to(&buf, &addr).await;
|
||||
if ret.is_err() {
|
||||
@@ -232,7 +233,7 @@ impl UdpTunnelListenerData {
|
||||
tracing::info!(?conn_id, ?remote_addr, "udp connection accept handling",);
|
||||
let socket = self.socket.as_ref().unwrap().clone();
|
||||
|
||||
let sack_buf = new_sack_packet(conn_id, magic).into_bytes(ZCPacketType::UDP);
|
||||
let sack_buf = new_sack_packet(conn_id, magic).into_bytes();
|
||||
if let Err(e) = socket.send_to(&sack_buf, remote_addr).await {
|
||||
tracing::error!(?e, "udp send sack packet error");
|
||||
return;
|
||||
@@ -436,6 +437,7 @@ impl TunnelListener for UdpTunnelListener {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct UdpTunnelConnector {
|
||||
addr: url::Url,
|
||||
bind_addrs: Vec<SocketAddr>,
|
||||
@@ -613,7 +615,7 @@ impl UdpTunnelConnector {
|
||||
// send syn
|
||||
let conn_id = rand::random();
|
||||
let magic = rand::random();
|
||||
let udp_packet = new_syn_packet(conn_id, magic).into_bytes(ZCPacketType::UDP);
|
||||
let udp_packet = new_syn_packet(conn_id, magic).into_bytes();
|
||||
let ret = socket.send_to(&udp_packet, &addr).await?;
|
||||
tracing::warn!(?udp_packet, ?ret, "udp send syn");
|
||||
|
||||
|
||||
@@ -138,17 +138,15 @@ impl Debug for WgPeerData {
|
||||
|
||||
impl WgPeerData {
|
||||
#[tracing::instrument]
|
||||
async fn handle_one_packet_from_me(
|
||||
&self,
|
||||
mut zc_packet: ZCPacket,
|
||||
) -> Result<(), anyhow::Error> {
|
||||
async fn handle_one_packet_from_me(&self, zc_packet: ZCPacket) -> Result<(), anyhow::Error> {
|
||||
let mut send_buf = vec![0u8; MAX_PACKET];
|
||||
|
||||
let packet = if matches!(self.wg_type, WgType::InternalUse) {
|
||||
let mut zc_packet = zc_packet.convert_type(ZCPacketType::WG);
|
||||
Self::fill_ip_header(&mut zc_packet);
|
||||
zc_packet.into_bytes(ZCPacketType::WG)
|
||||
zc_packet.into_bytes()
|
||||
} else {
|
||||
zc_packet.into_bytes(ZCPacketType::WG)
|
||||
zc_packet.convert_type(ZCPacketType::WG).into_bytes()
|
||||
};
|
||||
tracing::trace!(?packet, "Sending packet to peer");
|
||||
|
||||
@@ -650,7 +648,7 @@ impl WgTunnelConnector {
|
||||
let mut buf = vec![0u8; MAX_PACKET];
|
||||
let (n, recv_addr) = data.udp.recv_from(&mut buf).await.unwrap();
|
||||
if recv_addr != addr {
|
||||
continue;
|
||||
tracing::warn!(?recv_addr, "Received packet from changed address");
|
||||
}
|
||||
data.handle_one_packet_from_peer(&mut sink, &buf[..n]).await;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user