command ack for ExitNode setting

Cargo.lock

@@ -2141,7 +2141,7 @@ dependencies = [
 
 [[package]]
 name = "punchnet"
-version = "1.0.3"
+version = "1.0.4"
 dependencies = [
  "ahash",
  "arc-swap",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "punchnet"
-version = "1.0.3"
+version = "1.0.4"
 edition = "2021"
 
 [dependencies]

proto/message.proto

@@ -107,12 +107,15 @@ message SDLPolicyResponse {
 
 // 事件定义
 
-message SDLNatChangedEvent {
+message SDLEvent {
+    // nat映射变化
+    message NatChanged {
         bytes mac = 1;
         uint32 ip = 2;
     }
 
-message SDLSendRegisterEvent {
+    // 发送register消息
+    message SendRegister {
         bytes dst_mac = 1;
         uint32 nat_ip = 2;
         uint32 nat_port = 3;
@@ -120,10 +123,39 @@ message SDLSendRegisterEvent {
         optional SDLV6Info v6_info = 5;
     }
 
-message SDLNetworkShutdownEvent {
+    // 网络关闭
+    message NetworkShutdown {
         string message = 1;
     }
 
+    oneof event {
+        NatChanged nat_changed = 1;
+        SendRegister send_register = 2;
+        NetworkShutdown shutdown = 3;
+    }
+}
+
+// Command指令
+message SDLCommand {
+    uint32 pkt_id = 1;
+    // 出口节点控制
+    message ExitNodeControl {
+        int32 action = 1;    // 必选：操作类型
+        string remark = 2;    // 可选：备注（方便日志/调试）
+    }
+
+    oneof command {
+        ExitNodeControl exit_node = 2;
+    }
+}
+
+message SDLCommandAck {
+    uint32 pkt_id = 1;
+    int32 code = 2;
+    string message = 3;
+    bytes data = 4;
+}
+
 // UDP通讯消息
 
 // client和stun之间的心跳包，客户端需要和super的udp之间的存活逻辑

src/network/mod.rs

@@ -25,4 +25,7 @@ pub const DNS_IP: u32 = (100<<24) + (100<<16) + (100<<8) + 100;
 mod tun;
 pub use tun::{get_install_channel, restore_dns, arp_reply_arrived};
 
+#[cfg(target_os = "linux")]
+pub use tun::{set_allow_routing, set_disallow_routing};
+
 mod device;

src/network/node.rs

@@ -739,7 +739,7 @@ pub struct NodeConfig {
     // node name
     pub name: String,
     // 允许路由
-    pub allow_routing: bool,
+    pub allow_routing: AtomicBool,
 
     // 丢弃多播，广播消息
     pub _drop_multicast: bool,

src/network/route.rs

@@ -10,7 +10,6 @@ use tracing::{debug, error};
 
 use crate::{RouteTableTrie, network::tun::{add_route, del_route}, pb::{SdlArpResponse, SdlStunReply}};
 
-
 pub struct RouteTable2 {
     pub cache_table: DashMap<(Ipv4Net, Ipv4Addr), AtomicBool, RandomState>,
     pub route_table: RouteTableTrie,

src/network/tun_linux.rs

@@ -170,7 +170,7 @@ impl Iface {
 
             node.route_table.apply_system();
 
-            if node.config.allow_routing {
+            if node.config.allow_routing.load(Ordering::Relaxed) {
                 set_allow_routing();
             }
 
@@ -933,8 +933,23 @@ pub fn add_route(net: &Ipv4Net, gw: &Ipv4Addr) -> Result<()> {
     Ok(())
 }
 
+pub fn set_disallow_routing() {
+    let _ = Command::new("sysctl")
+        .arg("-w")
+        .arg("net.ipv4.ip_forward=0")
+        .output();
 
-fn set_allow_routing() {
+    let _ = Command::new("iptables")
+        .arg("-t")
+        .arg("nat")
+        .arg("-D")
+        .arg("POSTROUTING")
+        .arg("-j")
+        .arg("MASQUERADE")
+        .output();
+}
+
+pub fn set_allow_routing() {
     let _ = Command::new("sysctl")
         .arg("-w")
         .arg("net.ipv4.ip_forward=1")

src/pb/message.rs

@@ -151,15 +151,25 @@ pub struct SdlPolicyResponse {
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
-pub struct SdlNatChangedEvent {
+pub struct SdlEvent {
+    #[prost(oneof = "sdl_event::Event", tags = "1, 2, 3")]
+    pub event: ::core::option::Option<sdl_event::Event>,
+}
+/// Nested message and enum types in `SDLEvent`.
+pub mod sdl_event {
+    /// nat映射变化
+    #[allow(clippy::derive_partial_eq_without_eq)]
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct NatChanged {
         #[prost(bytes = "vec", tag = "1")]
         pub mac: ::prost::alloc::vec::Vec<u8>,
         #[prost(uint32, tag = "2")]
         pub ip: u32,
     }
+    /// 发送register消息
     #[allow(clippy::derive_partial_eq_without_eq)]
     #[derive(Clone, PartialEq, ::prost::Message)]
-pub struct SdlSendRegisterEvent {
+    pub struct SendRegister {
         #[prost(bytes = "vec", tag = "1")]
         pub dst_mac: ::prost::alloc::vec::Vec<u8>,
         #[prost(uint32, tag = "2")]
@@ -169,13 +179,66 @@ pub struct SdlSendRegisterEvent {
         #[prost(uint32, tag = "4")]
         pub nat_type: u32,
         #[prost(message, optional, tag = "5")]
-    pub v6_info: ::core::option::Option<Sdlv6Info>,
+        pub v6_info: ::core::option::Option<super::Sdlv6Info>,
+    }
+    /// 网络关闭
+    #[allow(clippy::derive_partial_eq_without_eq)]
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct NetworkShutdown {
+        #[prost(string, tag = "1")]
+        pub message: ::prost::alloc::string::String,
+    }
+    #[allow(clippy::derive_partial_eq_without_eq)]
+    #[derive(Clone, PartialEq, ::prost::Oneof)]
+    pub enum Event {
+        #[prost(message, tag = "1")]
+        NatChanged(NatChanged),
+        #[prost(message, tag = "2")]
+        SendRegister(SendRegister),
+        #[prost(message, tag = "3")]
+        Shutdown(NetworkShutdown),
+    }
+}
+/// Command指令
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct SdlCommand {
+    #[prost(uint32, tag = "1")]
+    pub pkt_id: u32,
+    #[prost(oneof = "sdl_command::Command", tags = "2")]
+    pub command: ::core::option::Option<sdl_command::Command>,
+}
+/// Nested message and enum types in `SDLCommand`.
+pub mod sdl_command {
+    /// 出口节点控制
+    #[allow(clippy::derive_partial_eq_without_eq)]
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct ExitNodeControl {
+        /// 必选：操作类型
+        #[prost(int32, tag = "1")]
+        pub action: i32,
+        /// 可选：备注（方便日志/调试）
+        #[prost(string, tag = "2")]
+        pub remark: ::prost::alloc::string::String,
+    }
+    #[allow(clippy::derive_partial_eq_without_eq)]
+    #[derive(Clone, PartialEq, ::prost::Oneof)]
+    pub enum Command {
+        #[prost(message, tag = "2")]
+        ExitNode(ExitNodeControl),
+    }
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
-pub struct SdlNetworkShutdownEvent {
-    #[prost(string, tag = "1")]
+pub struct SdlCommandAck {
+    #[prost(uint32, tag = "1")]
+    pub pkt_id: u32,
+    #[prost(int32, tag = "2")]
+    pub code: i32,
+    #[prost(string, tag = "3")]
     pub message: ::prost::alloc::string::String,
+    #[prost(bytes = "vec", tag = "4")]
+    pub data: ::prost::alloc::vec::Vec<u8>,
 }
 /// client和stun之间的心跳包，客户端需要和super的udp之间的存活逻辑
 #[allow(clippy::derive_partial_eq_without_eq)]

src/tcp/identity_cache.rs

@@ -17,10 +17,6 @@ static RULE_CACHE: OnceLock<DashMap<IdentityID, (u64, HashMap<Port, HashSet<Prot
 
 // static RULE_CACHE: OnceLock<DashMap<IdentityID, HashMap<Port, HashMap<Proto, AtomicU64>>>> = OnceLock::new();
 
-pub fn init_identity_cache() {
-    RULE_CACHE.set(DashMap::new()).unwrap();
-}
-
 pub fn set_identity_cache(identity: IdentityID, infos: Vec<RuleInfo>) {
     debug!("setting identity cache for identity={}, infos: {:?}", identity, infos);
 

src/tcp/quic.rs

@@ -1,6 +1,7 @@
 use std::{net::IpAddr, sync::{Arc, OnceLock, atomic::{AtomicBool, AtomicU64, Ordering}}, time::Duration};
 
 use futures_util::pin_mut;
+use libc::PR_SET_THP_DISABLE;
 use prost::Message;
 use quinn::SendStream;
 use sdlan_sn_rs::{config::AF_INET, peer::{SdlanSock, V6Info}, utils::{Result, SDLanError, get_current_timestamp, ip_to_string, rsa_decrypt}};
@@ -9,7 +10,9 @@ 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, RuleFromServer, config::{NULL_MAC, TCP_PING_TIME}, get_edge, network::{ARP_REPLY, ArpHdr, EthHdr, Node, RegisterSuperFeedback, StartStopInfo, arp_reply_arrived, 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}};
+#[cfg(target_os = "linux")]
+use crate::network::{set_allow_routing, set_disallow_routing};
+use crate::{AesEncryptor, Chacha20Encryptor, ConnectionInfo, ConnectionState, MyEncryptor, RuleFromServer, config::{NULL_MAC, TCP_PING_TIME}, get_edge, network::{ARP_REPLY, ArpHdr, EthHdr, Node, RegisterSuperFeedback, StartStopInfo, arp_reply_arrived, check_peer_registration_needed, handle_packet_peer_info}, pb::{SdlArpResponse, SdlCommand, SdlCommandAck, SdlEvent, SdlPolicyResponse, SdlRegisterSuper, SdlRegisterSuperAck, SdlRegisterSuperNak, encode_to_tcp_message, sdl_command, sdl_event::{self, Event, SendRegister}}, tcp::{EventType, NakMsgCode, NatType, PacketType, SdlanTcp, read_a_packet, send_stun_request}};
 
 static GLOBAL_QUIC_HANDLE: OnceLock<ReadWriterHandle> = OnceLock::new();
 
@@ -316,18 +319,14 @@ async fn handle_tcp_message(msg: SdlanTcp) {
                 error!("malformed COMMAND received");
                 return;
             }
-            handle_tcp_command(edge, msg.current_packet[0], &msg.current_packet[1..]).await;
+            handle_tcp_command(edge,  &msg.current_packet[..]).await;
         }
         PacketType::Event => {
             if msg.current_packet.len() < 1 {
                 error!("malformed EVENT received");
                 return;
             }
-            let Ok(event) = msg.current_packet[0].try_into() else {
-                error!("failed to parse event type");
-                return;
-            };
-            handle_tcp_event(edge, event, &msg.current_packet[1..]).await;
+            handle_tcp_event(edge, &msg.current_packet[..]).await;
         }
         PacketType::PeerInfo => {
             let _ = handle_packet_peer_info(edge, &msg.current_packet[..]).await;
@@ -344,15 +343,69 @@ async fn handle_tcp_message(msg: SdlanTcp) {
 }
 
 
-async fn handle_tcp_command(_edge: &Node, _cmdtype: u8, _cmdprotobuf: &[u8]) {}
-
-async fn handle_tcp_event(edge: &'static Node, eventtype: EventType, eventprotobuf: &[u8]) {
-    match eventtype {
-        EventType::SendRegister => {
-            let Ok(reg) = SdlSendRegisterEvent::decode(eventprotobuf) else {
-                error!("failed to decode SendRegister Event");
+async fn handle_tcp_command(edge: &Node, cmdprotobuf: &[u8]) {
+    let Ok(cmd) = SdlCommand::decode(cmdprotobuf) else {
+        error!("failed to decode SdlCommand");
         return;
     };
+
+    let pkt_id = cmd.pkt_id;
+    let Some(command) = cmd.command else {
+        error!("command type is none");
+        return;
+    };
+
+    match command {
+        sdl_command::Command::ExitNode(node) => {
+            debug!("got exit node command: {:?}", node);
+            // println!("got exit node command: {:?}", node);
+            // std::process::exit(0);
+            if node.action == 0 {
+                // stop
+                let origin = edge.config.allow_routing.fetch_and(false, Ordering::Relaxed);
+                if origin {
+                    #[cfg(target_os = "linux")]
+                    set_disallow_routing();
+                }
+            } else {
+                // start
+                let origin = edge.config.allow_routing.fetch_or(true, Ordering::Relaxed);
+                if !origin {
+                    #[cfg(target_os = "linux")]
+                    set_allow_routing();
+                }
+            }
+
+            let ack = SdlCommandAck{
+                pkt_id,
+                code: 0,
+                message: "ok".to_owned(),
+                data: vec![],
+            };
+            let msg = encode_to_tcp_message(Some(ack), PacketType::CommandACK as u8).unwrap();
+
+            let stream = get_quic_write_conn();
+            if let Err(e) = stream.send(msg).await {
+                error!("failed to write command ack to quic: {}", e.as_str());
+            }
+        }
+    }
+
+}
+
+async fn handle_tcp_event(edge: &'static Node, eventprotobuf: &[u8]) {
+    let Ok(sdl_event) = SdlEvent::decode(eventprotobuf) else {
+        error!("failed to decode SdlEvent");
+        return;
+    };
+
+    let Some(event) = sdl_event.event else {
+        error!("event type is none");
+        return;
+    };
+
+    match event {
+        sdl_event::Event::SendRegister(reg) => {
             let v4 = reg.nat_ip.to_be_bytes();
             let mut v6_sock = None;
             if let Some(v6_info) = reg.v6_info {
@@ -391,10 +444,12 @@ async fn handle_tcp_event(edge: &'static Node, eventtype: EventType, eventprotob
             )
             .await;
         }
+
         other => {
             debug!("unhandled event {:?}", other);
         }
     }
+
 }