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

弃用SkipMap

Browse Source
This commit is contained in:
lubeilin
2023-08-27 16:29:44 +08:00
parent 756539c3bd
commit 45d8c6717d
10 changed files with 67 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vnt/Cargo.toml
+1 -1
View File
@@ -11,7 +11,7 @@ bytes = "1.3.0"
log = "0.4.17" log = "0.4.17"
libc = "0.2.137" libc = "0.2.137"
crossbeam-utils = "0.8" crossbeam-utils = "0.8"
crossbeam-skiplist = "0.1.1" dashmap = "5.5.1"
parking_lot = "0.12.1" parking_lot = "0.12.1"
rand = "0.8.5" rand = "0.8.5"
sha2 = { version = "0.10.6", features = ["oid"] } sha2 = { version = "0.10.6", features = ["oid"] }
vnt/src/channel/channel.rs
+22 -30
View File
@@ -2,9 +2,8 @@ use std::io;
use std::net::{Ipv4Addr, SocketAddr}; use std::net::{Ipv4Addr, SocketAddr};
use std::sync::Arc; use std::sync::Arc;
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use crossbeam_skiplist::SkipMap;
use crossbeam_utils::atomic::AtomicCell; use crossbeam_utils::atomic::AtomicCell;
use parking_lot::Mutex; use dashmap::DashMap;
use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpStream, UdpSocket}; use tokio::net::{TcpStream, UdpSocket};
use tokio::net::tcp::OwnedReadHalf; use tokio::net::tcp::OwnedReadHalf;
@@ -16,16 +15,15 @@ use crate::handle::CurrentDeviceInfo;
use crate::handle::recv_handler::ChannelDataHandler; use crate::handle::recv_handler::ChannelDataHandler;
pub struct ContextInner { pub struct ContextInner {
pub(crate) lock: Mutex<()>,
//udp用于打洞、服务端通信(可选) //udp用于打洞、服务端通信(可选)
pub(crate) main_channel: Arc<UdpSocket>, pub(crate) main_channel: Arc<UdpSocket>,
//在udp的基础上，可以选择使用tcp和服务端通信 //在udp的基础上，可以选择使用tcp和服务端通信
pub(crate) main_tcp_channel: Option<tokio::sync::mpsc::Sender<Vec<u8>>>, pub(crate) main_tcp_channel: Option<tokio::sync::mpsc::Sender<Vec<u8>>>,
pub(crate) route_table: SkipMap<Ipv4Addr, Vec<Route>>, pub(crate) route_table: DashMap<Ipv4Addr, Vec<Route>>,
pub(crate) route_table_time: SkipMap<(RouteKey, Ipv4Addr), AtomicCell<Instant>>, pub(crate) route_table_time: DashMap<(RouteKey, Ipv4Addr), AtomicCell<Instant>>,
pub(crate) status_receiver: Receiver<Status>, pub(crate) status_receiver: Receiver<Status>,
pub(crate) status_sender: Sender<Status>, pub(crate) status_sender: Sender<Status>,
pub(crate) udp_map: SkipMap<usize, Arc<UdpSocket>>, pub(crate) udp_map: DashMap<usize, Arc<UdpSocket>>,
pub(crate) channel_num: usize, pub(crate) channel_num: usize,
current_device: Arc<AtomicCell<CurrentDeviceInfo>>, current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
} }
@@ -41,14 +39,13 @@ impl Context {
let channel_num = 1; let channel_num = 1;
let (status_sender, status_receiver) = channel(Status::Cone); let (status_sender, status_receiver) = channel(Status::Cone);
let inner = Arc::new(ContextInner { let inner = Arc::new(ContextInner {
lock: Mutex::new(()),
main_channel, main_channel,
main_tcp_channel, main_tcp_channel,
route_table: SkipMap::new(), route_table: DashMap::with_capacity(16),
route_table_time: SkipMap::new(), route_table_time: DashMap::with_capacity(16),
status_receiver, status_receiver,
status_sender, status_sender,
udp_map: SkipMap::new(), udp_map: DashMap::new(),
channel_num, channel_num,
current_device, current_device,
}); });
@@ -117,8 +114,10 @@ impl Context {
} }
pub(crate) async fn send_all(&self, buf: &[u8], addr: SocketAddr) -> io::Result<()> { pub(crate) async fn send_all(&self, buf: &[u8], addr: SocketAddr) -> io::Result<()> {
for udp in self.inner.udp_map.iter() { for udp_ref in self.inner.udp_map.iter() {
udp.value().send_to(buf, addr).await?; let udp = udp_ref.clone();
drop(udp_ref);
udp.send_to(buf, addr).await?;
} }
Ok(()) Ok(())
} }
@@ -139,8 +138,10 @@ impl Context {
} }
} }
if let Some(udp) = self.inner.udp_map.get(&route.index) { if let Some(udp_ref) = self.inner.udp_map.get(&route.index) {
return udp.value().send_to(buf, route.addr).await; let udp = udp_ref.value().clone();
drop(udp_ref);
return udp.send_to(buf, route.addr).await;
} }
} }
Err(io::Error::new(io::ErrorKind::NotFound, "route not found")) Err(io::Error::new(io::ErrorKind::NotFound, "route not found"))
@@ -171,8 +172,10 @@ impl Context {
}; };
} }
} }
if let Some(udp) = self.inner.udp_map.get(&route_key.index) { if let Some(udp_ref) = self.inner.udp_map.get(&route_key.index) {
return udp.value().send_to(buf, route_key.addr).await; let udp = udp_ref.value().clone();
drop(udp_ref);
return udp.send_to(buf, route_key.addr).await;
} }
Err(io::Error::new(io::ErrorKind::NotFound, "route not found")) Err(io::Error::new(io::ErrorKind::NotFound, "route not found"))
} }
@@ -201,12 +204,7 @@ impl Context {
} }
fn add_route_(&self, id: Ipv4Addr, route: Route, only_if_absent: bool) { fn add_route_(&self, id: Ipv4Addr, route: Route, only_if_absent: bool) {
let key = route.route_key(); let key = route.route_key();
let guard = self.inner.lock.lock(); let mut list = self.inner.route_table.entry(id).or_insert_with(||Vec::with_capacity(4));
let mut list = if let Some(entry) = self.inner.route_table.get(&id) {
entry.value().clone()
} else {
Vec::with_capacity(4)
};
let mut exist = false; let mut exist = false;
for x in list.iter_mut() { for x in list.iter_mut() {
if x.metric < route.metric { if x.metric < route.metric {
@@ -237,9 +235,7 @@ impl Context {
list.truncate(max_len); list.truncate(max_len);
} }
} }
self.inner.route_table.insert(id, list);
self.inner.route_table_time.insert((key, id), AtomicCell::new(Instant::now())); self.inner.route_table_time.insert((key, id), AtomicCell::new(Instant::now()));
drop(guard);
} }
pub fn route(&self, id: &Ipv4Addr) -> Option<Vec<Route>> { pub fn route(&self, id: &Ipv4Addr) -> Option<Vec<Route>> {
if let Some(v) = self.inner.route_table.get(id) { if let Some(v) = self.inner.route_table.get(id) {
@@ -297,16 +293,13 @@ impl Context {
v v
} }
pub fn remove_route_all(&self, id: &Ipv4Addr) { pub fn remove_route_all(&self, id: &Ipv4Addr) {
let guard = self.inner.lock.lock(); if let Some((_,routes)) = self.inner.route_table.remove(id) {
if let Some(v) = self.inner.route_table.remove(id) { for x in routes {
for x in v.value() {
self.inner.route_table_time.remove(&(x.route_key(), *id)); self.inner.route_table_time.remove(&(x.route_key(), *id));
} }
} }
drop(guard);
} }
pub fn remove_route(&self, id: &Ipv4Addr, route_key: RouteKey) { pub fn remove_route(&self, id: &Ipv4Addr, route_key: RouteKey) {
let guard = self.inner.lock.lock();
if let Some(v) = self.inner.route_table.get(id) { if let Some(v) = self.inner.route_table.get(id) {
let mut routes = v.value().clone(); let mut routes = v.value().clone();
drop(v); drop(v);
@@ -314,7 +307,6 @@ impl Context {
self.inner.route_table.insert(*id, routes); self.inner.route_table.insert(*id, routes);
} }
self.inner.route_table_time.remove(&(route_key, *id)); self.inner.route_table_time.remove(&(route_key, *id));
drop(guard);
} }
pub fn update_read_time(&self, id: &Ipv4Addr, route_key: &RouteKey) { pub fn update_read_time(&self, id: &Ipv4Addr, route_key: &RouteKey) {
if let Some(time) = self.inner.route_table_time.get(&(*route_key, *id)) { if let Some(time) = self.inner.route_table_time.get(&(*route_key, *id)) {
vnt/src/channel/idle.rs
-1
View File
@@ -29,7 +29,6 @@ impl Idle {
for entry in self.context.inner.route_table_time.iter() { for entry in self.context.inner.route_table_time.iter() {
let last_read = entry.value().load().elapsed(); let last_read = entry.value().load().elapsed();
if last_read >= self.read_idle { if last_read >= self.read_idle {
entry.remove();
return Ok((entry.key().1.clone(), entry.key().0.clone())); return Ok((entry.key().1.clone(), entry.key().0.clone()));
} else { } else {
if max < last_read { if max < last_read {
vnt/src/core/mod.rs
+3 -3
View File
@@ -3,8 +3,8 @@ use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use crossbeam_skiplist::SkipMap;
use crossbeam_utils::atomic::AtomicCell; use crossbeam_utils::atomic::AtomicCell;
use dashmap::DashMap;
use parking_lot::Mutex; use parking_lot::Mutex;
use rand::Rng; use rand::Rng;
use tokio::net::{TcpStream, UdpSocket}; use tokio::net::{TcpStream, UdpSocket};
@@ -48,7 +48,7 @@ pub struct Vnt {
device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>, device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
nat_test: NatTest, nat_test: NatTest,
connect_status: Arc<AtomicCell<ConnectStatus>>, connect_status: Arc<AtomicCell<ConnectStatus>>,
peer_nat_info_map: Arc<SkipMap<Ipv4Addr, NatInfo>>, peer_nat_info_map: Arc<DashMap<Ipv4Addr, NatInfo>>,
} }
pub struct VntUtil { pub struct VntUtil {
@@ -208,7 +208,7 @@ impl VntUtil {
config.server_address, config.token.clone(), config.server_address, config.token.clone(),
config.device_id.clone(), config.name.clone(),config.password.is_some())); config.device_id.clone(), config.name.clone(),config.password.is_some()));
let device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>> = Arc::new(Mutex::new((response.epoch, response.device_info_list))); let device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>> = Arc::new(Mutex::new((response.epoch, response.device_info_list)));
let peer_nat_info_map: Arc<SkipMap<Ipv4Addr, NatInfo>> = Arc::new(SkipMap::new()); let peer_nat_info_map: Arc<DashMap<Ipv4Addr, NatInfo>> = Arc::new(DashMap::new());
let connect_status = Arc::new(AtomicCell::new(ConnectStatus::Connected)); let connect_status = Arc::new(AtomicCell::new(ConnectStatus::Connected));
vnt/src/handle/recv_handler.rs
+5 -9
View File
@@ -1,8 +1,8 @@
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::sync::Arc; use std::sync::Arc;
use crossbeam_skiplist::SkipMap;
use crossbeam_utils::atomic::AtomicCell; use crossbeam_utils::atomic::AtomicCell;
use dashmap::DashMap;
use parking_lot::Mutex; use parking_lot::Mutex;
use protobuf::Message; use protobuf::Message;
use tokio::sync::mpsc::Sender; use tokio::sync::mpsc::Sender;
@@ -41,7 +41,7 @@ pub struct ChannelDataHandler {
igmp_server: Option<IgmpServer>, igmp_server: Option<IgmpServer>,
device_writer: DeviceWriter, device_writer: DeviceWriter,
connect_status: Arc<AtomicCell<ConnectStatus>>, connect_status: Arc<AtomicCell<ConnectStatus>>,
peer_nat_info_map: Arc<SkipMap<Ipv4Addr, NatInfo>>, peer_nat_info_map: Arc<DashMap<Ipv4Addr, NatInfo>>,
ip_proxy_map: Option<IpProxyMap>, ip_proxy_map: Option<IpProxyMap>,
out_external_route: AllowExternalRoute, out_external_route: AllowExternalRoute,
cone_sender: Sender<(Ipv4Addr, NatInfo)>, cone_sender: Sender<(Ipv4Addr, NatInfo)>,
@@ -61,7 +61,7 @@ impl ChannelDataHandler {
igmp_server: Option<IgmpServer>, igmp_server: Option<IgmpServer>,
device_writer: DeviceWriter, device_writer: DeviceWriter,
connect_status: Arc<AtomicCell<ConnectStatus>>, connect_status: Arc<AtomicCell<ConnectStatus>>,
peer_nat_info_map: Arc<SkipMap<Ipv4Addr, NatInfo>>, peer_nat_info_map: Arc<DashMap<Ipv4Addr, NatInfo>>,
ip_proxy_map: Option<IpProxyMap>, ip_proxy_map: Option<IpProxyMap>,
out_external_route: AllowExternalRoute, out_external_route: AllowExternalRoute,
cone_sender: Sender<(Ipv4Addr, NatInfo)>, cone_sender: Sender<(Ipv4Addr, NatInfo)>,
@@ -196,9 +196,7 @@ impl ChannelDataHandler {
ipv4.update_checksum(); ipv4.update_checksum();
let key = SocketAddrV4::new(source, source_port); let key = SocketAddrV4::new(source, source_port);
//https://github.com/crossbeam-rs/crossbeam/issues/1023 //https://github.com/crossbeam-rs/crossbeam/issues/1023
if !ip_proxy_map.tcp_proxy_map.contains_key(&key){ ip_proxy_map.tcp_proxy_map.insert(key, SocketAddrV4::new(dest_ip, dest_port));
ip_proxy_map.tcp_proxy_map.insert(key, SocketAddrV4::new(dest_ip, dest_port));
}
} }
ipv4::protocol::Protocol::Udp => { ipv4::protocol::Protocol::Udp => {
let dest_ip = ipv4.destination_ip(); let dest_ip = ipv4.destination_ip();
@@ -211,9 +209,7 @@ impl ChannelDataHandler {
ipv4.set_destination_ip(destination); ipv4.set_destination_ip(destination);
ipv4.update_checksum(); ipv4.update_checksum();
let key = SocketAddrV4::new(source, source_port); let key = SocketAddrV4::new(source, source_port);
if !ip_proxy_map.udp_proxy_map.contains_key(&key){ ip_proxy_map.udp_proxy_map.insert(key, SocketAddrV4::new(dest_ip, dest_port));
ip_proxy_map.udp_proxy_map.insert(key, SocketAddrV4::new(dest_ip, dest_port));
}
} }
ipv4::protocol::Protocol::Icmp => { ipv4::protocol::Protocol::Icmp => {
let dest_ip = ipv4.destination_ip(); let dest_ip = ipv4.destination_ip();
vnt/src/igmp_server/mod.rs
+12 -10
View File
@@ -2,7 +2,7 @@ use std::collections::{HashMap, HashSet};
use std::net::Ipv4Addr; use std::net::Ipv4Addr;
use std::sync::Arc; use std::sync::Arc;
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use crossbeam_skiplist::SkipMap; use dashmap::DashMap;
use parking_lot::RwLock; use parking_lot::RwLock;
use packet::igmp::igmp_v2::IgmpV2Packet; use packet::igmp::igmp_v2::IgmpV2Packet;
use packet::igmp::igmp_v3::{IgmpV3QueryPacket, IgmpV3RecordType, IgmpV3ReportPacket}; use packet::igmp::igmp_v3::{IgmpV3QueryPacket, IgmpV3RecordType, IgmpV3ReportPacket};
@@ -47,12 +47,12 @@ impl Multicast {
#[derive(Clone)] #[derive(Clone)]
pub struct IgmpServer { pub struct IgmpServer {
multicast: Arc<SkipMap<Ipv4Addr, Arc<RwLock<Multicast>>>>, multicast: Arc<DashMap<Ipv4Addr, Arc<RwLock<Multicast>>>>,
} }
impl IgmpServer { impl IgmpServer {
pub fn new(device_writer: DeviceWriter) -> Self { pub fn new(device_writer: DeviceWriter) -> Self {
let multicast: Arc<SkipMap<Ipv4Addr, Arc<RwLock<Multicast>>>> = Arc::new(SkipMap::new()); let multicast: Arc<DashMap<Ipv4Addr, Arc<RwLock<Multicast>>>> = Arc::new(DashMap::new());
std::thread::spawn(move || { std::thread::spawn(move || {
//预留以太网帧头和ip头 //预留以太网帧头和ip头
let mut buf = [0; 14 + 24 + 12]; let mut buf = [0; 14 + 24 + 12];
@@ -124,10 +124,12 @@ impl IgmpServer {
if !multicast_addr.is_multicast() { if !multicast_addr.is_multicast() {
return Ok(()); return Ok(());
} }
let multi = self.multicast.get_or_insert_with(multicast_addr, || { let multi = {
Arc::new(RwLock::new(Multicast::new())) self.multicast.entry(multicast_addr).or_insert_with(|| {
}); Arc::new(RwLock::new(Multicast::new()))
let mut guard = multi.value().write(); }).value().clone()
};
let mut guard = multi.write();
guard.members.insert(source, Instant::now()); guard.members.insert(source, Instant::now());
} }
IgmpType::LeaveV2 => { IgmpType::LeaveV2 => {
@@ -151,10 +153,10 @@ impl IgmpServer {
if !multicast_addr.is_multicast() { if !multicast_addr.is_multicast() {
return Ok(()); return Ok(());
} }
let multi = self.multicast.get_or_insert_with(multicast_addr, || { let multi = self.multicast.entry(multicast_addr).or_insert_with(|| {
Arc::new(RwLock::new(Multicast::new())) Arc::new(RwLock::new(Multicast::new()))
}); }).value().clone();
let mut guard = multi.value().write(); let mut guard = multi.write();
match group_record.record_type() { match group_record.record_type() {
IgmpV3RecordType::ModeIsInclude | IgmpV3RecordType::ChangeToIncludeMode => { IgmpV3RecordType::ModeIsInclude | IgmpV3RecordType::ChangeToIncludeMode => {
vnt/src/ip_proxy/icmp_proxy.rs
+4 -3
View File
@@ -3,8 +3,8 @@ use std::mem::MaybeUninit;
use std::net::{IpAddr, Ipv4Addr, SocketAddrV4}; use std::net::{IpAddr, Ipv4Addr, SocketAddrV4};
use std::sync::Arc; use std::sync::Arc;
use crossbeam_utils::atomic::AtomicCell; use crossbeam_utils::atomic::AtomicCell;
use dashmap::DashMap;
use crossbeam_skiplist::SkipMap;
use socket2::{Domain, SockAddr, Socket, Type}; use socket2::{Domain, SockAddr, Socket, Type};
use packet::icmp::icmp; use packet::icmp::icmp;
@@ -19,14 +19,14 @@ use crate::protocol::body::ENCRYPTION_RESERVED;
pub struct IcmpProxy { pub struct IcmpProxy {
icmp_socket: Arc<Socket>, icmp_socket: Arc<Socket>,
// 对端-> 真实来源 // 对端-> 真实来源
icmp_proxy_map: Arc<SkipMap<(Ipv4Addr, u16, u16), Ipv4Addr>>, icmp_proxy_map: Arc<DashMap<(Ipv4Addr, u16, u16), Ipv4Addr>>,
sender: ChannelSender, sender: ChannelSender,
current_device: Arc<AtomicCell<CurrentDeviceInfo>>, current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
client_cipher: Cipher, client_cipher: Cipher,
} }
impl IcmpProxy { impl IcmpProxy {
pub fn new(addr: SocketAddrV4, icmp_proxy_map: Arc<SkipMap<(Ipv4Addr, u16, u16), Ipv4Addr>>, pub fn new(addr: SocketAddrV4, icmp_proxy_map: Arc<DashMap<(Ipv4Addr, u16, u16), Ipv4Addr>>,
sender: ChannelSender, current_device: Arc<AtomicCell<CurrentDeviceInfo>>, client_cipher: Cipher) -> io::Result<IcmpProxy> { sender: ChannelSender, current_device: Arc<AtomicCell<CurrentDeviceInfo>>, client_cipher: Cipher) -> io::Result<IcmpProxy> {
let icmp_socket = Arc::new(Socket::new(Domain::IPV4, Type::RAW, Some(socket2::Protocol::ICMPV4))?); let icmp_socket = Arc::new(Socket::new(Domain::IPV4, Type::RAW, Some(socket2::Protocol::ICMPV4))?);
icmp_socket.bind(&SockAddr::from(addr))?; icmp_socket.bind(&SockAddr::from(addr))?;
@@ -60,6 +60,7 @@ impl IcmpProxy {
if let Some(entry) = self.icmp_proxy_map.get(&(peer_ip, id, seq)) { if let Some(entry) = self.icmp_proxy_map.get(&(peer_ip, id, seq)) {
//将数据发送到真实的来源 //将数据发送到真实的来源
let dest_ip = *entry.value(); let dest_ip = *entry.value();
drop(entry);
ipv4_packet.set_destination_ip(dest_ip); ipv4_packet.set_destination_ip(dest_ip);
ipv4_packet.update_checksum(); ipv4_packet.update_checksum();
let current_device = self.current_device.load(); let current_device = self.current_device.load();
vnt/src/ip_proxy/mod.rs
+7 -7
View File
@@ -2,7 +2,7 @@ use std::{io, thread};
use std::net::{Ipv4Addr, SocketAddrV4}; use std::net::{Ipv4Addr, SocketAddrV4};
use std::sync::Arc; use std::sync::Arc;
use crossbeam_utils::atomic::AtomicCell; use crossbeam_utils::atomic::AtomicCell;
use crossbeam_skiplist::SkipMap; use dashmap::DashMap;
use socket2::{SockAddr, Socket}; use socket2::{SockAddr, Socket};
use tokio::net::{TcpListener, UdpSocket}; use tokio::net::{TcpListener, UdpSocket};
use crate::channel::sender::ChannelSender; use crate::channel::sender::ChannelSender;
@@ -28,10 +28,10 @@ pub struct IpProxyMap {
pub(crate) tcp_proxy_port: u16, pub(crate) tcp_proxy_port: u16,
pub(crate) udp_proxy_port: u16, pub(crate) udp_proxy_port: u16,
//真实源地址 -> 目的地址 //真实源地址 -> 目的地址
pub(crate) tcp_proxy_map: Arc<SkipMap<SocketAddrV4, SocketAddrV4>>, pub(crate) tcp_proxy_map: Arc<DashMap<SocketAddrV4, SocketAddrV4>>,
pub(crate) udp_proxy_map: Arc<SkipMap<SocketAddrV4, SocketAddrV4>>, pub(crate) udp_proxy_map: Arc<DashMap<SocketAddrV4, SocketAddrV4>>,
// icmp用Identifier来区分，没有Identifier的一律不转发 // icmp用Identifier来区分，没有Identifier的一律不转发
pub(crate) icmp_proxy_map: Arc<SkipMap<(Ipv4Addr, u16, u16), Ipv4Addr>>, pub(crate) icmp_proxy_map: Arc<DashMap<(Ipv4Addr, u16, u16), Ipv4Addr>>,
icmp_socket: Arc<Socket>, icmp_socket: Arc<Socket>,
} }
@@ -42,9 +42,9 @@ impl IpProxyMap {
} }
pub async fn init_proxy(sender: ChannelSender, current_device: Arc<AtomicCell<CurrentDeviceInfo>>, client_cipher: Cipher,) -> io::Result<(TcpProxy, UdpProxy, IpProxyMap)> { pub async fn init_proxy(sender: ChannelSender, current_device: Arc<AtomicCell<CurrentDeviceInfo>>, client_cipher: Cipher,) -> io::Result<(TcpProxy, UdpProxy, IpProxyMap)> {
let tcp_proxy_map: Arc<SkipMap<SocketAddrV4, SocketAddrV4>> = Arc::new(SkipMap::new()); let tcp_proxy_map: Arc<DashMap<SocketAddrV4, SocketAddrV4>> = Arc::new(DashMap::new());
let udp_proxy_map: Arc<SkipMap<SocketAddrV4, SocketAddrV4>> = Arc::new(SkipMap::new()); let udp_proxy_map: Arc<DashMap<SocketAddrV4, SocketAddrV4>> = Arc::new(DashMap::new());
let icmp_proxy_map: Arc<SkipMap<(Ipv4Addr, u16, u16), Ipv4Addr>> = Arc::new(SkipMap::new()); let icmp_proxy_map: Arc<DashMap<(Ipv4Addr, u16, u16), Ipv4Addr>> = Arc::new(DashMap::new());
let tcp_listener = TcpListener::bind("0.0.0.0:0").await?; let tcp_listener = TcpListener::bind("0.0.0.0:0").await?;
let udp_socket = UdpSocket::bind("0.0.0.0:0").await?; let udp_socket = UdpSocket::bind("0.0.0.0:0").await?;
let tcp_proxy_port = tcp_listener.local_addr()?.port(); let tcp_proxy_port = tcp_listener.local_addr()?.port();
vnt/src/ip_proxy/tcp_proxy.rs
+4 -3
View File
@@ -1,17 +1,17 @@
use std::io; use std::io;
use std::net::{SocketAddr, SocketAddrV4}; use std::net::{SocketAddr, SocketAddrV4};
use std::sync::Arc; use std::sync::Arc;
use dashmap::DashMap;
use crossbeam_skiplist::SkipMap;
use tokio::net::{TcpListener, TcpStream}; use tokio::net::{TcpListener, TcpStream};
pub struct TcpProxy { pub struct TcpProxy {
tcp_listener: TcpListener, tcp_listener: TcpListener,
tcp_proxy_map: Arc<SkipMap<SocketAddrV4, SocketAddrV4>>, tcp_proxy_map: Arc<DashMap<SocketAddrV4, SocketAddrV4>>,
} }
impl TcpProxy { impl TcpProxy {
pub fn new(tcp_listener: TcpListener, tcp_proxy_map: Arc<SkipMap<SocketAddrV4, SocketAddrV4>>) -> Self { pub fn new(tcp_listener: TcpListener, tcp_proxy_map: Arc<DashMap<SocketAddrV4, SocketAddrV4>>) -> Self {
Self { Self {
tcp_listener, tcp_listener,
tcp_proxy_map, tcp_proxy_map,
@@ -27,6 +27,7 @@ impl TcpProxy {
SocketAddr::V4(sender_addr) => { SocketAddr::V4(sender_addr) => {
if let Some(entry) = tcp_proxy_map.get(&sender_addr) { if let Some(entry) = tcp_proxy_map.get(&sender_addr) {
let dest_addr = *entry.value(); let dest_addr = *entry.value();
drop(entry);
let peer_tcp_stream = match TcpStream::connect(dest_addr).await { let peer_tcp_stream = match TcpStream::connect(dest_addr).await {
Ok(peer_tcp_stream) => { peer_tcp_stream } Ok(peer_tcp_stream) => { peer_tcp_stream }
Err(e) => { Err(e) => {
vnt/src/ip_proxy/udp_proxy.rs
+9 -6
View File
@@ -2,17 +2,17 @@ use std::io;
use std::net::{SocketAddr, SocketAddrV4}; use std::net::{SocketAddr, SocketAddrV4};
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use crossbeam_skiplist::SkipMap; use dashmap::DashMap;
use tokio::net::UdpSocket; use tokio::net::UdpSocket;
/// 一个udp代理，作用是利用系统协议栈，将udp数据报解析出来再转发到目的地址 /// 一个udp代理，作用是利用系统协议栈，将udp数据报解析出来再转发到目的地址
pub struct UdpProxy { pub struct UdpProxy {
udp_socket: Arc<UdpSocket>, udp_socket: Arc<UdpSocket>,
map: Arc<SkipMap<SocketAddrV4, SocketAddrV4>>, map: Arc<DashMap<SocketAddrV4, SocketAddrV4>>,
} }
impl UdpProxy { impl UdpProxy {
pub fn new(udp_socket: UdpSocket, map: Arc<SkipMap<SocketAddrV4, SocketAddrV4>>) -> Self { pub fn new(udp_socket: UdpSocket, map: Arc<DashMap<SocketAddrV4, SocketAddrV4>>) -> Self {
let udp_socket = Arc::new(udp_socket); let udp_socket = Arc::new(udp_socket);
Self { Self {
udp_socket, udp_socket,
@@ -23,7 +23,7 @@ impl UdpProxy {
let map = self.map; let map = self.map;
let udp_socket = self.udp_socket; let udp_socket = self.udp_socket;
let mut buf = [0u8; 65536]; let mut buf = [0u8; 65536];
let inner_map: Arc<SkipMap<SocketAddrV4, Arc<UdpSocket>>> = Arc::new(SkipMap::new()); let inner_map: Arc<DashMap<SocketAddrV4, Arc<UdpSocket>>> = Arc::new(DashMap::new());
loop { loop {
match udp_socket.recv_from(&mut buf).await { match udp_socket.recv_from(&mut buf).await {
@@ -48,11 +48,14 @@ impl UdpProxy {
} }
} }
async fn start0(buf: &[u8], sender_addr: SocketAddrV4, inner_map: &Arc<SkipMap<SocketAddrV4, Arc<UdpSocket>>>, map: &Arc<SkipMap<SocketAddrV4, SocketAddrV4>>, udp_socket: &Arc<UdpSocket>) -> io::Result<()> { async fn start0(buf: &[u8], sender_addr: SocketAddrV4, inner_map: &Arc<DashMap<SocketAddrV4, Arc<UdpSocket>>>, map: &Arc<DashMap<SocketAddrV4, SocketAddrV4>>, udp_socket: &Arc<UdpSocket>) -> io::Result<()> {
if let Some(entry) = inner_map.get(&sender_addr) { if let Some(entry) = inner_map.get(&sender_addr) {
entry.value().send(buf).await?; let udp = entry.value().clone();
drop(entry);
udp.send(buf).await?;
} else if let Some(entry) = map.get(&sender_addr) { } else if let Some(entry) = map.get(&sender_addr) {
let dest_addr = *entry.value(); let dest_addr = *entry.value();
drop(entry);
let peer_udp_socket = UdpSocket::bind("0.0.0.0:0").await?; let peer_udp_socket = UdpSocket::bind("0.0.0.0:0").await?;
peer_udp_socket.connect(dest_addr).await?; peer_udp_socket.connect(dest_addr).await?;
peer_udp_socket.send(buf).await?; peer_udp_socket.send(buf).await?;