From 74a549dea90a4ee01f47b09719639130dccf2bb4 Mon Sep 17 00:00:00 2001 From: alex Date: Thu, 2 Apr 2026 16:45:31 +0800 Subject: [PATCH] command ack for ExitNode setting --- Cargo.lock | 2 +- Cargo.toml | 2 +- proto/message.proto | 54 ++++++++++++++++---- src/network/mod.rs | 3 ++ src/network/node.rs | 2 +- src/network/route.rs | 1 - src/network/tun_linux.rs | 19 ++++++- src/pb/message.rs | 105 ++++++++++++++++++++++++++++++-------- src/tcp/identity_cache.rs | 4 -- src/tcp/quic.rs | 83 +++++++++++++++++++++++++----- 10 files changed, 219 insertions(+), 56 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0d01f04..539380e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2141,7 +2141,7 @@ dependencies = [ [[package]] name = "punchnet" -version = "1.0.3" +version = "1.0.4" dependencies = [ "ahash", "arc-swap", diff --git a/Cargo.toml b/Cargo.toml index 5943295..e4897ce 100755 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "punchnet" -version = "1.0.3" +version = "1.0.4" edition = "2021" [dependencies] diff --git a/proto/message.proto b/proto/message.proto index 8f10ab5..5121ee4 100644 --- a/proto/message.proto +++ b/proto/message.proto @@ -107,21 +107,53 @@ message SDLPolicyResponse { // 事件定义 -message SDLNatChangedEvent { - bytes mac = 1; - uint32 ip = 2; +message SDLEvent { + // nat映射变化 + message NatChanged { + bytes mac = 1; + uint32 ip = 2; + } + + // 发送register消息 + message SendRegister { + bytes dst_mac = 1; + uint32 nat_ip = 2; + uint32 nat_port = 3; + uint32 nat_type = 4; + optional SDLV6Info v6_info = 5; + } + + // 网络关闭 + message NetworkShutdown { + string message = 1; + } + + oneof event { + NatChanged nat_changed = 1; + SendRegister send_register = 2; + NetworkShutdown shutdown = 3; + } } -message SDLSendRegisterEvent { - bytes dst_mac = 1; - uint32 nat_ip = 2; - uint32 nat_port = 3; - uint32 nat_type = 4; - optional SDLV6Info v6_info = 5; +// Command指令 +message SDLCommand { + uint32 pkt_id = 1; + // 出口节点控制 + message ExitNodeControl { + int32 action = 1; // 必选：操作类型 + string remark = 2; // 可选：备注（方便日志/调试） + } + + oneof command { + ExitNodeControl exit_node = 2; + } } -message SDLNetworkShutdownEvent { - string message = 1; +message SDLCommandAck { + uint32 pkt_id = 1; + int32 code = 2; + string message = 3; + bytes data = 4; } // UDP通讯消息 diff --git a/src/network/mod.rs b/src/network/mod.rs index c213354..0b38f2c 100755 --- a/src/network/mod.rs +++ b/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; diff --git a/src/network/node.rs b/src/network/node.rs index 8626f3f..9697adf 100755 --- a/src/network/node.rs +++ b/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, diff --git a/src/network/route.rs b/src/network/route.rs index c2861ee..a244f26 100755 --- a/src/network/route.rs +++ b/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, diff --git a/src/network/tun_linux.rs b/src/network/tun_linux.rs index 76e6cde..8c888cc 100755 --- a/src/network/tun_linux.rs +++ b/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") diff --git a/src/pb/message.rs b/src/pb/message.rs index 3c15ada..38918c4 100644 --- a/src/pb/message.rs +++ b/src/pb/message.rs @@ -151,31 +151,94 @@ pub struct SdlPolicyResponse { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] -pub struct SdlNatChangedEvent { - #[prost(bytes = "vec", tag = "1")] - pub mac: ::prost::alloc::vec::Vec, - #[prost(uint32, tag = "2")] - pub ip: u32, +pub struct SdlEvent { + #[prost(oneof = "sdl_event::Event", tags = "1, 2, 3")] + pub event: ::core::option::Option, +} +/// 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, + #[prost(uint32, tag = "2")] + pub ip: u32, + } + /// 发送register消息 + #[allow(clippy::derive_partial_eq_without_eq)] + #[derive(Clone, PartialEq, ::prost::Message)] + pub struct SendRegister { + #[prost(bytes = "vec", tag = "1")] + pub dst_mac: ::prost::alloc::vec::Vec, + #[prost(uint32, tag = "2")] + pub nat_ip: u32, + #[prost(uint32, tag = "3")] + pub nat_port: u32, + #[prost(uint32, tag = "4")] + pub nat_type: u32, + #[prost(message, optional, tag = "5")] + pub v6_info: ::core::option::Option, + } + /// 网络关闭 + #[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, +} +/// 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 SdlSendRegisterEvent { - #[prost(bytes = "vec", tag = "1")] - pub dst_mac: ::prost::alloc::vec::Vec, - #[prost(uint32, tag = "2")] - pub nat_ip: u32, - #[prost(uint32, tag = "3")] - pub nat_port: u32, - #[prost(uint32, tag = "4")] - pub nat_type: u32, - #[prost(message, optional, tag = "5")] - pub v6_info: ::core::option::Option, -} -#[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, } /// client和stun之间的心跳包，客户端需要和super的udp之间的存活逻辑 #[allow(clippy::derive_partial_eq_without_eq)] diff --git a/src/tcp/identity_cache.rs b/src/tcp/identity_cache.rs index fa78a4e..dbf8b5e 100644 --- a/src/tcp/identity_cache.rs +++ b/src/tcp/identity_cache.rs @@ -17,10 +17,6 @@ static RULE_CACHE: OnceLock>>> = OnceLock::new(); -pub fn init_identity_cache() { - RULE_CACHE.set(DashMap::new()).unwrap(); -} - pub fn set_identity_cache(identity: IdentityID, infos: Vec) { debug!("setting identity cache for identity={}, infos: {:?}", identity, infos); diff --git a/src/tcp/quic.rs b/src/tcp/quic.rs index 81bec2f..0a2afe6 100644 --- a/src/tcp/quic.rs +++ b/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 = 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_command(edge: &Node, cmdprotobuf: &[u8]) { + let Ok(cmd) = SdlCommand::decode(cmdprotobuf) else { + error!("failed to decode SdlCommand"); + return; + }; -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"); - 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); } } + }