[service] Fix check for invalid kext handle (#1716)

* [service] Fix check for invalid kext handle

* [windows_kext] Use BTreeMap as cache structure

* [windows_kext] Fix synchronization bug

* Update windows_kext/kextinterface/kext_file.go

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

* Update windows_kext/kextinterface/kext_file.go

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

* Update windows_kext/kextinterface/kext_file.go

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
This commit is contained in:
Vladimir Stoilov
2024-10-16 12:19:08 +03:00
parent 706ce222d0
commit 05f4f3fc6d
10 changed files with 86 additions and 196 deletions

View File

@@ -1,14 +1,10 @@
use alloc::collections::BTreeMap;
use protocol::info::{BandwidthValueV4, BandwidthValueV6, Info};
use smoltcp::wire::{IpProtocol, Ipv4Address, Ipv6Address};
use wdk::rw_spin_lock::RwSpinLock;
use crate::driver_hashmap::DeviceHashMap;
#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Default)]
pub struct Key<Address>
where
Address: Eq + PartialEq,
{
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Default)]
pub struct Key<Address: Ord> {
pub local_ip: Address,
pub local_port: u16,
pub remote_ip: Address,
@@ -25,32 +21,32 @@ enum Direction {
Rx(usize),
}
pub struct Bandwidth {
stats_tcp_v4: DeviceHashMap<Key<Ipv4Address>, Value>,
stats_tcp_v4: BTreeMap<Key<Ipv4Address>, Value>,
stats_tcp_v4_lock: RwSpinLock,
stats_tcp_v6: DeviceHashMap<Key<Ipv6Address>, Value>,
stats_tcp_v6: BTreeMap<Key<Ipv6Address>, Value>,
stats_tcp_v6_lock: RwSpinLock,
stats_udp_v4: DeviceHashMap<Key<Ipv4Address>, Value>,
stats_udp_v4: BTreeMap<Key<Ipv4Address>, Value>,
stats_udp_v4_lock: RwSpinLock,
stats_udp_v6: DeviceHashMap<Key<Ipv6Address>, Value>,
stats_udp_v6: BTreeMap<Key<Ipv6Address>, Value>,
stats_udp_v6_lock: RwSpinLock,
}
impl Bandwidth {
pub fn new() -> Self {
Self {
stats_tcp_v4: DeviceHashMap::new(),
stats_tcp_v4: BTreeMap::new(),
stats_tcp_v4_lock: RwSpinLock::default(),
stats_tcp_v6: DeviceHashMap::new(),
stats_tcp_v6: BTreeMap::new(),
stats_tcp_v6_lock: RwSpinLock::default(),
stats_udp_v4: DeviceHashMap::new(),
stats_udp_v4: BTreeMap::new(),
stats_udp_v4_lock: RwSpinLock::default(),
stats_udp_v6: DeviceHashMap::new(),
stats_udp_v6: BTreeMap::new(),
stats_udp_v6_lock: RwSpinLock::default(),
}
}
@@ -62,7 +58,7 @@ impl Bandwidth {
if self.stats_tcp_v4.is_empty() {
return None;
}
stats_map = core::mem::replace(&mut self.stats_tcp_v4, DeviceHashMap::new());
stats_map = core::mem::replace(&mut self.stats_tcp_v4, BTreeMap::new());
}
let mut values = alloc::vec::Vec::with_capacity(stats_map.len());
@@ -89,7 +85,7 @@ impl Bandwidth {
if self.stats_tcp_v6.is_empty() {
return None;
}
stats_map = core::mem::replace(&mut self.stats_tcp_v6, DeviceHashMap::new());
stats_map = core::mem::replace(&mut self.stats_tcp_v6, BTreeMap::new());
}
let mut values = alloc::vec::Vec::with_capacity(stats_map.len());
@@ -116,7 +112,7 @@ impl Bandwidth {
if self.stats_udp_v4.is_empty() {
return None;
}
stats_map = core::mem::replace(&mut self.stats_udp_v4, DeviceHashMap::new());
stats_map = core::mem::replace(&mut self.stats_udp_v4, BTreeMap::new());
}
let mut values = alloc::vec::Vec::with_capacity(stats_map.len());
@@ -140,10 +136,10 @@ impl Bandwidth {
let stats_map;
{
let _guard = self.stats_udp_v6_lock.write_lock();
if self.stats_tcp_v6.is_empty() {
if self.stats_udp_v6.is_empty() {
return None;
}
stats_map = core::mem::replace(&mut self.stats_tcp_v6, DeviceHashMap::new());
stats_map = core::mem::replace(&mut self.stats_udp_v6, BTreeMap::new());
}
let mut values = alloc::vec::Vec::with_capacity(stats_map.len());
@@ -235,8 +231,8 @@ impl Bandwidth {
);
}
fn update<Address: Eq + PartialEq + core::hash::Hash>(
map: &mut DeviceHashMap<Key<Address>, Value>,
fn update<Address: Ord>(
map: &mut BTreeMap<Key<Address>, Value>,
lock: &mut RwSpinLock,
key: Key<Address>,
bytes: Direction,