diff --git a/.vscode/settings.json b/.vscode/settings.json index c4c8773..01706d9 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,5 +1,5 @@ { // "rust-analyzer.cargo.target": "x86_64-pc-windows-gnu", - "rust-analyzer.cargo.target": "x86_64-unknown-linux-gnu", + // "rust-analyzer.cargo.target": "x86_64-unknown-linux-gnu", // "rust-analyzer.cargo.features": ["tun"] } \ No newline at end of file diff --git a/src/network/tun_win.rs b/src/network/tun_win.rs index b19f3f9..0cc400d 100755 --- a/src/network/tun_win.rs +++ b/src/network/tun_win.rs @@ -495,7 +495,17 @@ impl TunTapPacketHandler for Iface { if let Some(transport) = sliced_packet.transport { match transport { TransportSlice::Tcp(tcp) => { - // just fall back + use crate::FiveTuple; + use std::net::IpAddr; + + let out_five_tuple = FiveTuple { + src_ip: IpAddr::V4(ipv4.header().source_addr()), + dst_ip: IpAddr::V4(ipv4.header().destination_addr()), + src_port: tcp.source_port(), + dst_port: tcp.destination_port(), + proto: etherparse::IpNumber::TCP.0, + }; + eee.rule_cache.touch_packet(out_five_tuple); } TransportSlice::Udp(udp) => { if dstip == DNS_IP { @@ -512,6 +522,18 @@ impl TunTapPacketHandler for Iface { .await; // edge.udp_sock_for_dns.send_to() return Ok(()); + } else { + use crate::FiveTuple; + use std::net::IpAddr; + + let out_five_tuple = FiveTuple { + src_ip: IpAddr::V4(ipv4.header().source_addr()), + dst_ip: IpAddr::V4(ipv4.header().destination_addr()), + src_port: udp.source_port(), + dst_port: udp.destination_port(), + proto: etherparse::IpNumber::UDP.0, + }; + eee.rule_cache.touch_packet(out_five_tuple); } } _other => {} diff --git a/src/utils/acl_session.rs b/src/utils/acl_session.rs index 1fde0e3..20b7d83 100644 --- a/src/utils/acl_session.rs +++ b/src/utils/acl_session.rs @@ -1,4 +1,11 @@ -use std::{net::IpAddr, sync::{Arc, atomic::{AtomicU64, Ordering}}, time::{Duration, SystemTime, UNIX_EPOCH}}; +use std::{ + net::IpAddr, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; use ahash::RandomState; use dashmap::{DashMap, DashSet}; @@ -6,7 +13,7 @@ use tracing::{debug, error}; const RULE_VALID_TIME_IN_SECS: u64 = 60; -#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] pub struct FiveTuple { pub src_ip: IpAddr, pub dst_ip: IpAddr, @@ -55,8 +62,13 @@ impl SessionTable { if let Some(info) = self.table.get(&key) { info.last_active.store(now_secs(), Ordering::Relaxed); return; - } - self.table.insert(key, SessionInfo { last_active: AtomicU64::new(now_secs()) }); + } + self.table.insert( + key, + SessionInfo { + last_active: AtomicU64::new(now_secs()), + }, + ); } pub fn process_packet(&self, key: &FiveTuple) -> bool { @@ -71,9 +83,9 @@ impl SessionTable { pub fn retain(&self) { let now = now_secs(); debug!("retain session"); - self.table.retain(|_, info|{ + self.table.retain(|_, info| { let last = info.last_active.load(Ordering::Relaxed); - now-last < self.timeout_secs + now - last < self.timeout_secs }); } } @@ -87,8 +99,6 @@ type Proto = u8; type RuleInfo = (Port, Proto); - - #[derive(Debug)] pub struct RuleFromServer { pub proto: Proto, @@ -96,7 +106,7 @@ pub struct RuleFromServer { } pub struct RuleCache { - pub rule_info: DashMap), RandomState>, + pub rule_info: DashMap), RandomState>, pub rule_valid_secs: u64, pub session_table: Arc, } @@ -105,7 +115,7 @@ type ShouldRenew = bool; impl RuleCache { pub fn new() -> Self { - let session_table= Arc::new(SessionTable::new(300)); + let session_table = Arc::new(SessionTable::new(300)); let table_cleaner = Arc::clone(&session_table); tokio::spawn(async move { loop { @@ -121,7 +131,10 @@ impl RuleCache { } pub fn set_identity_cache(&self, identity: IdentityID, infos: Vec) { - debug!("setting identity cache for identity={}, infos: {:?}", identity, infos); + debug!( + "setting identity cache for identity={}, infos: {:?}", + identity, infos + ); let now = now_secs(); @@ -130,7 +143,8 @@ impl RuleCache { // let mut protomap = HashSet::new(); now_sets.insert((info.port, info.proto)); } - self.rule_info.insert(identity, (AtomicU64::new(now) , now_sets)); + self.rule_info + .insert(identity, (AtomicU64::new(now), now_sets)); } pub fn touch_packet(&self, info: FiveTuple) { @@ -138,7 +152,12 @@ impl RuleCache { self.session_table.add_session_info(info); } - pub fn is_identity_ok(&self, allow_routing: bool, identity: IdentityID, info: FiveTuple) -> (bool, ShouldRenew) { + pub fn is_identity_ok( + &self, + allow_routing: bool, + identity: IdentityID, + info: FiveTuple, + ) -> (bool, ShouldRenew) { // return (true, false); if allow_routing { return (true, false); @@ -162,8 +181,14 @@ impl RuleCache { } } - self.rule_info.insert(identity, (AtomicU64::new(now), DashSet::with_hasher(RandomState::new()))); + self.rule_info.insert( + identity, + ( + AtomicU64::new(now), + DashSet::with_hasher(RandomState::new()), + ), + ); (false, true) } -} \ No newline at end of file +}