diff --git a/message.proto b/message.proto index 6fe18eb..b0cdc53 100644 --- a/message.proto +++ b/message.proto @@ -55,8 +55,12 @@ message SDLRegisterSuper { // 部分逻辑会脱离quic去通讯,增加session_token校验 message SDLRegisterSuperAck { uint32 pkt_id = 1; - bytes aes_key = 2; - bytes session_token = 3; + // 目前支持aes, chacha20 + string algorithm = 2; + bytes key = 3; + // 逻辑分段,chacha20加密算法需要使用该字段 + uint32 region_id = 4; + bytes session_token = 5; } message SDLRegisterSuperNak { diff --git a/src/network/async_main.rs b/src/network/async_main.rs index e8cbbc5..fabd9c7 100755 --- a/src/network/async_main.rs +++ b/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) { async fn edge_send_packet_to_net(eee: &Node, data: Vec) { // 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()); diff --git a/src/network/node.rs b/src/network/node.rs index 668f817..1a76743 100755 --- a/src/network/node.rs +++ b/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, + pub network_id: AtomicU32, pub network_domain: RwLock, @@ -192,7 +194,7 @@ pub struct Node { // authorize related pub authorized: AtomicBool, // pub header_key: RwLock>>, - pub encrypt_key: RwLock>>, + // pub encrypt_key: RwLock>>, 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) { + 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> { self.encrypt_key.read().unwrap().clone() } + */ /* pub fn sn_is_known(&self, sock: &SdlanSock) -> bool { diff --git a/src/network/packet.rs b/src/network/packet.rs index 23598e8..2f8c02e 100755 --- a/src/network/packet.rs +++ b/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(); - if key.len() == 0 { + //let key = eee.get_encrypt_key(); + + // if key.len() == 0 { // check the encrypt key - error!("packet encrypt key not provided"); - return; - } + // error!("packet encrypt key not provided"); + // 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()); diff --git a/src/network/tun_linux.rs b/src/network/tun_linux.rs index 83cf0d9..2ed55bd 100755 --- a/src/network/tun_linux.rs +++ b/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, - 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(()); }; diff --git a/src/network/tuntap.rs b/src/network/tuntap.rs index ee5aba9..0bd42da 100755 --- a/src/network/tuntap.rs +++ b/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_device(&self, data: Vec, 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) -> std::io::Result<()>; } static ARP_WAIT_LIST: OnceCell = 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; }; diff --git a/src/pb/message.rs b/src/pb/message.rs index 65c4587..e3c188a 100644 --- a/src/pb/message.rs +++ b/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")] - pub aes_key: ::prost::alloc::vec::Vec, + /// 目前支持aes, chacha20 + #[prost(string, tag = "2")] + pub algorithm: ::prost::alloc::string::String, #[prost(bytes = "vec", tag = "3")] + pub key: ::prost::alloc::vec::Vec, + /// 逻辑分段,chacha20加密算法需要使用该字段 + #[prost(uint32, tag = "4")] + pub region_id: u32, + #[prost(bytes = "vec", tag = "5")] pub session_token: ::prost::alloc::vec::Vec, } #[allow(clippy::derive_partial_eq_without_eq)] diff --git a/src/tcp/quic.rs b/src/tcp/quic.rs index be50d1d..2f54b80 100644 --- a/src/tcp/quic.rs +++ b/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 = 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, stream: &mut SendStream, pkt_id: Option) { diff --git a/src/utils/encrypter.rs b/src/utils/encrypter.rs new file mode 100644 index 0000000..f86b26e --- /dev/null +++ b/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); + fn encrypt(&self, data: &[u8]) -> Result>; + fn decrypt(&self, ciphered: &[u8]) -> Result>; +} + +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) { + 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> { + 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> { + 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, + is_setted: bool, + next_counter: AtomicU32, + region_id: u32, +} + +impl Chacha20Encryptor { + pub fn new(key: Vec, 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) { + self.key = key; + self.region_id = region_id; + } + + fn encrypt(&self, data: &[u8]) -> Result> { + 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(®ion_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> { + 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, + is_setted: bool, +} + +impl AesEncryptor { + pub fn new(key: Vec) -> Self { + Self { + key: Vec::new(), + is_setted: true, + } + } +} + +impl Encryptor for AesEncryptor { + fn decrypt(&self, ciphered: &[u8]) -> Result> { + aes_decrypt(&self.key, ciphered) + } + + fn encrypt(&self, data: &[u8]) -> Result> { + aes_encrypt(&self.key, data) + } + + fn is_setted(&self) -> bool { + self.is_setted + } + + fn set_key(&mut self, _region_id: u32, key:Vec) { + self.key = key; + self.is_setted = true; + } +} + diff --git a/src/utils/mod.rs b/src/utils/mod.rs index fb60be9..8aa2372 100755 --- a/src/utils/mod.rs +++ b/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;