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

去除不安全的计数器

Browse Source
This commit is contained in:
lbl8603
2024-06-30 23:02:50 +08:00
parent 8658bd3601
commit 1df5a1cfb1
7 changed files with 29 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vnt/src/core/conn.rs
+5 -8
View File
@@ -27,9 +27,7 @@ use crate::nat::NatTest;
#[cfg(feature = "integrated_tun")]
use crate::tun_tap_device::tun_create_helper::{DeviceAdapter, TunDeviceHelper};
use crate::tun_tap_device::vnt_device::DeviceWrite;
use crate::util::{
Scheduler, SingleU64Adder, StopManager, U64Adder, WatchSingleU64Adder, WatchU64Adder,
};
use crate::util::{Scheduler, StopManager, U64Adder, WatchU64Adder};
use crate::{nat, VntCallback};
#[derive(Clone)]
@@ -68,7 +66,7 @@ pub struct VntInner {
context: Arc<Mutex<Option<ChannelContext>>>,
peer_nat_info_map: Arc<RwLock<HashMap<Ipv4Addr, NatInfo>>>,
down_count_watcher: WatchU64Adder,
up_count_watcher: WatchSingleU64Adder,
up_count_watcher: WatchU64Adder,
client_secret_hash: Option<[u8; 16]>,
compressor: Compressor,
client_cipher: Cipher,
@@ -201,14 +199,13 @@ impl VntInner {
let (punch_sender, punch_receiver) = maintain::punch_channel();
let peer_nat_info_map: Arc<RwLock<HashMap<Ipv4Addr, NatInfo>>> =
Arc::new(RwLock::new(HashMap::with_capacity(16)));
let down_counter =
U64Adder::with_capacity(config.ports.as_ref().map(|v| v.len()).unwrap_or_default() + 8);
let down_counter = U64Adder::default();
let down_count_watcher = down_counter.watch();
let handshake = Handshake::new(
#[cfg(feature = "server_encrypt")]
rsa_cipher.clone(),
);
let up_counter = SingleU64Adder::new();
let up_counter = U64Adder::default();
let up_count_watcher = up_counter.watch();
#[cfg(feature = "integrated_tun")]
let tun_device_helper = {
@@ -348,7 +345,7 @@ pub fn start<Call: VntCallback>(
punch: Punch,
callback: Call,
down_count_watcher: WatchU64Adder,
up_count_watcher: WatchSingleU64Adder,
up_count_watcher: WatchU64Adder,
) {
// 定时心跳
maintain::heartbeat(
vnt/src/handle/maintain/up_status.rs
+4 -4
View File
@@ -3,7 +3,7 @@ use crate::handle::CurrentDeviceInfo;
use crate::proto::message::{ClientStatusInfo, PunchNatType, RouteItem};
use crate::protocol::body::ENCRYPTION_RESERVED;
use crate::protocol::{service_packet, NetPacket, Protocol, HEAD_LEN, MAX_TTL};
use crate::util::{Scheduler, WatchSingleU64Adder, WatchU64Adder};
use crate::util::{Scheduler, WatchU64Adder};
use crossbeam_utils::atomic::AtomicCell;
use protobuf::Message;
use std::io;
@@ -16,7 +16,7 @@ pub fn up_status(
context: ChannelContext,
current_device_info: Arc<AtomicCell<CurrentDeviceInfo>>,
down_count_watcher: WatchU64Adder,
up_count_watcher: WatchSingleU64Adder,
up_count_watcher: WatchU64Adder,
) {
let _ = scheduler.timeout(Duration::from_secs(60), move |x| {
up_status0(
@@ -34,7 +34,7 @@ fn up_status0(
context: ChannelContext,
current_device_info: Arc<AtomicCell<CurrentDeviceInfo>>,
down_count_watcher: WatchU64Adder,
up_count_watcher: WatchSingleU64Adder,
up_count_watcher: WatchU64Adder,
) {
if let Err(e) = send_up_status_packet(
&context,
@@ -62,7 +62,7 @@ fn send_up_status_packet(
context: &ChannelContext,
current_device_info: &AtomicCell<CurrentDeviceInfo>,
down_count_watcher: &WatchU64Adder,
up_count_watcher: &WatchSingleU64Adder,
up_count_watcher: &WatchU64Adder,
) -> io::Result<()> {
let device_info = current_device_info.load();
if device_info.status.offline() {
vnt/src/handle/tun_tap/tun_handler.rs
+3 -3
View File
@@ -26,7 +26,7 @@ use crate::protocol;
use crate::protocol::body::ENCRYPTION_RESERVED;
use crate::protocol::ip_turn_packet::BroadcastPacket;
use crate::protocol::{ip_turn_packet, NetPacket, MAX_TTL};
use crate::util::{SingleU64Adder, StopManager};
use crate::util::{StopManager, U64Adder};
fn icmp(device_writer: &Device, mut ipv4_packet: IpV4Packet<&mut [u8]>) -> anyhow::Result<()> {
if ipv4_packet.protocol() == Protocol::Icmp {
@@ -53,7 +53,7 @@ pub fn start(
#[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
client_cipher: Cipher,
server_cipher: Cipher,
mut up_counter: SingleU64Adder,
up_counter: U64Adder,
device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
compressor: Compressor,
device_stop: DeviceStop,
@@ -71,7 +71,7 @@ pub fn start(
ip_proxy_map,
client_cipher,
server_cipher,
&mut up_counter,
&up_counter,
device_list,
compressor,
device_stop,
vnt/src/handle/tun_tap/unix.rs
+3 -3
View File
@@ -7,7 +7,7 @@ use crate::handle::tun_tap::DeviceStop;
use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo};
#[cfg(feature = "ip_proxy")]
use crate::ip_proxy::IpProxyMap;
use crate::util::{SingleU64Adder, StopManager};
use crate::util::{StopManager, U64Adder};
use crossbeam_utils::atomic::AtomicCell;
use mio::event::Source;
use mio::unix::SourceFd;
@@ -30,7 +30,7 @@ pub(crate) fn start_simple(
#[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
client_cipher: Cipher,
server_cipher: Cipher,
up_counter: &mut SingleU64Adder,
up_counter: &U64Adder,
device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
compressor: Compressor,
device_stop: DeviceStop,
@@ -85,7 +85,7 @@ fn start_simple0(
#[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
client_cipher: Cipher,
server_cipher: Cipher,
up_counter: &mut SingleU64Adder,
up_counter: &U64Adder,
device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
compressor: Compressor,
) -> anyhow::Result<()> {
vnt/src/handle/tun_tap/windows.rs
+3 -3
View File
@@ -7,7 +7,7 @@ use crate::handle::tun_tap::DeviceStop;
use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo};
#[cfg(feature = "ip_proxy")]
use crate::ip_proxy::IpProxyMap;
use crate::util::{SingleU64Adder, StopManager};
use crate::util::{StopManager, U64Adder};
use crossbeam_utils::atomic::AtomicCell;
use parking_lot::Mutex;
use std::sync::Arc;
@@ -23,7 +23,7 @@ pub(crate) fn start_simple(
#[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
client_cipher: Cipher,
server_cipher: Cipher,
up_counter: &mut SingleU64Adder,
up_counter: &U64Adder,
device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
compressor: Compressor,
device_stop: DeviceStop,
@@ -76,7 +76,7 @@ fn start_simple0(
#[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
client_cipher: Cipher,
server_cipher: Cipher,
up_counter: &mut SingleU64Adder,
up_counter: &U64Adder,
device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
compressor: Compressor,
) -> anyhow::Result<()> {
vnt/src/tun_tap_device/tun_create_helper.rs
+3 -3
View File
@@ -16,7 +16,7 @@ use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo};
#[cfg(feature = "ip_proxy")]
use crate::ip_proxy::IpProxyMap;
use crate::tun_tap_device::vnt_device::DeviceWrite;
use crate::util::{SingleU64Adder, StopManager};
use crate::util::{StopManager, U64Adder};
#[repr(transparent)]
#[derive(Clone, Default)]
@@ -67,7 +67,7 @@ struct TunDeviceHelperInner {
ip_proxy_map: Option<IpProxyMap>,
client_cipher: Cipher,
server_cipher: Cipher,
up_counter: SingleU64Adder,
up_counter: U64Adder,
device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
compressor: Compressor,
}
@@ -81,7 +81,7 @@ impl TunDeviceHelper {
#[cfg(feature = "ip_proxy")] ip_proxy_map: Option<IpProxyMap>,
client_cipher: Cipher,
server_cipher: Cipher,
up_counter: SingleU64Adder,
up_counter: U64Adder,
device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
compressor: Compressor,
device_adapter: DeviceAdapter,
vnt/src/util/counter/adder.rs
+8 -114
View File
@@ -1,139 +1,33 @@
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
/// 不安全的并发计数器，谨慎使用
use crossbeam_utils::atomic::AtomicCell;
#[derive(Clone, Default)]
pub struct U64Adder {
global_index: Arc<AtomicUsize>,
inner: Arc<U64AdderInner>,
index: usize,
}
#[derive(Clone)]
pub struct SingleU64Adder {
inner: Arc<SingleU64AdderInner>,
}
impl SingleU64Adder {
pub fn new() -> Self {
Self {
inner: Arc::new(SingleU64AdderInner::new()),
}
}
pub fn add(&mut self, num: u64) {
self.inner.add(num);
}
pub fn get(&self) -> u64 {
self.inner.get()
}
pub fn watch(&self) -> WatchSingleU64Adder {
WatchSingleU64Adder {
inner: self.inner.clone(),
}
}
}
struct SingleU64AdderInner {
ptr: *mut u64,
}
impl SingleU64AdderInner {
fn new() -> Self {
Self {
ptr: Box::into_raw(Box::new(0)),
}
}
#[inline(always)]
fn add(&self, num: u64) {
unsafe { *self.ptr += num }
}
fn get(&self) -> u64 {
unsafe { *self.ptr }
}
}
impl Drop for SingleU64AdderInner {
fn drop(&mut self) {
unsafe {
let _ = Box::from_raw(self.ptr);
}
}
}
unsafe impl Send for SingleU64AdderInner {}
unsafe impl Sync for SingleU64AdderInner {}
struct U64AdderInner {
base: Vec<SingleU64AdderInner>,
}
impl U64AdderInner {
pub fn get(&self) -> u64 {
let mut count = 0;
for counter in self.base.iter() {
count += counter.get()
}
count
}
count: Arc<AtomicCell<u64>>,
}
impl U64Adder {
/// 计数槽容量
pub fn with_capacity(capacity: usize) -> Self {
let mut base = Vec::with_capacity(capacity);
for _ in 0..capacity {
base.push(SingleU64AdderInner::new())
}
let inner = Arc::new(U64AdderInner { base });
U64Adder {
global_index: Arc::new(AtomicUsize::new(1)),
inner,
index: 0,
}
}
pub fn add(&self, num: u64) {
self.inner.base[self.index].add(num);
self.count.fetch_add(num);
}
pub fn get(&self) -> u64 {
self.inner.get()
self.count.load()
}
pub fn watch(&self) -> WatchU64Adder {
WatchU64Adder {
inner: self.inner.clone(),
}
}
}
impl Clone for U64Adder {
fn clone(&self) -> Self {
let index = self.global_index.fetch_add(1, Ordering::AcqRel);
if index > self.inner.base.len() {
panic!()
}
Self {
global_index: self.global_index.clone(),
inner: self.inner.clone(),
index,
count: self.count.clone(),
}
}
}
#[derive(Clone)]
pub struct WatchU64Adder {
inner: Arc<U64AdderInner>,
count: Arc<AtomicCell<u64>>,
}
impl WatchU64Adder {
pub fn get(&self) -> u64 {
self.inner.get()
}
}
#[derive(Clone)]
pub struct WatchSingleU64Adder {
inner: Arc<SingleU64AdderInner>,
}
impl WatchSingleU64Adder {
pub fn get(&self) -> u64 {
self.inner.get()
self.count.load()
}
}