chacha20 and aes is supported, need further test

message.proto

@@ -55,8 +55,12 @@ message SDLRegisterSuper {
 // 部分逻辑会脱离quic去通讯，增加session_token校验
 message SDLRegisterSuperAck {
     uint32 pkt_id = 1;
-    bytes aes_key = 2;
+    // 目前支持aes， chacha20
-    bytes session_token = 3;
+    string algorithm = 2;
+    bytes key = 3;
+    // 逻辑分段，chacha20加密算法需要使用该字段
+    uint32 region_id = 4;
+    bytes session_token = 5;
 }
 
 message SDLRegisterSuperNak {

src/network/async_main.rs

@@ -252,7 +252,7 @@ async fn loop_tap(eee: &'static Node, cancel: CancellationToken) {
                 packet.extend_from_slice(&reply);
 
                 /// TODO: check the packet should
-                if let Err(_e) = eee.device.handle_packet_from_net(&packet, &Vec::new()).await {
+                if let Err(_e) = eee.device.handle_packet_from_net(&packet).await {
                     error!("failed to write dns packet to device");
                 }
             }
@@ -311,14 +311,16 @@ async fn read_and_parse_tun_packet(eee: &'static Node, buf: Vec<u8>) {
 async fn edge_send_packet_to_net(eee: &Node, data: Vec<u8>) {
     // debug!("edge send packet to net({} bytes): {:?}", data.len(), data);
 
+    /* 
     let encrypt_key = eee.get_encrypt_key();
     if encrypt_key.len() == 0 {
         error!("drop tun packet due to encrypt key len is 0");
         return;
     }
+    */
     if let Err(e) = eee
         .device
-        .handle_packet_from_device(data, encrypt_key.as_slice())
+        .handle_packet_from_device(data)
         .await
     {
         error!("failed to handle packet from device: {}", e.to_string());

src/network/node.rs

@@ -13,7 +13,7 @@ use tokio::sync::oneshot;
 use tracing::{debug, error};
 
 use crate::quic::quic_init;
-use crate::{ConnectionInfo, get_base_dir};
+use crate::{ConnectionInfo, Encryptor, MyEncryptor, get_base_dir};
 use crate::pb::{
     encode_to_tcp_message, encode_to_udp_message, SdlEmpty, SdlStunProbe, SdlStunProbeReply,
 };
@@ -162,6 +162,8 @@ impl IdentityID {
 pub struct Node {
     packet_id: AtomicU32,
 
+    pub encryptor: RwLock<MyEncryptor>,
+
     pub network_id: AtomicU32,
     pub network_domain: RwLock<String>,
 
@@ -192,7 +194,7 @@ pub struct Node {
     // authorize related
     pub authorized: AtomicBool,
     // pub header_key: RwLock<Arc<Vec<u8>>>,
-    pub encrypt_key: RwLock<Arc<Vec<u8>>>,
+    // pub encrypt_key: RwLock<Arc<Vec<u8>>>,
 
     pub rsa_pubkey: String,
     pub rsa_private: RsaPrivateKey,
@@ -384,6 +386,8 @@ impl Node {
 
         Self {
             packet_id: AtomicU32::new(1),
+            encryptor: RwLock::new(MyEncryptor::new()),
+
             network_id: AtomicU32::new(0),
             hostname: RwLock::new(hostname),
 
@@ -411,7 +415,7 @@ impl Node {
             device: new_iface("dev", mode),
 
             authorized: AtomicBool::new(false),
-            encrypt_key: RwLock::new(Arc::new(Vec::new())),
+            // encrypt_key: RwLock::new(Arc::new(Vec::new())),
             // rsa_pubkey:
             rsa_pubkey: pubkey,
             rsa_private: private,
@@ -456,9 +460,9 @@ impl Node {
         self.authorized.load(Ordering::Relaxed)
     }
 
-    pub fn set_authorized(&self, authorized: bool, encrypt_key: Vec<u8>) {
+    pub fn set_authorized(&self, authorized: bool) {
         self.authorized.store(authorized, Ordering::Relaxed);
-        *(self.encrypt_key.write().unwrap()) = Arc::new(encrypt_key);
+        // *(self.encrypt_key.write().unwrap()) = Arc::new(encrypt_key);
     }
 
     /*
@@ -467,9 +471,11 @@ impl Node {
     }
     */
 
+    /* 
     pub fn get_encrypt_key(&self) -> Arc<Vec<u8>> {
         self.encrypt_key.read().unwrap().clone()
     }
+    */
 
     /*
     pub fn sn_is_known(&self, sock: &SdlanSock) -> bool {

src/network/packet.rs

@@ -883,15 +883,18 @@ async fn handle_tun_packet(
     pkt: SdlData, //orig_sender: &SdlanSock
 ) {
     let payload = pkt.data;
-    let key = eee.get_encrypt_key();
+    //let key = eee.get_encrypt_key();
-    if key.len() == 0 {
+
+    // if key.len() == 0 {
         // check the encrypt key
-        error!("packet encrypt key not provided");
+        // error!("packet encrypt key not provided");
-        return;
+        // return;
-    }
+    // }
 
     // test_aes(key.as_slice());
-    let origin = aes_decrypt(key.as_slice(), &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");
         return;
@@ -947,7 +950,7 @@ async fn handle_tun_packet(
     debug!("sending packet to tun, {} bytes", data.len());
     if let Err(e) = eee
         .device
-        .handle_packet_from_net(&data, key.as_slice())
+        .handle_packet_from_net(&data)
         .await
     {
         error!("failed to handle packet from net: {}", e.to_string());

src/network/tun_linux.rs

@@ -197,7 +197,7 @@ impl Iface {
 
 #[cfg(not(feature = "tun"))]
 impl TunTapPacketHandler for Iface {
-    async fn handle_packet_from_net(&self, data: &[u8], _: &[u8]) -> std::io::Result<()> {
+    async fn handle_packet_from_net(&self, data: &[u8]) -> std::io::Result<()> {
         // debug!("in tap mode, got data: {:?}", data);
         match self.send(data) {
             Err(e) => {
@@ -210,7 +210,7 @@ impl TunTapPacketHandler for Iface {
     async fn handle_packet_from_device(
         &self,
         data: Vec<u8>,
-        encrypt_key: &[u8],
+        // encrypt_key: &[u8],
     ) -> std::io::Result<()> {
         use etherparse::PacketHeaders;
 
@@ -288,7 +288,8 @@ impl TunTapPacketHandler for Iface {
                 }
                 let size = data.len();
 
-                let Ok(encrypted) = aes_encrypt(encrypt_key, &data) else {
+                let Ok(encrypted) = edge.encryptor.read().unwrap().decrypt(&data) else {
+                // let Ok(encrypted) = aes_encrypt(encrypt_key, &data) else {
                     error!("failed to encrypt packet request");
                     return Ok(());
                 };

src/network/tuntap.rs

@@ -24,8 +24,8 @@ use super::get_edge;
 pub const MAX_WAIT_PACKETS: usize = 100;
 
 pub trait TunTapPacketHandler {
-    async fn handle_packet_from_net(&self, data: &[u8], key: &[u8]) -> std::io::Result<()>;
+    async fn handle_packet_from_net(&self, data: &[u8]) -> std::io::Result<()>;
-    async fn handle_packet_from_device(&self, data: Vec<u8>, key: &[u8]) -> std::io::Result<()>;
+    async fn handle_packet_from_device(&self, data: Vec<u8>) -> std::io::Result<()>;
 }
 
 static ARP_WAIT_LIST: OnceCell<ArpWaitList> = OnceCell::new();
@@ -75,7 +75,7 @@ impl ArpWaitList {
             return;
         }
 
-        let encrypt_key = edge.get_encrypt_key();
+        // let encrypt_key = edge.get_encrypt_key();
         let network_id = edge.network_id.load(Ordering::Relaxed);
 
         let src_mac = edge.device_config.get_mac();
@@ -88,7 +88,8 @@ impl ArpWaitList {
 
             let pkt_size = packet.len();
 
-            let Ok(encrypted) = aes_encrypt(&encrypt_key, &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/pb/message.rs

@@ -73,9 +73,15 @@ pub struct SdlRegisterSuper {
 pub struct SdlRegisterSuperAck {
     #[prost(uint32, tag = "1")]
     pub pkt_id: u32,
-    #[prost(bytes = "vec", tag = "2")]
+    /// 目前支持aes， chacha20
-    pub aes_key: ::prost::alloc::vec::Vec<u8>,
+    #[prost(string, tag = "2")]
+    pub algorithm: ::prost::alloc::string::String,
     #[prost(bytes = "vec", tag = "3")]
+    pub key: ::prost::alloc::vec::Vec<u8>,
+    /// 逻辑分段，chacha20加密算法需要使用该字段
+    #[prost(uint32, tag = "4")]
+    pub region_id: u32,
+    #[prost(bytes = "vec", tag = "5")]
     pub session_token: ::prost::alloc::vec::Vec<u8>,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]

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::{ConnectionInfo, ConnectionState, 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, 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}};
 
 static GLOBAL_QUIC_HANDLE: OnceLock<ReadWriterHandle> = OnceLock::new();
 
@@ -126,11 +126,23 @@ async fn handle_tcp_message(msg: SdlanTcp) {
             );
             debug!("got register super ack: {:?}", ack);
             edge.session_token.set(ack.session_token);
-            let Ok(aes) = rsa_decrypt(&edge.rsa_private, &ack.aes_key) else {
+            let Ok(key) = rsa_decrypt(&edge.rsa_private, &ack.key) else {
                 error!("failed to rsa decrypt aes key");
                 return;
             };
 
+            match ack.algorithm.to_ascii_lowercase().as_str() {
+                "chacha20" => {
+                    *edge.encryptor.write().unwrap() = MyEncryptor::ChaChao20(Chacha20Encryptor::new(key, ack.region_id));
+                }
+                "aes" => {
+                    *edge.encryptor.write().unwrap() = MyEncryptor::Aes(AesEncryptor::new(key));
+                }
+                _other => {
+
+                }
+            }
+
 
             /* 
             let Some(dev) = ack.dev_addr else {
@@ -170,7 +182,8 @@ async fn handle_tcp_message(msg: SdlanTcp) {
             // edge.device.reload_config(&edge.device_config, &dev.network_domain);
             edge.device.reload_config(&edge.device_config, &edge.network_domain.read().unwrap().clone());
 
-            edge.set_authorized(true, aes);
+            edge.set_authorized(true);
+
             send_stun_request(edge).await;
             tokio::spawn(async {
                 let nattype = edge.probe_nat_type().await;
@@ -316,7 +329,8 @@ async fn handle_tcp_message(msg: SdlanTcp) {
                 message: "failed to decode REGISTER SUPER NAK".to_owned(),
             });
             */
-            edge.set_authorized(false, Vec::new());
+            edge.set_authorized(false);
+            *edge.encryptor.write().unwrap() = MyEncryptor::Invalid;
             // std::process::exit(0);
         }
         PacketType::Command => {
@@ -674,7 +688,8 @@ impl ReadWriteActor {
 
 async fn on_disconnected_callback() {
     let edge = get_edge();
-    edge.set_authorized(false, vec![]);
+    edge.set_authorized(false);
+    *edge.encryptor.write().unwrap() = MyEncryptor::Invalid;
 }
 
 async fn on_connected_callback(local_ip: Option<IpAddr>, stream: &mut SendStream, pkt_id: Option<u32>) {

src/utils/encrypter.rs

@@ -0,0 +1,176 @@
+use std::{sync::{Arc, OnceLock, RwLock, atomic::{AtomicBool, AtomicU32, Ordering}}, time::{SystemTime, UNIX_EPOCH}};
+
+use chacha20poly1305::{KeyInit, aead::Aead};
+use dashmap::DashSet;
+use sdlan_sn_rs::utils::{Result, SDLanError, aes_decrypt, aes_encrypt};
+
+const CounterMask: u32 = (1<<24) - 1;
+
+pub trait Encryptor {
+    fn is_setted(&self) -> bool;
+    fn set_key(&mut self, region_id: u32, key:Vec<u8>);
+    fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>>;
+    fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>>;
+}
+
+pub enum MyEncryptor {
+    Invalid,
+    ChaChao20(Chacha20Encryptor),
+    Aes(AesEncryptor),
+}
+
+impl MyEncryptor {
+    pub fn new() -> Self {
+        Self::Invalid
+    }
+
+    pub fn is_setted(&self) -> bool {
+        match self {
+            Self::Invalid => false,
+            Self::Aes(aes) => {
+                aes.is_setted()
+            }
+            Self::ChaChao20(cha) => {
+                cha.is_setted()
+            }
+        }
+    }
+
+    pub fn set_key(&mut self, region_id: u32, key:Vec<u8>) {
+        match self {
+            Self::Invalid => {}
+            Self::Aes(aes) => {
+                aes.set_key(region_id, key);
+            }
+            Self::ChaChao20(cha) => {
+                cha.set_key(region_id, key);
+            }
+        }
+    }
+
+    pub fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> {
+        match self {
+            Self::Invalid => {
+                Err(SDLanError::EncryptError("invalid encryptor".to_owned()))
+            }
+            Self::Aes(aes) => {
+                aes.encrypt(data)
+            }
+            Self::ChaChao20(cha) => {
+                cha.encrypt(data)
+            }
+        }
+    }
+    pub fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>> {
+        match self {
+            Self::Invalid => {
+                Err(SDLanError::EncryptError("invalid encryptor".to_owned()))
+            }
+            Self::Aes(aes) => {
+                aes.decrypt(ciphered)
+            }
+            Self::ChaChao20(cha) => {
+                cha.decrypt(ciphered)
+            }
+        }
+    }
+}
+
+pub struct Chacha20Encryptor {
+    key: Vec<u8>,
+    is_setted: bool,
+    next_counter: AtomicU32,
+    region_id: u32,
+}
+
+impl Chacha20Encryptor {
+    pub fn new(key: Vec<u8>, region_id: u32) -> Self {
+        Self {
+            key,
+            is_setted: true,
+            next_counter: AtomicU32::new(0),
+            region_id,
+        }
+    }
+}
+
+impl Encryptor for Chacha20Encryptor {
+    fn set_key(&mut self, region_id: u32, key:Vec<u8>) {
+        self.key = key;
+        self.region_id = region_id;
+    }
+
+    fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> {
+        let cipher = chacha20poly1305::ChaCha20Poly1305::new(self.key.as_slice().into());
+        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64;
+
+        let next_counter = self.next_counter.fetch_update(Ordering::Release, Ordering::Acquire, |current| {
+            Some((current + 1) & CounterMask)
+        }).unwrap() as u64;
+
+        let mut nonce = Vec::new();
+        let region_id = self.region_id.to_be_bytes();
+        nonce.extend_from_slice(&region_id);
+        let next_data = (now<<24) | next_counter;
+        nonce.extend_from_slice(&next_data.to_be_bytes());
+
+        match cipher.encrypt(nonce.as_slice().into(), data) {
+            Ok(data) => Ok(data),
+            Err(e) => {
+                Err(SDLanError::EncryptError(e.to_string()))
+            }
+        }
+    }
+
+    fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>> {
+        if ciphered.len() < 12 {
+            return Err(SDLanError::EncryptError("ciphered text size error".to_owned()))
+        }
+        let cipher = chacha20poly1305::ChaCha20Poly1305::new(self.key.as_slice().into());
+        let nonce = &ciphered[0..11];
+        match cipher.decrypt(nonce.into(), &ciphered[12..]) {
+            Ok(data) => Ok(data),
+            Err(e) => {
+                Err(SDLanError::EncryptError(e.to_string()))
+            }
+        }
+    }
+
+    fn is_setted(&self) -> bool {
+        self.is_setted
+    }
+}
+
+pub struct AesEncryptor {
+    key: Vec<u8>,
+    is_setted: bool,
+}
+
+impl AesEncryptor {
+    pub fn new(key: Vec<u8>) -> Self {
+        Self {
+            key: Vec::new(),
+            is_setted: true,
+        }
+    }
+}
+
+impl Encryptor for AesEncryptor {
+    fn decrypt(&self, ciphered: &[u8]) -> Result<Vec<u8>> {
+        aes_decrypt(&self.key, ciphered)
+    }
+
+    fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> {
+        aes_encrypt(&self.key, data)
+    }
+
+    fn is_setted(&self) -> bool {
+        self.is_setted
+    }
+
+    fn set_key(&mut self, _region_id: u32, key:Vec<u8>) {
+        self.key = key;
+        self.is_setted = true;
+    }
+}
+

src/utils/mod.rs

@@ -1,7 +1,9 @@
 mod command;
+mod encrypter;
 
 use std::{fs::OpenOptions, io::Write, net::Ipv4Addr, path::Path};
 
+pub use encrypter::*;
 pub use command::*;
 
 mod socks;