base proxy ready

common/src/crypto/aead.rs

@@ -1,6 +1,6 @@
-use chacha20poly1305::aead::Buffer;
+use bytes::{Bytes, BytesMut};
 
 pub trait AeadPacker {
-    fn encrypt<B: Buffer>(&mut self, data: &mut B) -> Result<(), chacha20poly1305::aead::Error>;
-    fn decrypt<B: Buffer>(&mut self, data: &mut B) -> Result<(), chacha20poly1305::aead::Error>;
+    fn encrypt(&mut self, data: &mut BytesMut) -> Result<Bytes, chacha20poly1305::aead::Error>;
+    fn decrypt(&mut self, data: &mut BytesMut) -> Result<Bytes, chacha20poly1305::aead::Error>;
 }

common/src/crypto/chacha.rs

@@ -1,80 +1,139 @@
+use bytes::{Bytes, BytesMut};
 use chacha20poly1305::aead::generic_array::GenericArray;
-use chacha20poly1305::aead::{Buffer, OsRng};
-use chacha20poly1305::{
-    AeadCore, AeadInPlace, ChaCha20Poly1305, ChaChaPoly1305, Key, KeyInit, Nonce,
-};
+use chacha20poly1305::{AeadInPlace, ChaCha20Poly1305, Key, KeyInit, Nonce};
 
 use crate::crypto::aead::AeadPacker;
 
 pub struct NonceState {
     counter: u64,
-    nonce: Nonce,
-    handshake: bool,
+    base_iv: [u8; 12],
 }
 
 impl NonceState {
-    pub fn new() -> Self {
-        let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
+    pub fn new(base_iv: [u8; 12]) -> Self {
         Self {
             counter: 0,
-            nonce,
-            handshake: false,
+            base_iv,
         }
     }
 
-    pub fn get_handshake(&self) -> bool {
-        return self.handshake;
-    }
-
-    pub fn set_nonce(&mut self, nonce: Nonce) {
-        self.nonce = nonce;
-        self.handshake = true
-    }
-
-    pub fn increase_counter(&mut self) {
-        self.counter += 1;
+    // Возвращает nonce для текущего пакета и ПЕРЕХОДИТ к следующему
+    pub fn next_nonce(&mut self) -> Nonce {
+        let mut iv = self.base_iv;
+        // В TLS 1.3 используется Little Endian для счетчика при XOR
+        // Но если ты сам пишешь протокол, BE тоже пойдет.
+        // Главное — единообразие.
         let counter_bytes = self.counter.to_be_bytes();
-        self.nonce[4..12].copy_from_slice(&counter_bytes);
+
+        for i in 0..8 {
+            iv[i + 4] ^= counter_bytes[i];
+        }
+
+        self.counter += 1;
+        *GenericArray::from_slice(&iv)
     }
 }
 
 pub struct ChaChaCipher {
-    key: Key,
+    pub encrypt_cipher: ChaCha20Poly1305,
+    pub decrypt_cipher: ChaCha20Poly1305,
     pub encrypt_state: NonceState,
     pub decrypt_state: NonceState,
-    pub cipher: ChaCha20Poly1305,
 }
 
 impl ChaChaCipher {
     pub fn new() -> Self {
-        let key = GenericArray::clone_from_slice(&[0; 32]);
-        let cipher = ChaCha20Poly1305::new(&key);
+        let start_key = Key::from([0u8; 32]);
+        let encrypt_cipher = ChaCha20Poly1305::new(&start_key);
+        let decrypt_cipher = ChaCha20Poly1305::new(&start_key);
         Self {
-            key,
-            encrypt_state: NonceState::new(),
-            decrypt_state: NonceState::new(),
-            cipher,
+            encrypt_state: NonceState::new([0u8; 12]),
+            decrypt_state: NonceState::new([0u8; 12]),
+            encrypt_cipher,
+            decrypt_cipher,
         }
     }
 
-    pub fn set_key(&mut self, key: Key) -> () {
-        self.key = key;
-        self.cipher = ChaChaPoly1305::new(&self.key);
+    pub fn set_keys(
+        &mut self,
+        w_key: [u8; 32],
+        w_iv: [u8; 12], // Write (исходящие)
+        r_key: [u8; 32],
+        r_iv: [u8; 12], // Read (входящие)
+    ) {
+        self.encrypt_cipher = ChaCha20Poly1305::new(Key::from_slice(&w_key));
+        self.decrypt_cipher = ChaCha20Poly1305::new(Key::from_slice(&r_key));
+
+        self.encrypt_state = NonceState::new(w_iv);
+        self.decrypt_state = NonceState::new(r_iv);
+
+        tracing::debug!("Cipher keys and IVs updated for both directions");
     }
 }
-
 impl AeadPacker for ChaChaCipher {
-    fn encrypt<B: Buffer>(&mut self, data: &mut B) -> Result<(), chacha20poly1305::aead::Error> {
-        self.cipher
-            .encrypt_in_place(&self.encrypt_state.nonce, &[], data)?;
-        self.encrypt_state.increase_counter();
-        Ok(())
+    fn encrypt(&mut self, data: &mut BytesMut) -> Result<Bytes, chacha20poly1305::aead::Error> {
+        // Сначала получаем текущий counter для лога (до того, как next_nonce его инкрементирует)
+        let current_counter = self.encrypt_state.counter;
+        let nonce = self.encrypt_state.next_nonce();
+        let nonce_hex = hex::encode(nonce);
+        let data_len = data.len();
+
+        match self.encrypt_cipher.encrypt_in_place(&nonce, &[], data) {
+            Ok(_) => {
+                tracing::trace!(
+                    counter = current_counter,
+                    nonce = %nonce_hex,
+                    len = data_len,
+                    "Encryption successful"
+                );
+                Ok(data.split().freeze())
+            }
+            Err(e) => {
+                tracing::error!(
+                    counter = current_counter,
+                    nonce = %nonce_hex,
+                    len = data_len,
+                    error = ?e,
+                    "AEAD encryption failure"
+                );
+                Err(e)
+            }
+        }
     }
 
-    fn decrypt<B: Buffer>(&mut self, data: &mut B) -> Result<(), chacha20poly1305::aead::Error> {
-        self.cipher
-            .decrypt_in_place(&self.decrypt_state.nonce, &[], data)?;
-        self.decrypt_state.increase_counter();
-        Ok(())
+    fn decrypt(&mut self, data: &mut BytesMut) -> Result<Bytes, chacha20poly1305::aead::Error> {
+        let current_counter = self.decrypt_state.counter;
+        let nonce = self.decrypt_state.next_nonce();
+        let nonce_hex = hex::encode(nonce);
+        let data_len = data.len();
+
+        let data_prefix = if data.len() >= 8 {
+            hex::encode(&data[..8])
+        } else {
+            hex::encode(data.as_ref())
+        };
+
+        match self.decrypt_cipher.decrypt_in_place(&nonce, &[], data) {
+            Ok(_) => {
+                tracing::trace!(
+                    counter = current_counter,
+                    nonce = %nonce_hex,
+                    len = data_len,
+                    "Decryption successful"
+                );
+                Ok(data.split().freeze())
+            }
+            Err(e) => {
+                tracing::error!(
+                    counter = current_counter,
+                    nonce = %nonce_hex,
+                    len = data_len,
+                    prefix = %data_prefix,
+                    error = ?e,
+                    "AEAD decryption failure! Verification failed or data malformed"
+                );
+                Err(e)
+            }
+        }
     }
 }

common/src/crypto/mod.rs

@@ -1,6 +1,7 @@
 pub mod aead;
 pub mod chacha;
-pub mod ecdh;
-pub mod hkdf;
-pub mod hmac;
-pub mod salt_pair;
+mod ecdh;
+mod hkdf;
+mod hmac;
+mod salt_pair;
+pub mod session;

common/src/crypto/session.rs

@@ -0,0 +1,128 @@
+use x25519_dalek::PublicKey;
+
+use crate::{
+    crypto::{ecdh::ECDH, hkdf::HKDF, salt_pair::SaltPair},
+    tlseng::extension::ExtensionStack,
+};
+
+pub struct SessionKeys {
+    pub salt: SaltPair,
+    pub ecdh: ECDH,
+}
+
+impl SessionKeys {
+    pub fn new(is_initiator: bool) -> Self {
+        Self {
+            salt: SaltPair::new(is_initiator),
+            ecdh: ECDH::new(),
+        }
+    }
+
+    pub fn update_keys(
+        &mut self,
+        salt: [u8; 32],
+        extensions: &ExtensionStack,
+        is_server: bool, // true если мы Сервер (парсим ClientHello), false если Клиент
+    ) -> Result<([u8; 32], [u8; 12], [u8; 32], [u8; 12]), String> {
+        // Сохраняем соль от удаленной стороны
+        self.salt.set_remote_salt(salt);
+
+        tracing::debug!(
+            remote_salt = %hex::encode(&salt[..8]),
+            local_salt = %hex::encode(&self.salt.get_local()[..8]),
+            total_salt = %hex::encode(&self.salt.get_total()[28..36]),
+            "Updating keys with new salt"
+        );
+
+        const EXT_KEY_SHARE: u16 = 0x0033;
+
+        if let Some(dh_data) = extensions.find_by_type(EXT_KEY_SHARE) {
+            let mut key_bytes = [0u8; 32];
+
+            if is_server {
+                // МЫ СЕРВЕР: Парсим ClientHello KeyShare (там список)
+                // Минимум: 2 (длина списка) + 2 (группа) + 2 (длина ключа) + 32 (ключ) = 38
+                if dh_data.len() < 38 {
+                    return Err(format!(
+                        "Client KeyShare too short: {} bytes",
+                        dh_data.len()
+                    ));
+                }
+
+                let mut found = false;
+                // Ищем маркер x25519 (00 1d) и длину 32 (00 20)
+                for i in 2..=(dh_data.len() - 34) {
+                    if dh_data[i..i + 4] == [0x00, 0x1d, 0x00, 0x20] {
+                        key_bytes.copy_from_slice(&dh_data[i + 4..i + 36]);
+                        found = true;
+                        break;
+                    }
+                }
+
+                if !found {
+                    return Err("Could not find x25519 key in ClientHello".into());
+                }
+            } else {
+                // МЫ КЛИЕНТ: Парсим ServerHello KeyShare (там сразу группа и ключ)
+                // [Group:2] [KeyLen:2] [Key:32] = 36 байт
+                if dh_data.len() < 36 {
+                    return Err("Server KeyShare too short".into());
+                }
+                key_bytes.copy_from_slice(&dh_data[4..36]);
+            }
+
+            // Проверка на "кривой" ключ
+            if key_bytes.iter().all(|&x| x == 0) {
+                return Err("Extracted remote public key is all ZEROS!".into());
+            }
+
+            let public_key = PublicKey::from(key_bytes);
+
+            tracing::debug!(
+                remote_pub = %hex::encode(&key_bytes[..4]),
+                role = if is_server { "Server" } else { "Client" },
+                "Key exchange successful, deriving material..."
+            );
+
+            // Вызываем генерацию, которая вернет (w_key, w_iv, r_key, r_iv)
+            self.generate_keys(&public_key, is_server)
+        } else {
+            Err("No KeyShare extension found in handshake".into())
+        }
+    }
+    fn generate_keys(
+        &mut self,
+        public_key: &PublicKey,
+        is_server: bool,
+    ) -> Result<([u8; 32], [u8; 12], [u8; 32], [u8; 12]), String> {
+        let shared_key = self
+            .ecdh
+            .get_shared(public_key)
+            .ok_or_else(|| "No shared secret".to_string())?;
+
+        tracing::debug!(
+            shared_prefix = %hex::encode(&shared_key[..8]),
+            "DH Shared secret derived"
+        );
+
+        let hkdf = HKDF::extract_key(&self.salt.get_total(), &shared_key);
+
+        let c_key = HKDF::expand_key::<32>(&hkdf, b"client_aead").map_err(|e| e.to_string())?;
+        let c_iv = HKDF::expand_key::<12>(&hkdf, b"client_iv").map_err(|e| e.to_string())?;
+        let s_key = HKDF::expand_key::<32>(&hkdf, b"server_aead").map_err(|e| e.to_string())?;
+        let s_iv = HKDF::expand_key::<12>(&hkdf, b"server_iv").map_err(|e| e.to_string())?;
+
+        tracing::info!(
+            client_key_short = %hex::encode(&c_key[..4]),
+            server_key_short = %hex::encode(&s_key[..4]),
+            "HKDF expansion complete"
+        );
+
+        // Распределяем: (write_key, write_iv, read_key, read_iv)
+        if is_server {
+            Ok((s_key, s_iv, c_key, c_iv)) // Сервер пишет своим, читает клиентским
+        } else {
+            Ok((c_key, c_iv, s_key, s_iv)) // Клиент пишет своим, читает серверным
+        }
+    }
+}