use my own key exchange (diffie hellman) crate :D

src/lib.rs

@@ -1,106 +1,85 @@
 use chacha::Block;
+use diffie_hellman::{Public, Secret};
+use rand::{RngCore, rng};
 use std::io::{Error, Read, Write};
 use std::net::{Shutdown, TcpStream};
-use rand::{rng, RngCore};
-use x25519_dalek::{EphemeralSecret, PublicKey};
 
 #[inline]
 fn u8_to_u32(to: &mut [u32], from: &[u8]) {
     for i in 0..to.len() {
-        to[i] = u32::from_le_bytes(from[4*i..][..4].try_into().unwrap());
+        to[i] = u32::from_le_bytes(from[4 * i..][..4].try_into().unwrap());
     }
 }
 
-/*
-#[repr(u8)]
-enum Exchange {
-    DHCurve25519 = 0
-}
-
-impl TryFrom<u8> for Exchange {
-    type Error = Error;
-
-    fn try_from(value: u8) -> Result<Self, Self::Error> {
-        match value {
-            0 => Ok(Exchange::DHCurve25519),
-            _ => Err(Error::new(ErrorKind::InvalidData, "Invalid exchange code")),
-        }
-    }
-}
-
-#[repr(u8)]
-enum Encryption {
-    ChaCha = 0
-}
-
-impl TryFrom<u8> for Encryption {
-    type Error = Error;
-
-    fn try_from(value: u8) -> Result<Self, Self::Error> {
-        match value {
-            0 => Ok(Encryption::ChaCha),
-            _ => Err(Error::new(ErrorKind::InvalidData, "Invalid encryption code")),
-        }
-    }
-}
-
- */
-
 #[derive(Debug)]
 pub struct EncryptedStream {
     socket: TcpStream,
     block: Block,
 }
 
+const KEY_DERIVATION_CONTEXT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
+
 impl EncryptedStream {
+    /// connect and negotiate an encrypted stream as a client
     pub fn connect(address: &str) -> Result<EncryptedStream, Error> {
-        let secret = EphemeralSecret::random();
-        let public = PublicKey::from(&secret);
+        let secret = Secret::random();
+        let public = Public::from(&secret);
 
         let mut socket = TcpStream::connect(address)?;
 
-        socket.write(public.as_bytes())?;
+        socket.write(&public.as_bytes())?;
 
         let mut their_public_bytes = [0u8; 32];
         socket.read_exact(&mut their_public_bytes)?;
 
-        let their_public = PublicKey::from(their_public_bytes);
+        let their_public = Public::from(their_public_bytes);
 
         let shared_secret = secret.diffie_hellman(&their_public);
 
         let mut nonce_bytes = [0u8; 12];
         socket.read_exact(&mut nonce_bytes)?;
 
+        let key = blake3::derive_key(KEY_DERIVATION_CONTEXT, &shared_secret.to_bytes());
+
+        Ok(EncryptedStream::wrap(socket, key, nonce_bytes))
+    }
+
+    pub fn wrap(socket: TcpStream, key: [u8; 32], nonce: [u8; 12]) -> EncryptedStream {
         let mut nonce_u32s = [0u32; 3];
-        u8_to_u32(&mut nonce_u32s, &nonce_bytes);
+        u8_to_u32(&mut nonce_u32s, &nonce);
 
         let mut key_u32s = [0u32; 8];
-        u8_to_u32(&mut key_u32s, &shared_secret.to_bytes());
+        u8_to_u32(&mut key_u32s, &key);
 
-        Ok(EncryptedStream {
+        EncryptedStream {
             socket,
-            block: Block::new(key_u32s, 0, nonce_u32s, 20)
-        })
+            block: Block::new(key_u32s, 0, nonce_u32s, 20),
+        }
     }
 
     pub fn shutdown(&self, how: Shutdown) -> Result<(), Error> {
         self.socket.shutdown(how)
     }
+
+    pub fn reset(&mut self, next_nonce: [u32; 3]) {
+        self.block.reset(next_nonce)
+    }
 }
 
 impl TryFrom<TcpStream> for EncryptedStream {
     type Error = Error;
 
+    /// Accept and negotiate an encrypted connection as a server
     fn try_from(mut socket: TcpStream) -> Result<Self, Self::Error> {
-        let secret = EphemeralSecret::random();
-        let public = PublicKey::from(&secret);
+        let secret = Secret::random();
+        let public = Public::from(&secret);
 
         let mut their_public_bytes = [0u8; 32];
         socket.read(&mut their_public_bytes)?;
 
-        socket.write(public.as_bytes())?;
+        socket.write(&public.as_bytes())?;
 
-        let their_public = PublicKey::from(their_public_bytes);
+        let their_public = Public::from(their_public_bytes);
 
         let shared_secret = secret.diffie_hellman(&their_public);
 
@@ -109,16 +88,9 @@ impl TryFrom<TcpStream> for EncryptedStream {
 
         socket.write(&nonce_bytes)?;
 
-        let mut nonce_u32s = [0u32; 3];
-        u8_to_u32(&mut nonce_u32s, &nonce_bytes);
+        let key = blake3::derive_key(KEY_DERIVATION_CONTEXT, &shared_secret.to_bytes());
 
-        let mut key_u32s = [0u32; 8];
-        u8_to_u32(&mut key_u32s, &shared_secret.to_bytes());
-
-        Ok(EncryptedStream {
-            socket,
-            block: Block::new(key_u32s, 0, nonce_u32s, 20)
-        })
+        Ok(EncryptedStream::wrap(socket, key, nonce_bytes))
     }
 }
 
@@ -154,6 +126,5 @@ impl Write for EncryptedStream {
     }
 }
 
-
 #[cfg(test)]
 mod tests {}

src/main.rs

@@ -1,7 +1,7 @@
 use hermes::EncryptedStream;
 use std::env::args;
 use std::io::{Read, Write};
-use std::net::{Shutdown, TcpListener, TcpStream};
+use std::net::{Shutdown, TcpListener};
 use std::process::exit;
 
 fn main() {
@@ -21,7 +21,9 @@ fn main() {
                 let mut encrypted = EncryptedStream::try_from(incoming.unwrap()).unwrap();
                 println!("successfully encrypted stream");
 
-                encrypted.write(b"the fitness gram pacer test is a multi-stage aerobic fitness test...").unwrap();
+                encrypted
+                    .write(b"the fitness gram pacer test is a multi-stage aerobic fitness test...")
+                    .unwrap();
                 println!("successfully wrote data");
 
                 encrypted.flush().unwrap();