merge cookie_match and packet_id_match to one

Cargo.lock

@@ -155,6 +155,15 @@ version = "0.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d67af77d68a931ecd5cbd8a3b5987d63a1d1d1278f7f6a60ae33db485cdebb69"
 
+[[package]]
+name = "arc-swap"
+version = "1.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a07d1f37ff60921c83bdfc7407723bdefe89b44b98a9b772f225c8f9d67141a6"
+dependencies = [
+ "rustversion",
+]
+
 [[package]]
 name = "arrayvec"
 version = "0.7.6"
@@ -2134,6 +2143,8 @@ dependencies = [
 name = "punchnet"
 version = "1.0.3"
 dependencies = [
+ "ahash",
+ "arc-swap",
  "bytes",
  "cargo-deb",
  "chacha20poly1305",
@@ -2147,6 +2158,7 @@ dependencies = [
  "futures-util",
  "hex",
  "hmac",
+ "ipnet",
  "libc",
  "local-ip-address",
  "md-5",

Cargo.toml

@@ -39,6 +39,9 @@ chacha20poly1305 = "0.10.1"
 hmac = "0.12.1"
 md-5 = "0.10.6"
 hex = "0.4.3"
+ahash = "0.8.12"
+ipnet = "2.12.0"
+arc-swap = "1.9.0"
 # rolling-file = { path = "../rolling-file" }
 
 [target.'cfg(unix)'.dependencies]

Makefile

@@ -1,4 +1,4 @@
-self:
+self: libtun-so
 	RUSTFLAGS="-L ." cargo build --release
 
 linux: libtun-so

src/bin/punchnet/main.rs

@@ -117,7 +117,7 @@ async fn daemonize_me(
 
     let (tx, rx) = std::sync::mpsc::channel();
 
-    let hostname = "118.178.229.213".to_owned();
+    let hostname = "root.punchsky.com".to_owned();
     let host = format!("{}:80", hostname);
     let mut server = String::new();
     if let Ok(addrs) = host.to_socket_addrs() {
@@ -354,7 +354,7 @@ fn main() {
         }
     }
 
-    let should_daemonize = true;
+    let should_daemonize = false;
 
     #[cfg(not(target_os = "windows"))]
     if should_daemonize {
@@ -398,7 +398,7 @@ fn main() {
 fn run_it(cmd: CommandLineInput2, client_id: String, mac: Mac, system: &str, version: &str) {
     let rt = Runtime::new().unwrap();
     match &cmd.cmd {
-        Commands::Start => {
+        Commands::Start(rtinfo) => {
             rt.block_on(async move {
                 let remembered_token = get_access_token();
                 if remembered_token.is_none() {

src/network/arp.rs

@@ -15,11 +15,10 @@ use tokio::sync::{
     oneshot,
 };
 
-use super::{get_edge, get_route_table, init_arp_wait_list, init_route};
+use super::{get_edge, init_arp_wait_list};
 
 static GLOBAL_ARP: OnceCell<ArpActor> = OnceCell::new();
 pub fn init_arp() {
-    init_route();
     init_arp_wait_list();
     let actor = ArpActor::new();
     GLOBAL_ARP.set(actor).unwrap();
@@ -198,9 +197,10 @@ impl ArpInfo {
         }
 
         if target_ip == 0 {
-            let route_table = get_route_table();
-            if let Some(gateway_ip) = route_table.get_gateway_ip(ip) {
-                target_ip = gateway_ip;
+            // let route_table = get_route_table();
+            if let Some((_prefix, gateway_ip)) = edge.route_table.route_table.lookup(ip) {
+            // if let Some(gateway_ip) = route_table.get_gateway_ip(ip) {
+                target_ip = gateway_ip.into();
             }
         }
         if target_ip == 0 {

src/network/async_main.rs

@@ -7,7 +7,7 @@ use crate::network::ipv6::run_ipv6;
 use crate::network::{
     get_edge, ping_to_sn, read_and_parse_packet, TunTapPacketHandler,
 };
-use crate::tcp::{init_identity_cache, init_quic_conn, send_stun_request};
+use crate::tcp::{init_quic_conn, send_stun_request};
 use crate::utils::{send_to_sock, CommandLine};
 use crate::{ConnectionInfo};
 use sdlan_sn_rs::peer::{SdlanSock};
@@ -31,8 +31,6 @@ pub async fn async_main(
     // let _ = PidRecorder::new(".pid");
     let edge = get_edge();
 
-    init_identity_cache();
-
     // let token = args.token.clone();
     let cancel_tcp = cancel.clone();
     let (ipv6_network_restarter, rx) = channel(10);
@@ -252,7 +250,7 @@ async fn loop_tap(eee: &'static Node, cancel: CancellationToken) {
             }
         }
     }
-    debug!("loop_tap exited");
+    error!("loop_tap exited");
 }
 
 async fn get_tun_flow(eee: &'static Node, tx: Sender<Vec<u8>>) {

src/network/node.rs

@@ -1,8 +1,12 @@
+use arc_swap::ArcSwap;
 use dashmap::DashMap;
+use prost::Message;
 use quinn::Endpoint;
 use rsa::RsaPrivateKey;
 use sdlan_sn_rs::config::{AF_INET, AF_INET6};
 use tokio::net::UdpSocket;
+use std::any::Any;
+use std::future::Future;
 use std::net::SocketAddr;
 use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicU8, Ordering};
 use std::sync::{Arc, Mutex, RwLock};
@@ -12,8 +16,9 @@ use tokio::sync::mpsc::Sender;
 use tokio::sync::oneshot;
 use tracing::{debug, error};
 
+use crate::network::RouteTable2;
 use crate::quic::quic_init;
-use crate::{ConnectionInfo, MyEncryptor, get_base_dir};
+use crate::{ConnectionInfo, MyEncryptor, RuleCache, get_base_dir};
 use crate::pb::{
     encode_to_tcp_message, encode_to_udp_message, SdlEmpty, SdlStunProbe, SdlStunProbeReply,
 };
@@ -162,13 +167,16 @@ impl IdentityID {
 pub struct Node {
     packet_id: AtomicU32,
 
-    pub encryptor: RwLock<MyEncryptor>,
+    pub encryptor: ArcSwap<MyEncryptor>,
 
     pub network_id: AtomicU32,
     pub network_domain: RwLock<String>,
 
     pub identity_id: IdentityID,
 
+    pub rule_cache: RuleCache,
+    pub route_table: RouteTable2,
+
     pub access_token: StringToken<String>,
     pub session_token: StringToken<Vec<u8>>,
     
@@ -228,23 +236,18 @@ pub struct Node {
     nat_type: Mutex<NatType>,
 
     nat_cookie: AtomicU32,
-    cookie_match: DashMap<u32, oneshot::Sender<SdlStunProbeReply>>,
 
-    packet_id_match: DashMap<u32, oneshot::Sender<RegisterSuperFeedback>>,
+    //cookie_match: DashMap<u32, oneshot::Sender<SdlStunProbeReply>>,
+    cookie_match: Queryer,
+
+    // packet_id_match: DashMap<u32, oneshot::Sender<RegisterSuperFeedback>>,
 }
 
 unsafe impl Sync for Node {}
 
 impl Node {
     pub fn send_register_super_feedback(&self, pktid: u32, feed: RegisterSuperFeedback) {
-        match self.packet_id_match.remove(&pktid) {
-            Some(sender) => {
-                let _ = sender.1.send(feed);
-            }
-            None => {
-                return;
-            }
-        }
+        self.cookie_match.write_feedback(pktid, feed);
     }
 
     pub fn get_nat_type(&self) -> NatType {
@@ -320,30 +323,27 @@ impl Node {
         self.identity_id.store(identity_id);
         // *self._token.lock().unwrap() = token;
         // *self.network_code.lock().unwrap() = network_code;
-        let (tx, rx) = oneshot::channel();
         let id = self.get_next_packet_id();
-        self.packet_id_match.insert(id, tx);
-        let _ = self
-            .start_stop_sender
-            .send(StartStopInfo {
-                is_start: true,
-                pkt_id: Some(id),
-            })
-            .await;
-        debug!("start with feedback");
 
-        tokio::select! {
-            rx_info = rx => {
-                if let Ok(result) = rx_info {
-                    self.packet_id_match.remove(&id);
-                    Ok(result)
-                } else {
-                    Err(SDLanError::NormalError("rx closed"))
-                }
-            }
-            _ = tokio::time::sleep(timeout) => {
-                Err(SDLanError::NormalError("timed out"))
-            }
+        let res = self.cookie_match.do_action_and_wait_for(
+            id, 
+            || async {
+                let _ = self
+                    .start_stop_sender
+                    .send(StartStopInfo {
+                        is_start: true,
+                        pkt_id: Some(id),
+                    })
+                    .await;
+                debug!("start with feedback");
+            },
+            timeout
+        ).await?;
+
+        if let Ok(res) = res.downcast() {
+            Ok(*res)
+        } else {
+            Err(SDLanError::ConvertError("failed to convert feedback to RSFeedback".to_owned()))
         }
     }
 
@@ -386,11 +386,16 @@ impl Node {
 
         Self {
             packet_id: AtomicU32::new(1),
-            encryptor: RwLock::new(MyEncryptor::new()),
+            encryptor: ArcSwap::from(Arc::new(MyEncryptor::new())),
+            // encryptor: RwLock::new(MyEncryptor::new()),
 
             network_id: AtomicU32::new(0),
             hostname: RwLock::new(hostname),
 
+            rule_cache: RuleCache::new(),
+
+            route_table: RouteTable2::new(),
+
             network_domain: RwLock::new(String::new()),
 
             udp_sock_for_dns: udpsock_for_dns,
@@ -443,9 +448,9 @@ impl Node {
             stats: NodeStats::new(),
             _last_register_req: AtomicU64::new(0),
 
-            packet_id_match: DashMap::new(),
+            // packet_id_match: DashMap::new(),
             nat_cookie: AtomicU32::new(1),
-            cookie_match: DashMap::new(),
+            cookie_match: Queryer::new(),
             server_ip,
             install_channel,
         }
@@ -527,11 +532,7 @@ impl Node {
     */
 
     pub async fn send_nat_probe_reply(&self, cookie: u32, buf: SdlStunProbeReply) {
-        if let Some((_key, chan)) = self.cookie_match.remove(&cookie) {
-            let _ = chan.send(buf);
-            return;
-        }
-        error!("failed to get such cookie stun probe");
+        self.cookie_match.write_feedback(cookie, buf);
     }
 
     pub async fn probe_nat_type(&self) -> NatType {
@@ -614,33 +615,14 @@ impl Node {
             cookie,
             step: 0,
         };
+
+        let result = self.cookie_match.send_message_to_udp_and_wait_for(&self.udp_sock_v4, cookie, probe, PacketType::StunProbe as u8, to_server, Duration::from_secs(3)).await?;
+
+        if let Ok(res) = result.downcast() {
+            return Ok(*res);
+        }
+        return Err(SDLanError::ConvertError("failed to convert to StunprobeReply".to_owned()))
         // println!("==> sending probe request: {:?}", probe);
-
-        let (tx, rx) = oneshot::channel();
-        self.cookie_match.insert(cookie, tx);
-        // let cookie = msg.cookie;
-        let msg = encode_to_udp_message(Some(probe), PacketType::StunProbe as u8).unwrap();
-        if let Err(_e) = self.udp_sock_v4.send_to(&msg, to_server).await {
-            self.cookie_match.remove(&cookie);
-            return Err(SDLanError::NormalError("send error"));
-        }
-
-        tokio::select! {
-            _ = tokio::time::sleep(Duration::from_secs(3)) => {
-                self.cookie_match.remove(&cookie);
-                return Err(SDLanError::NormalError("timed out"));
-            }
-            reply = rx => {
-                self.cookie_match.remove(&cookie);
-                if let Ok(reply) = reply {
-                    // reply received,
-                    return Ok(reply);
-                    // println!("got nat ip: {}:{}", ip_to_string(&reply.ip), reply.port);
-                }
-                return Err(SDLanError::NormalError("reply recv error"));
-                // step 1 received
-            }
-        }
     }
 }
 
@@ -823,3 +805,107 @@ impl EdgePeer {
         self.nat_type
     }
 }
+
+type BoxedProstMessage = Box<dyn Any + Send + Sync + 'static>;
+
+pub struct Queryer {
+    pub mailbox: DashMap<u32, tokio::sync::oneshot::Sender<BoxedProstMessage>>,
+}
+
+impl Queryer {
+    pub fn new() -> Self {
+        Self {
+            mailbox: DashMap::new(),
+        }
+    }
+
+    pub fn write_feedback<T: Any + Sync + Send + 'static>(&self, id: u32, data: T) {
+        if let Some((_, tx)) = self.mailbox.remove(&id) {
+            if let Err(e) = tx.send(Box::new(data)) {
+                error!("failed to write feedback");
+            }
+        }
+    }
+
+    pub async fn send_message_to_udp_and_wait_for<T: Message>(&self, sock: &Socket, id: u32, message: T, packet_type: u8, to_server: &SocketAddr, timeout: Duration) -> Result<BoxedProstMessage> {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.mailbox.insert(id, tx);
+
+        let content = encode_to_udp_message(Some(message), packet_type)?;
+
+        if let Err(_e) = sock.send_to(&content, to_server).await {
+            self.mailbox.remove(&id);
+            return Err(SDLanError::NormalError("send error"));
+        }
+
+        tokio::select! {
+            data = rx => {
+                if let Ok(data) = data {
+                    self.mailbox.remove(&id);
+                    Ok(data)
+                } else {
+                    self.mailbox.remove(&id);
+                    Err(SDLanError::IOError("rx receive failed".to_string()))
+                }
+            }
+            _ = tokio::time::sleep(timeout) => {
+                self.mailbox.remove(&id);
+                Err(SDLanError::IOError("timed out".to_string()))
+            }
+        }
+
+    }
+
+    pub async fn do_action_and_wait_for<T, F>(&self, id: u32, action: T, timeout: Duration) -> Result<BoxedProstMessage> 
+    where 
+        F: Future<Output = ()>,
+        T: Fn() -> F,
+    {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.mailbox.insert(id, tx);
+
+        action().await;
+
+        tokio::select! {
+            data = rx => {
+                if let Ok(data) = data {
+                    self.mailbox.remove(&id);
+                    Ok(data)
+                } else {
+                    self.mailbox.remove(&id);
+                    Err(SDLanError::IOError("rx receive failed".to_string()))
+                }
+            }
+            _ = tokio::time::sleep(timeout) => {
+                self.mailbox.remove(&id);
+                Err(SDLanError::IOError("timed out".to_string()))
+            }
+        }
+
+    }
+
+    pub async fn send_message_to_quic_and_wait_for<T: Message>(&self, id: u32, message: T, packet_type: u8, timeout: Duration) -> Result<BoxedProstMessage> {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.mailbox.insert(id, tx);
+
+        let content = encode_to_tcp_message(Some(message), packet_type)?;
+        let quic_conn = get_quic_write_conn();
+        quic_conn.send(content).await?;
+
+        tokio::select! {
+            data = rx => {
+                if let Ok(data) = data {
+                    self.mailbox.remove(&id);
+                    Ok(data)
+                } else {
+                    self.mailbox.remove(&id);
+                    Err(SDLanError::IOError("rx receive failed".to_string()))
+                }
+            }
+            _ = tokio::time::sleep(timeout) => {
+                self.mailbox.remove(&id);
+                Err(SDLanError::IOError("timed out".to_string()))
+            }
+        }
+    }
+}

src/network/packet.rs

@@ -1,7 +1,8 @@
 use std::{net::SocketAddr, sync::atomic::Ordering, time::Duration};
 
+use crate::FiveTuple;
 use crate::pb::SdlPolicyRequest;
-use crate::tcp::{NatType, get_quic_write_conn, is_identity_ok};
+use crate::tcp::{NatType, get_quic_write_conn};
 use crate::{network::TunTapPacketHandler, utils::mac_to_string};
 
 use crate::{
@@ -13,7 +14,7 @@ use crate::{
     tcp::{PacketType},
     utils::{send_to_sock, Socket},
 };
-use etherparse::{Ethernet2Header, PacketHeaders, ip_number};
+use etherparse::{Ethernet2Header, IpNumber, PacketHeaders, ip_number};
 use prost::Message;
 use sdlan_sn_rs::utils::{BROADCAST_MAC};
 use sdlan_sn_rs::{
@@ -843,38 +844,6 @@ async fn renew_identity_request(eee: &Node, identity: u32) {
     }
 }
 
-async fn check_identity_is_ok(eee: &Node, identity: u32, protocol: u8, port: u16) -> bool{
-    true
-}
-
-async fn check_identity_is_ok2(eee: &Node, identity: u32, protocol: u8, port: u16) -> bool{
-    let result = is_identity_ok(identity, protocol, port);
-    if result.1 {
-        renew_identity_request(eee, identity).await;
-    }
-    match result.0 {
-        Some(true) => {
-            // identity is ok
-            true
-        }
-        Some(false) => {
-            // identity is not allowed
-            warn!("identity is not allowed for protocol={:?}, port={}", protocol, port);
-            false
-        }
-        None => {
-            if !result.1 {
-                renew_identity_request(eee, identity).await;
-            } else {
-                // has been sent
-            }
-            false
-            // no such identity, should request for it
-        }
-    }
-
-}
-
 async fn handle_tun_packet(
     eee: &Node,
     _from_sock: &SdlanSock,
@@ -892,7 +861,8 @@ async fn handle_tun_packet(
 
     // test_aes(key.as_slice());
 
-    let origin = eee.encryptor.read().unwrap().decrypt(&payload);
+    let origin = eee.encryptor.load().decrypt(&payload);
+    // let origin = eee.encryptor.read().unwrap().decrypt(&payload);
     // let origin = aes_decrypt(&payload);
     if let Err(_e) = origin {
         error!("failed to decrypt original data");
@@ -918,17 +888,35 @@ async fn handle_tun_packet(
                 match protocol {
                     ip_number::TCP => {
                         let tcp_header = headers.transport.unwrap().tcp().unwrap();
-                        let port = tcp_header.destination_port;
-                        let src_port = tcp_header.source_port;
-                        println!("tcp srcport={}, dstport={}", src_port, port);
-                        if !check_identity_is_ok(eee, pkt.identity_id, protocol.0, port).await {
+                        let five_tuple = FiveTuple {
+                            src_ip: ipv4.destination.into(),
+                            dst_ip: ipv4.source.into(),
+                            src_port: tcp_header.destination_port,
+                            dst_port: tcp_header.source_port,
+                            proto:IpNumber::TCP.0,
+                        };
+                        let (valid, need_refresh) = eee.rule_cache.is_identity_ok(pkt.identity_id, five_tuple);
+                        if need_refresh {
+                            renew_identity_request(eee, pkt.identity_id).await;
+                        }
+                        if !valid {
                             return;
                         }
                     }
                     ip_number::UDP => {
                         let udp_header = headers.transport.unwrap().udp().unwrap();
-                        let port = udp_header.destination_port;
-                        if !check_identity_is_ok(eee, pkt.identity_id, protocol.0, port).await {
+                        let five_tuple = FiveTuple {
+                            src_ip: ipv4.destination.into(),
+                            dst_ip: ipv4.source.into(),
+                            src_port: udp_header.destination_port,
+                            dst_port: udp_header.source_port,
+                            proto:IpNumber::UDP.0,
+                        };
+                        let (valid, need_refresh) = eee.rule_cache.is_identity_ok(pkt.identity_id, five_tuple);
+                        if need_refresh {
+                            renew_identity_request(eee, pkt.identity_id).await;
+                        }
+                        if !valid {
                             return;
                         }
                     }

src/network/route.rs

@@ -1,127 +1,83 @@
-use std::{net::Ipv4Addr, sync::RwLock};
+use std::{collections::HashMap, net::Ipv4Addr, sync::atomic::{AtomicBool, Ordering}, time::Duration};
 
-use once_cell::sync::OnceCell;
-use sdlan_sn_rs::utils::net_bit_len_to_mask;
+use ahash::RandomState;
+use dashmap::{DashMap};
+use ipnet::Ipv4Net;
+use sdlan_sn_rs::utils::{Result, SDLanError};
 
+use tokio::sync::oneshot::{Receiver, Sender, channel};
 use tracing::{debug, error};
 
-#[derive(Debug)]
-pub struct RouteTable {
-    content: RwLock<Vec<RouteInfo>>,
+use crate::{RouteTableTrie, network::tun::add_route, pb::{SdlArpResponse, SdlStunReply}};
+
+
+pub struct RouteTable2 {
+    pub cache_table: DashMap<(Ipv4Net, Ipv4Addr), AtomicBool, RandomState>,
+    pub route_table: RouteTableTrie,
 }
 
-static ROUTETABLE: OnceCell<RouteTable> = OnceCell::new();
-
-pub fn init_route() {
-    let rt = RouteTable::new();
-    ROUTETABLE.set(rt).unwrap();
-}
-
-pub fn get_route_table() -> &'static RouteTable {
-    ROUTETABLE.get().unwrap()
-}
-
-impl RouteTable {
+impl RouteTable2 {
     pub fn new() -> Self {
         Self {
-            content: RwLock::new(Vec::new()),
+            cache_table: DashMap::with_hasher(RandomState::new()),
+            route_table: RouteTableTrie::new(),
         }
     }
 
-    pub fn get_gateway_ip(&self, net_ip: u32) -> Option<u32> {
-        let routes = self.content.read().unwrap();
-        for route in &*routes {
-            debug!("route: {:?}", route.to_string());
-            if (route.net_ip & route.net_mask) == (net_ip & route.net_mask) {
-                // found
-                return Some(route.gateway_ip);
+    pub fn parse_and_add_route(&self, route_str: &str) -> Result<()> {
+        let routes = parse_route(route_str);
+        for route in routes.keys() {
+            if self.cache_table.get(route).is_some() {
+                error!("route {} {} has been added", route.0.to_string(), route.1);
+                return Err(SDLanError::IOError(format!("route {} already added", route.0.to_string())));
             }
         }
-        None
-    }
 
-    pub fn del_route(&self, net_ip: u32, net_mask: u32) {
-        let mut routes = self.content.write().unwrap();
-        let mut remove_idx = routes.len();
-        for i in 0..routes.len() {
-            let route = &routes[i];
-            if route.net_ip == net_ip && route.net_mask == net_mask {
-                remove_idx = i;
-                break;
-            }
+        for route in routes.keys() {
+            self.cache_table.insert(*route, AtomicBool::new(false));
+            self.route_table.insert(route.0.addr().into(), route.0.prefix_len(), route.1);
         }
-        if remove_idx < routes.len() {
-            routes.remove(remove_idx);
-        }
-    }
-
-    pub fn add_route(&self, net_ip: u32, net_mask: u32, gateway_ip: u32) -> Result<(), String> {
-        {
-            let cnt = self.content.read().unwrap();
-            let net = net_ip & net_mask;
-            for route in &*cnt {
-                if (route.net_ip & route.net_mask) == net {
-                    return Err("route exists".to_owned());
-                }
-            }
-        }
-        {
-            let mut routes = self.content.write().unwrap();
-            routes.push(RouteInfo {
-                net_ip,
-                net_mask,
-                gateway_ip,
-            })
-        }
-
         Ok(())
     }
-}
 
-#[derive(Debug)]
-pub struct RouteInfo {
-    pub net_ip: u32,
-    pub net_mask: u32,
-    pub gateway_ip: u32,
-}
-
-impl RouteInfo {
-    pub fn to_string(&self) -> String {
-        format!(
-            "{:?} mask={:?}, gateway={:?}",
-            self.net_ip.to_be_bytes(),
-            self.net_mask.to_be_bytes(),
-            self.gateway_ip.to_be_bytes()
-        )
+    pub fn apply_system(&self) {
+        for route in &self.cache_table {
+            let origin = route.fetch_or(true, Ordering::Relaxed);
+            if !origin {
+                // should add to system
+                add_route(route.key().0, route.key().1);
+            }
+        }
     }
 }
 
-// ip, mask, gateway, cidr;gateway,cidr2;gateway2
-pub fn parse_route(route: String) -> Vec<(u32, u32, u32)> {
-    let mut result = Vec::new();
-    let routes: Vec<_> = route.split(",").collect();
-    for route in routes {
-        let route_info: Vec<_> = route.split(";").collect();
-        debug!("got route info: {:?}", route_info);
-        if route_info.len() != 2 {
+// ip, mask, gateway, cidr gateway,cidr2 gateway2
+pub fn parse_route(route: &str) -> HashMap<(Ipv4Net, Ipv4Addr), bool> {
+    let mut result = HashMap::new();
+    // let routes: Vec<_> = route.split(",").collect();
 
-            error!("route info format error");
+    for route in route.trim().split(",") {
+        let route_info: Vec<_> = route.trim().split_whitespace().collect();
+        if route_info.len() != 2 {
+            error!("route info format error: {}", route);
             continue;
         }
+        debug!("got route info: {:?}", route_info);
+
+        let Ok(gateway) = route_info[1].parse::<Ipv4Addr>() else {
+            error!("failed to parse gw: {}", route_info[1]);
+            continue;
+        };
+
         let cidr = route_info[0];
-        let gateway = route_info[1].parse::<Ipv4Addr>().unwrap();
-        let ip_and_mask: Vec<_> = cidr.split("/").collect();
-        if ip_and_mask.len() != 2 {
-            error!("route info ip/bit error");
+        let Ok(net )= cidr.parse::<Ipv4Net>()  else {
+            error!("failed to parse cidr: {}, skipping", cidr);
             continue;
+        };
+        let origin = result.insert((net, gateway), true);
+        if origin.is_some() {
+            error!("{} {} already added", net.to_string(), gateway.to_string());
         }
-        let ip = ip_and_mask[0].parse::<Ipv4Addr>().unwrap();
-        let maskbit = ip_and_mask[1].parse::<u8>().unwrap();
-        result.push((
-            u32::from_be_bytes(ip.octets()),
-            net_bit_len_to_mask(maskbit),
-            u32::from_be_bytes(gateway.octets()),
-        ));
     }
     result
 }

src/network/tun_linux.rs

@@ -1,4 +1,5 @@
 use etherparse::{Ethernet2Header};
+use ipnet::Ipv4Net;
 use sdlan_sn_rs::config::SDLAN_DEFAULT_TTL;
 use sdlan_sn_rs::utils::{
     ip_to_string, is_ipv6_multicast, net_bit_len_to_mask,
@@ -7,6 +8,7 @@ use sdlan_sn_rs::utils::{
 use std::ffi::CStr;
 use std::ffi::{c_char, c_int};
 use std::fs::{self, OpenOptions};
+use std::net::Ipv4Addr;
 use std::os::unix::fs::{MetadataExt, PermissionsExt};
 use std::path::Path;
 use std::ptr::null_mut;
@@ -260,10 +262,40 @@ impl TunTapPacketHandler for Iface {
                         etherparse::NetHeaders::Ipv4(ipv4, _) => {
                             use etherparse::ip_number::{self, ICMP};
 
+                            use crate::FiveTuple;
+                            use etherparse::IpNumber;
                             
                             if let Some(transport) = headers.transport {
-                                if let Some(tcp) = transport.tcp() {
+                                match ipv4.protocol {
+                                    IpNumber::TCP => {
+                                        let tcp = transport.tcp().unwrap();
+
+                                        let out_five_tuple = FiveTuple {
+                                            src_ip: ipv4.source.into(),
+                                            dst_ip: ipv4.destination.into(),
+                                            src_port: tcp.source_port,
+                                            dst_port: tcp.destination_port,
+                                            proto: IpNumber::TCP.0,
+                                        };
+
+                                        edge.rule_cache.touch_packet(out_five_tuple);
                                     // is tcp
+                                    }
+                                    IpNumber::UDP => {
+                                        let udp = transport.udp().unwrap();
+
+                                        let out_five_tuple = FiveTuple {
+                                            src_ip: ipv4.source.into(),
+                                            dst_ip: ipv4.destination.into(),
+                                            src_port: udp.source_port,
+                                            dst_port: udp.destination_port,
+                                            proto: IpNumber::UDP.0,
+                                        };
+                                        edge.rule_cache.touch_packet(out_five_tuple);
+                                    }
+                                    _other => {
+
+                                    }
                                 }
                             }
 
@@ -289,7 +321,8 @@ impl TunTapPacketHandler for Iface {
                 }
                 let size = data.len();
 
-                let Ok(encrypted) = edge.encryptor.read().unwrap().encrypt(&data) else {
+                let Ok(encrypted) = edge.encryptor.load().encrypt(&data) else {
+                // let Ok(encrypted) = edge.encryptor.read().unwrap().encrypt(&data) else {
                 // let Ok(encrypted) = aes_encrypt(encrypt_key, &data) else {
                     error!("failed to encrypt packet request");
                     return Ok(());
@@ -852,3 +885,19 @@ fn restore_resolv_conf() -> Result<()> {
 
     Ok(())
 }
+
+pub fn add_route(net: Ipv4Net, gw: Ipv4Addr) -> Result<()> {
+    let res = Command::new("route")
+        .arg("add")
+        .arg("-net")
+        .arg(net.to_string())
+        .arg("gw")
+        .arg(gw.to_string())
+        .output()?;
+
+    Ok(())
+}
+
+pub fn del_route() -> Result<()> {
+    Ok(())
+}

src/network/tuntap.rs

@@ -88,7 +88,8 @@ impl ArpWaitList {
 
             let pkt_size = packet.len();
 
-            let Ok(encrypted) = edge.encryptor.read().unwrap().encrypt(&packet) else {
+            let Ok(encrypted) = edge.encryptor.load().encrypt(&packet) else {
+            // let Ok(encrypted) = edge.encryptor.read().unwrap().encrypt(&packet) else {
             // let Ok(encrypted) = aes_encrypt(&encrypt_key, &packet) else {
                 error!("failed to encrypt packet request");
                 return;

src/tcp/mod.rs

@@ -2,11 +2,8 @@ mod tcp_codec;
 // mod tcp_conn;
 mod quic;
 
-mod identity_cache;
-
 pub use tcp_codec::*;
 pub use quic::*;
-pub use identity_cache::*;
 
 
 // pub use tcp_conn::*;

src/tcp/quic.rs

@@ -9,7 +9,7 @@ use tokio::{sync::mpsc::{Receiver, Sender, channel}, time::sleep};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, warn};
 
-use crate::{AesEncryptor, Chacha20Encryptor, ConnectionInfo, ConnectionState, MyEncryptor, config::{NULL_MAC, TCP_PING_TIME}, get_edge, network::{ARP_REPLY, ArpHdr, EthHdr, Node, RegisterSuperFeedback, StartStopInfo, check_peer_registration_needed, handle_packet_peer_info}, pb::{SdlArpResponse, SdlPolicyResponse, SdlRegisterSuper, SdlRegisterSuperAck, SdlRegisterSuperNak, SdlSendRegisterEvent, encode_to_tcp_message}, tcp::{EventType, NakMsgCode, NatType, PacketType, RuleInfo, SdlanTcp, read_a_packet, send_stun_request, set_identity_cache}};
+use crate::{AesEncryptor, Chacha20Encryptor, ConnectionInfo, ConnectionState, MyEncryptor, RuleFromServer, config::{NULL_MAC, TCP_PING_TIME}, get_edge, network::{ARP_REPLY, ArpHdr, EthHdr, Node, RegisterSuperFeedback, StartStopInfo, check_peer_registration_needed, handle_packet_peer_info}, pb::{SdlArpResponse, SdlPolicyResponse, SdlRegisterSuper, SdlRegisterSuperAck, SdlRegisterSuperNak, SdlSendRegisterEvent, encode_to_tcp_message}, tcp::{EventType, NakMsgCode, NatType, PacketType, SdlanTcp, read_a_packet, send_stun_request}};
 
 static GLOBAL_QUIC_HANDLE: OnceLock<ReadWriterHandle> = OnceLock::new();
 
@@ -36,7 +36,7 @@ impl ReadWriterHandle {
         Ok(())
     }
 
-    fn new<>(
+    fn new(
         cancel: CancellationToken,
         addr: &str,
         // on_connected: OnConnectedCallback<'a>,
@@ -78,8 +78,10 @@ impl ReadWriterHandle {
             loop {
                 if let Some(msg) = data_from_tcp.recv().await {
                     handle_tcp_message(msg).await;
+                    println!("handle_tcp_message ok");
                 } else {
                     error!("data from tcp exited");
+                    println!("data from tcp exited");
                     // eprintln!("data from tcp exited");
                     return;
                 }
@@ -109,6 +111,7 @@ async fn handle_tcp_message(msg: SdlanTcp) {
     // edge.tcp_pong.store(now, Ordering::Relaxed);
 
     debug!("got tcp message: {:?}", msg.packet_type);
+    println!("got tcp message: {:?}", msg.packet_type);
     match msg.packet_type {
         PacketType::RegisterSuperACK => {
             let Ok(ack) = SdlRegisterSuperAck::decode(&msg.current_packet[..]) else {
@@ -116,6 +119,7 @@ async fn handle_tcp_message(msg: SdlanTcp) {
                 return;
             };
 
+            debug!("got register super ack1: {:?}", ack);
             edge.send_register_super_feedback(
                 ack.pkt_id,
                 RegisterSuperFeedback {
@@ -133,10 +137,12 @@ async fn handle_tcp_message(msg: SdlanTcp) {
 
             match ack.algorithm.to_ascii_lowercase().as_str() {
                 "chacha20" => {
-                    *edge.encryptor.write().unwrap() = MyEncryptor::ChaChao20(Chacha20Encryptor::new(key, ack.region_id));
+                    edge.encryptor.store(Arc::new(MyEncryptor::ChaChao20(Chacha20Encryptor::new(key, ack.region_id))))
+                    // *edge.encryptor.write().unwrap() = MyEncryptor::ChaChao20(Chacha20Encryptor::new(key, ack.region_id));
                 }
                 "aes" => {
-                    *edge.encryptor.write().unwrap() = MyEncryptor::Aes(AesEncryptor::new(key));
+                    edge.encryptor.store(Arc::new(MyEncryptor::Aes(AesEncryptor::new(key))))
+                    // *edge.encryptor.write().unwrap() = MyEncryptor::Aes(AesEncryptor::new(key));
                 }
                 _other => {
 
@@ -190,12 +196,14 @@ async fn handle_tcp_message(msg: SdlanTcp) {
                 debug!("nat type is {:?}", nattype);
                 // println!("nat type is: {:?}", nattype);
             });
+            println!("register message handled");
         }
         PacketType::ArpResponse => {
             let Ok(resp) = SdlArpResponse::decode(&msg.current_packet[..]) else {
                 error!("failed to decode ARP RESPONSE");
                 return;
             };
+            debug!("got arp response: {:?}", resp);
             if resp.target_mac.len() != 6 {
                 // invalid target_mac
                 error!("invalid target_mac: {:?}, ip={}", resp.target_mac, ip_to_string(&resp.target_ip));
@@ -245,13 +253,13 @@ async fn handle_tcp_message(msg: SdlanTcp) {
                 let port = u16::from_be_bytes([policy.rules[start+1], policy.rules[start+2]]);
                 start += 3;
 
-                infos.push(RuleInfo{
+                infos.push(RuleFromServer{
                     proto,
                     port,
                 });
             }
 
-            set_identity_cache(identity, infos);
+            edge.rule_cache.set_identity_cache(identity, infos);
         }
 
         PacketType::RegisterSuperNAK => {
@@ -330,7 +338,8 @@ async fn handle_tcp_message(msg: SdlanTcp) {
             });
             */
             edge.set_authorized(false);
-            *edge.encryptor.write().unwrap() = MyEncryptor::Invalid;
+            edge.encryptor.store(Arc::new(MyEncryptor::Invalid));
+            // *edge.encryptor.write().unwrap() = MyEncryptor::Invalid;
             // std::process::exit(0);
         }
         PacketType::Command => {
@@ -515,6 +524,7 @@ impl ReadWriteActor {
                         }
                     } else {
                         // None, just return
+                        println!("start or stop is None");
                         return;
                     }
                 }
@@ -622,6 +632,7 @@ impl ReadWriteActor {
 
             let write_to_tcp = async {
                 while let Some(data) = to_tcp.recv().await {
+                    debug!("data size = {}", data.len());
                     match send.write(&data).await {
                         Ok(size) => {
                             debug!("{} bytes sent to tcp", size);
@@ -689,7 +700,7 @@ impl ReadWriteActor {
 async fn on_disconnected_callback() {
     let edge = get_edge();
     edge.set_authorized(false);
-    *edge.encryptor.write().unwrap() = MyEncryptor::Invalid;
+    edge.encryptor.store(Arc::new(MyEncryptor::Invalid));
 }
 
 async fn on_connected_callback(local_ip: Option<IpAddr>, stream: &mut SendStream, _pkt_id: Option<u32>) {

src/utils/acl_session.rs

@@ -0,0 +1,163 @@
+use std::{net::IpAddr, sync::{Arc, atomic::{AtomicU64, Ordering}}, time::{Duration, Instant, SystemTime, UNIX_EPOCH}};
+
+use ahash::RandomState;
+use dashmap::{DashMap, DashSet};
+use tracing::{debug, error};
+
+const RuleValidTimeInSecs: u64 = 60;
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] 
+pub struct FiveTuple {
+    pub src_ip: IpAddr,
+    pub dst_ip: IpAddr,
+    pub src_port: u16,
+    pub dst_port: u16,
+    pub proto: u8,
+}
+
+impl FiveTuple {
+    pub fn reverse(&self) -> Self {
+        Self {
+            src_ip: self.dst_ip,
+            dst_ip: self.src_ip,
+            dst_port: self.src_port,
+            src_port: self.dst_port,
+            proto: self.proto,
+        }
+    }
+}
+
+pub struct SessionInfo {
+    pub last_active: AtomicU64,
+}
+
+pub struct SessionTable {
+    table: Arc<DashMap<FiveTuple, SessionInfo, RandomState>>,
+    timeout_secs: u64,
+}
+
+fn now_secs() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs()
+}
+
+impl SessionTable {
+    pub fn new(timeout_secs: u64) -> Self {
+        Self {
+            table: Arc::new(DashMap::with_hasher(RandomState::new())),
+            timeout_secs,
+        }
+    }
+
+    pub fn add_session_info(&self, key: FiveTuple) {
+        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()) });
+    }
+
+    pub fn process_packet(&self, key: &FiveTuple) -> bool {
+        if let Some(session) = self.table.get_mut(&key) {
+            session.last_active.store(now_secs(), Ordering::Relaxed);
+            return true;
+        }
+
+        false
+    }
+
+    pub fn retain(&self) {
+        let now = now_secs();
+        debug!("retain session");
+        self.table.retain(|_, info|{
+            let last = info.last_active.load(Ordering::Relaxed);
+            now-last < self.timeout_secs
+        });
+    }
+}
+
+// ip's u32 representation
+type IdentityID = u32;
+// port representation
+type Port = u16;
+// proto, like icmp = 1, udp=16
+type Proto = u8;
+
+type RuleInfo = (Port, Proto);
+
+
+
+#[derive(Debug)]
+pub struct RuleFromServer {
+    pub proto: Proto,
+    pub port: Port,
+}
+
+pub struct RuleCache {
+    pub rule_info: DashMap<IdentityID, (AtomicU64, DashSet<RuleInfo,RandomState>), RandomState>,
+    pub rule_valid_secs: u64,
+    pub session_table: Arc<SessionTable>,
+}
+
+type ShouldRenew = bool;
+
+impl RuleCache {
+    pub fn new() -> Self {
+        let session_table= Arc::new(SessionTable::new(300));
+        let table_cleaner = Arc::clone(&session_table);
+        tokio::spawn(async move {
+            loop {
+                tokio::time::sleep(Duration::from_secs(20)).await;
+                table_cleaner.retain();
+            }
+        });
+        Self {
+            rule_info: DashMap::with_hasher(RandomState::new()),
+            rule_valid_secs: RuleValidTimeInSecs,
+            session_table,
+        }
+    }
+
+    pub fn set_identity_cache(&self, identity: IdentityID, infos: Vec<RuleFromServer>) {
+        debug!("setting identity cache for identity={}, infos: {:?}", identity, infos);
+
+        let now = now_secs();
+
+        let now_sets = DashSet::with_hasher(RandomState::new());
+        for info in &infos {
+            // let mut protomap = HashSet::new();
+            now_sets.insert((info.port, info.proto));
+        }
+        self.rule_info.insert(identity, (AtomicU64::new(now) , now_sets));
+    }
+
+    pub fn touch_packet(&self, info: FiveTuple) {
+        error!("touch a packet: {:?}", info);
+        self.session_table.add_session_info(info);
+    }
+
+    pub fn is_identity_ok(&self, identity: IdentityID, info: FiveTuple) -> (bool, ShouldRenew) {
+        error!("is identity ok? {:?}", info);
+        if self.session_table.process_packet(&info) {
+            error!("identity is ok");
+            return (true, false);
+        }
+
+        let now = now_secs();
+        if let Some(sets_info) = self.rule_info.get(&identity) {
+            let last_set_time = sets_info.0.load(Ordering::Relaxed);
+            if sets_info.1.contains(&(info.dst_port, info.proto)) {
+                if (now - last_set_time) > self.rule_valid_secs {
+                    return (true, true);
+                }
+                return (true, false);
+            }
+        }
+
+        self.rule_info.insert(identity, (AtomicU64::new(now), DashSet::with_hasher(RandomState::new())));
+
+        (false, true)
+    }
+}

src/utils/command.rs

@@ -22,14 +22,24 @@ pub enum Commands {
 
     /// after login, we can use start to 
     /// connect to the remote
-    Start,
+    Start(RouteCmdInfo),
 
     Info,
 
+    RouteAdd(RouteCmdInfo),
+    RouteDel(RouteCmdInfo),
+    RouteList,
+
     /// exits the 
     Stop,
 }
 
+#[derive(Args, Debug)]
+pub struct RouteCmdInfo {
+    #[arg(short, long, default_value="")]
+    pub route: String,
+}
+
 #[derive(Args, Debug)]
 pub struct UserLogin {
     #[arg(short, long, env = APP_USER_ENV_NAME)]
@@ -39,6 +49,14 @@ pub struct UserLogin {
     pub password: String,
 }
 
+#[derive(Args, Debug)]
+pub struct AutoRunTokenLogin {
+    #[arg(long, env=APP_TOKEN_ENV_NAME, required=false)]
+    pub token: String,
+
+    #[arg(short, long, default_value="")]
+    pub route: String,
+}
 
 #[derive(Args, Debug)]
 pub struct TokenLogin {

src/utils/mod.rs

@@ -1,10 +1,14 @@
 mod command;
+mod acl_session;
 mod encrypter;
+mod system_action;
 
 use std::{fs::OpenOptions, io::Write, net::Ipv4Addr, path::Path};
 
 pub use encrypter::*;
 pub use command::*;
+pub use acl_session::*;
+pub use system_action::*;
 
 mod socks;
 use rand::Rng;

src/utils/system_action.rs

@@ -0,0 +1,91 @@
+use std::{net::Ipv4Addr, sync::Arc};
+
+use arc_swap::ArcSwap;
+
+use crate::network::Node;
+
+#[derive(Default, Clone)]
+pub struct TrieNode {
+    child: [Option<Box<TrieNode>>; 2],
+    prefix_len: u8,
+    nexthop: Option<Ipv4Addr>,
+}
+
+#[derive(Default, Clone)]
+struct IpTrie {
+    root: TrieNode,
+}
+
+impl IpTrie {
+    fn new() -> Self {
+        Self {
+            root: TrieNode::default()
+        }
+    }
+
+    fn insert(&mut self, prefix: u32, prefix_len: u8, nexthop: Ipv4Addr) {
+        if prefix_len > 32 {
+            return;
+        }
+
+        let mut node = &mut self.root;
+        for i in 0..prefix_len {
+            let bit = ((prefix >> (31-i)) & 1) as usize;
+            node = node.child[bit].get_or_insert_with(|| Box::new(TrieNode::default()));
+
+        }
+        if prefix_len > node.prefix_len {
+            node.prefix_len = prefix_len;
+            node.nexthop = Some(nexthop);
+        }
+    }
+
+    fn lookup(&self, ip: u32) -> Option<(u8, Ipv4Addr)>{
+        let mut node = &self.root;
+        let mut best = None;
+
+        for i in 0..32 {
+            if node.nexthop.is_some() {
+                best = Some((node.prefix_len, node.nexthop.unwrap()));
+            }
+            let bit = ((ip>>(31-i)) & 1) as usize;
+            match &node.child[bit] {
+                Some(child) => {
+                    node = child;
+                }
+                None => {
+                    break;
+                }
+            }
+        }
+        if node.nexthop.is_some() {
+            best = Some((node.prefix_len, node.nexthop.unwrap()));
+        }
+        best
+    }
+}
+
+pub struct RouteTableTrie {
+    trie: ArcSwap<IpTrie>,
+}
+
+impl RouteTableTrie {
+    pub fn new() -> Self {
+        Self {
+            trie: ArcSwap::new(Arc::new(IpTrie::default()))
+        }
+    }
+
+    pub fn lookup(&self, ip: u32) -> Option<(u8, Ipv4Addr)> {
+        let trie = self.trie.load();
+        trie.lookup(ip)
+    }
+
+    pub fn insert(&self, prefix: u32, prefix_len: u8, nexthop: Ipv4Addr) {
+        let old = self.trie.load();
+        let mut new_trie = (*(*old)).clone();
+
+        new_trie.insert(prefix, prefix_len, nexthop);
+        self.trie.store(Arc::new(new_trie));
+    }
+}