Add secret updating after key exchange

Signed-off-by: trivernis <trivernis@protonmail.com>

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "vented"
 description = "Event driven encrypted tcp communicaton"
-version = "0.1.3"
+version = "0.2.0"
 authors = ["trivernis <trivernis@protonmail.com>"]
 edition = "2018"
 readme = "README.md"

src/crypto/mod.rs

@@ -1,6 +1,5 @@
 use std::io::{Read, Write};
 use std::net::TcpStream;
-use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 
 use byteorder::{BigEndian, ByteOrder};
@@ -13,6 +12,7 @@ use typenum::U24;
 use crate::event::Event;
 use crate::result::VentedResult;
 
+use crypto_box::ChaChaBox;
 pub use crypto_box::PublicKey;
 pub use crypto_box::SecretKey;
 
@@ -21,24 +21,28 @@ pub use crypto_box::SecretKey;
 pub struct CryptoStream {
     send_stream: Arc<Mutex<TcpStream>>,
     recv_stream: Arc<Mutex<TcpStream>>,
-    sent_count: Arc<AtomicUsize>,
-    recv_count: Arc<AtomicUsize>,
-    secret_box: Arc<Mutex<crypto_box::ChaChaBox>>,
+    send_secret: Arc<Mutex<EncryptionBox<ChaChaBox>>>,
+    recv_secret: Arc<Mutex<EncryptionBox<ChaChaBox>>>,
 }
 
 impl CryptoStream {
     /// Creates a new crypto stream from a given Tcp Stream and with a given secret
-    pub fn new(inner: TcpStream, secret_box: crypto_box::ChaChaBox) -> VentedResult<Self> {
+    pub fn new(
+        inner: TcpStream,
+        public_key: &PublicKey,
+        secret_key: &SecretKey,
+    ) -> VentedResult<Self> {
         inner.set_nonblocking(false)?;
         let send_stream = Arc::new(Mutex::new(inner.try_clone()?));
         let recv_stream = Arc::new(Mutex::new(inner));
+        let send_box = EncryptionBox::new(ChaChaBox::new(public_key, secret_key));
+        let recv_box = EncryptionBox::new(ChaChaBox::new(public_key, secret_key));
 
         Ok(Self {
             send_stream,
             recv_stream,
-            sent_count: Arc::new(AtomicUsize::new(0)),
-            recv_count: Arc::new(AtomicUsize::new(0)),
-            secret_box: Arc::new(Mutex::new(secret_box)),
+            send_secret: Arc::new(Mutex::new(send_box)),
+            recv_secret: Arc::new(Mutex::new(recv_box)),
         })
     }
 
@@ -47,16 +51,7 @@ impl CryptoStream {
     /// length: u64
     /// data: length
     pub fn send(&self, mut event: Event) -> VentedResult<()> {
-        let number = self.sent_count.fetch_add(1, Ordering::SeqCst);
-        let nonce = generate_nonce(number);
-
-        let ciphertext = self.secret_box.lock().encrypt(
-            &nonce,
-            Payload {
-                msg: &event.as_bytes(),
-                aad: &[],
-            },
-        )?;
+        let ciphertext = self.send_secret.lock().encrypt(&event.as_bytes())?;
         let mut stream = self.send_stream.lock();
         let mut length_raw = [0u8; 8];
         BigEndian::write_u64(&mut length_raw, ciphertext.len() as u64);
@@ -83,20 +78,78 @@ impl CryptoStream {
         stream.read(&mut ciphertext)?;
         log::trace!("Received raw message");
 
-        let number = self.recv_count.fetch_add(1, Ordering::SeqCst);
-        let nonce = generate_nonce(number);
-        let plaintext = self.secret_box.lock().decrypt(
+        let plaintext = self.recv_secret.lock().decrypt(&ciphertext)?;
+
+        let event = Event::from_bytes(&mut &plaintext[..])?;
+        log::trace!("Decoded message to event '{}'", event.name);
+
+        Ok(event)
+    }
+
+    /// Updates the keys in the inner encryption box
+    pub fn update_key(&self, secret_key: &SecretKey, public_key: &PublicKey) {
+        let send_box = ChaChaBox::new(public_key, secret_key);
+        let recv_box = ChaChaBox::new(public_key, secret_key);
+        self.send_secret.lock().swap_box(send_box);
+        self.recv_secret.lock().swap_box(recv_box);
+        log::trace!("Updated secret");
+    }
+}
+
+pub struct EncryptionBox<T>
+where
+    T: Aead,
+{
+    inner: T,
+    counter: usize,
+}
+
+impl<T> EncryptionBox<T>
+where
+    T: Aead,
+{
+    /// Creates a new encryption box with the given inner value
+    pub fn new(inner: T) -> Self {
+        Self { inner, counter: 0 }
+    }
+
+    /// Swaps the crypto box for a new one
+    pub fn swap_box(&mut self, new_box: T) {
+        self.inner = new_box;
+    }
+}
+
+impl EncryptionBox<ChaChaBox> {
+    /// Encrypts the given data by using the inner ChaCha box and nonce
+    pub fn encrypt(&mut self, data: &[u8]) -> VentedResult<Vec<u8>> {
+        let nonce = generate_nonce(self.counter);
+
+        let ciphertext = self.inner.encrypt(
             &nonce,
             Payload {
-                msg: &ciphertext,
                 aad: &[],
+                msg: data,
             },
         )?;
+        self.counter += 1;
 
-        let event = Event::from_bytes(&mut &plaintext[..])?;
-        log::trace!("Decoded message to event '{}'", event.name);
+        Ok(ciphertext)
+    }
 
-        Ok(event)
+    /// Decrypts the data by using the inner ChaCha box and nonce
+    pub fn decrypt(&mut self, data: &[u8]) -> VentedResult<Vec<u8>> {
+        let nonce = generate_nonce(self.counter);
+
+        let plaintext = self.inner.decrypt(
+            &nonce,
+            Payload {
+                msg: data,
+                aad: &[],
+            },
+        )?;
+        self.counter += 1;
+
+        Ok(plaintext)
     }
 }
 

src/server/mod.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 use std::net::{TcpListener, TcpStream};
 
-use crypto_box::{ChaChaBox, PublicKey, SecretKey};
+use crypto_box::{PublicKey, SecretKey};
 use scheduled_thread_pool::ScheduledThreadPool;
 
 use crate::crypto::CryptoStream;
@@ -16,6 +16,7 @@ use crate::server::server_events::{
 };
 use crossbeam_utils::sync::WaitGroup;
 use parking_lot::Mutex;
+use sha2::Digest;
 use std::io::Write;
 use std::sync::Arc;
 use std::thread;
@@ -327,7 +328,6 @@ impl VentedServer {
         }
 
         let public_key = PublicKey::from(public_key);
-        let secret_box = ChaChaBox::new(&public_key, &secret_key);
 
         let node_data = if let Some(data) = known_nodes.lock().iter().find(|n| n.id == node_id) {
             data.clone()
@@ -337,14 +337,21 @@ impl VentedServer {
             return Err(UnknownNode(node_id));
         };
 
-        let mut stream = CryptoStream::new(stream, secret_box)?;
+        let mut stream = CryptoStream::new(stream, &public_key, &secret_key)?;
 
         log::trace!("Authenticating recipient...");
-        Self::authenticate_other(&mut stream, node_data.public_key)?;
+        let key_a = Self::authenticate_other(&mut stream, node_data.public_key)?;
         log::trace!("Authenticating self...");
+        let key_b =
             Self::authenticate_self(&mut stream, StaticSecret::from(global_secret.to_bytes()))?;
         log::trace!("Connection fully authenticated.");
 
+        let pre_secret = StaticSecret::from(secret_key.to_bytes()).diffie_hellman(&public_key);
+        let final_secret =
+            Self::generate_final_secret(pre_secret.to_bytes().to_vec(), key_a, key_b);
+        let final_public = final_secret.public_key();
+        stream.update_key(&final_secret, &final_public);
+
         Ok((node_id, stream))
     }
 
@@ -400,20 +407,30 @@ impl VentedServer {
             .as_bytes(),
         )?;
         stream.flush()?;
-        let secret_box = ChaChaBox::new(&public_key, &secret_key);
-        let mut stream = CryptoStream::new(stream, secret_box)?;
+
+        let mut stream = CryptoStream::new(stream, &public_key, &secret_key)?;
 
         log::trace!("Authenticating self...");
+        let key_a =
             Self::authenticate_self(&mut stream, StaticSecret::from(global_secret.to_bytes()))?;
         log::trace!("Authenticating recipient...");
-        Self::authenticate_other(&mut stream, node_data.public_key)?;
+        let key_b = Self::authenticate_other(&mut stream, node_data.public_key)?;
         log::trace!("Connection fully authenticated.");
 
+        let pre_secret = StaticSecret::from(secret_key.to_bytes()).diffie_hellman(&public_key);
+        let final_secret =
+            Self::generate_final_secret(pre_secret.to_bytes().to_vec(), key_a, key_b);
+        let final_public = final_secret.public_key();
+        stream.update_key(&final_secret, &final_public);
+
         Ok((node_id, stream))
     }
 
     /// Performs the challenged side of the authentication challenge
-    fn authenticate_self(stream: &CryptoStream, static_secret: StaticSecret) -> VentedResult<()> {
+    fn authenticate_self(
+        stream: &CryptoStream,
+        static_secret: StaticSecret,
+    ) -> VentedResult<Vec<u8>> {
         let challenge_event = stream.read()?;
 
         if challenge_event.name != CHALLENGE_EVENT {
@@ -433,7 +450,7 @@ impl VentedServer {
         let response = stream.read()?;
 
         match response.name.as_str() {
-            ACCEPT_EVENT => Ok(()),
+            ACCEPT_EVENT => Ok(auth_key.to_bytes().to_vec()),
             REJECT_EVENT => Err(VentedError::Rejected),
             _ => {
                 stream.send(Event::new(REJECT_EVENT))?;
@@ -446,7 +463,7 @@ impl VentedServer {
     fn authenticate_other(
         stream: &CryptoStream,
         other_static_public: PublicKey,
-    ) -> VentedResult<()> {
+    ) -> VentedResult<Vec<u8>> {
         let auth_secret = SecretKey::generate(&mut rand::thread_rng());
         stream.send(Event::with_payload(
             CHALLENGE_EVENT,
@@ -470,7 +487,7 @@ impl VentedServer {
             Err(VentedError::AuthFailed)
         } else {
             stream.send(Event::new(ACCEPT_EVENT))?;
-            Ok(())
+            Ok(calculated_secret.to_vec())
         }
     }
 
@@ -481,4 +498,21 @@ impl VentedServer {
 
         parts_a.get(0) == parts_b.get(0) && parts_a.get(1) == parts_b.get(1)
     }
+
+    /// Generates a secret from handshake components
+    fn generate_final_secret(
+        mut pre_secret: Vec<u8>,
+        mut key_a: Vec<u8>,
+        mut key_b: Vec<u8>,
+    ) -> SecretKey {
+        let mut secret_data = Vec::new();
+        secret_data.append(&mut pre_secret);
+        secret_data.append(&mut key_a);
+        secret_data.append(&mut key_b);
+        let final_secret = sha2::Sha256::digest(&secret_data).to_vec();
+        let mut final_secret_arr = [0u8; 32];
+        final_secret_arr.copy_from_slice(&final_secret[..]);
+
+        SecretKey::from(final_secret_arr)
+    }
 }