Refactor codebase for reusability

Cargo.toml

@@ -4,11 +4,21 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-ring = "0.17.8"
 hmac = "0.12"
 sha2 = "0.10"
 curve25519-dalek = "4.1.1"
 x25519-dalek = {version="^2.0.1", features = ["static_secrets"] }
-rand_core = {version="0.6.4", features = ["std"]}
-rand = "0.9.1"
+
+# wasm incompatible
+#ring = "0.17.8"
+#rand_core = {version="0.6.4", features = ["std"]}
+#rand = "0.9.1"
+
+#wasm compatible
+getrandom = { version = "0.2", features = ["js"] }
+aes-gcm = "0.10"
+serde = { version = "1.0.219", features = ["derive"] }
+
+log = "0.4.26"
+hex = "0.4.3"
 

src/client.rs

@@ -0,0 +1,130 @@
+use x25519_dalek::PublicKey;
+use sha2::Sha256;
+use hmac::{Hmac, Mac};
+use hmac::digest::Digest;
+use crate::common::{
+    NTorCertificate,
+    InitSessionMessage,
+    InitSessionResponse,
+    NTorParty,
+    PrivatePublicKeyPair
+};
+use crate::helpers::{generate_private_public_key_pair};
+
+#[derive(Clone)]
+pub struct NTorClient {
+    ephemeral_key_pair: PrivatePublicKeyPair,
+    pub(crate) shared_secret: Option<Vec<u8>>,
+}
+
+impl NTorClient {
+    pub fn new() -> Self {
+        Self {
+            ephemeral_key_pair: PrivatePublicKeyPair {
+                private_key: None,
+                public_key: PublicKey::from([0; 32]),
+            },
+            shared_secret: None,
+        }
+    }
+
+    pub fn initialise_session(&mut self) -> InitSessionMessage {
+        self.ephemeral_key_pair = generate_private_public_key_pair();
+
+        InitSessionMessage {
+            client_ephemeral_public_key: self.ephemeral_key_pair.public_key,
+        }
+    }
+
+    // Steps 15 - 20 of the Goldberg 2012 paper.
+    pub fn handle_response_from_server(
+        &mut self,
+        server_certificate: &NTorCertificate,
+        msg: &InitSessionResponse,
+    ) -> bool {
+        println!("Client:");
+
+        // Step 18: Compute the shared secret.
+        let mut buffer: Vec<u8> = Vec::new();
+
+        // ECDH Client private ephemeral * server static public key
+        let taken_private_key = self.ephemeral_key_pair.private_key.take().unwrap();
+        let mut ecdh_result_1 = taken_private_key
+            .diffie_hellman(&msg.server_ephemeral_public_key)
+            .to_bytes()
+            .to_vec();
+        println!("[Debug] ECDH result 1: {:?}", ecdh_result_1);
+        buffer.append(&mut ecdh_result_1);
+
+        // ECDH Client private ephemeral * server ephemeral public Key
+        let mut ecdh_result_2 = taken_private_key
+            .diffie_hellman(&server_certificate.public_key)
+            .to_bytes()
+            .to_vec();
+        println!("[Debug] ECDH result 2: {:?}", ecdh_result_2);
+        buffer.append(&mut ecdh_result_2);
+
+        // Server id
+        buffer.append(&mut server_certificate.server_id.as_bytes().to_vec());
+
+        // Client ephemeral public
+        buffer.append(&mut self.ephemeral_key_pair.public_key.as_bytes().to_vec());
+
+        // Server ephemeral public
+        buffer.append(&mut msg.server_ephemeral_public_key.as_bytes().to_vec());
+
+        // "ntor" string identifier
+        buffer.append(&mut "ntor".as_bytes().to_vec());
+
+        // Instantiate and run hashing function
+        let mut hasher = Sha256::new();
+        hasher.update(buffer);
+        let sha256_hash = hasher.finalize();
+        let sha256_hash: &[u8; 32] = match sha256_hash.as_slice().try_into() {
+            Ok(array_ref) => array_ref,
+            Err(_) => {
+                panic!("Invalid sha256 hash length");
+            }
+        };
+
+        let secret_key_prime = &sha256_hash[0..16];
+        println!("[Debug] Client secret key prime: {:?}", secret_key_prime);
+
+        let secret_key = &sha256_hash[16..];
+
+        // Step 19: Compute HMAC (t_b in the paper)
+
+        let mut buffer: Vec<u8> = Vec::new();
+        buffer.append(&mut server_certificate.server_id.as_bytes().to_vec());
+        buffer.append(&mut msg.server_ephemeral_public_key.as_bytes().to_vec());
+        buffer.append(&mut self.ephemeral_key_pair.public_key.as_bytes().to_vec());
+        buffer.append(&mut "ntor".as_bytes().to_vec());
+        buffer.append(&mut "server".as_bytes().to_vec());
+
+        let mut hmac_hash = Hmac::<Sha256>::new_from_slice(&buffer).unwrap();
+        hmac_hash.update(secret_key_prime);
+        let computed_t_b_hash = hmac_hash.finalize().into_bytes().to_vec();
+
+        // assert that computed_t_b_hash equals t_hash generated by server
+        if computed_t_b_hash == msg.t_b_hash {
+            self.shared_secret = Some(secret_key.to_vec());
+
+            println!("Shared secret:");
+            println!("{:?}\n", secret_key);
+            true
+        } else {
+            println!("Failed to verify the shared secret: try again bro.");
+            false
+        }
+    }
+}
+
+impl NTorParty for NTorClient {
+    fn get_shared_secret(&self) -> Option<Vec<u8>> {
+        self.shared_secret.clone()
+    }
+
+    fn set_shared_secret(&mut self, shared_secret: Vec<u8>) {
+        self.shared_secret = Some(shared_secret);
+    }
+}

src/common.rs

@@ -0,0 +1,148 @@
+use x25519_dalek::{PublicKey, StaticSecret};
+use crate::helpers;
+use crate::helpers::{key_derivation, vec_to_array32};
+
+#[derive(Clone)]
+pub struct PrivatePublicKeyPair {
+    // In the future, type StaticSecret should be reserved for the server's static and the EphemeralSecret reserved for the ephemeral private key.
+    // However, as a quirk of the nTOR protocol, we also need to use StaticSecret for the client's ephemeral private key hence why it is adopted here.
+    pub(crate) private_key: Option<StaticSecret>,
+    pub(crate) public_key: PublicKey,
+}
+
+impl PrivatePublicKeyPair {
+    pub fn get_public_key(&self) -> PublicKey {
+        self.public_key
+    }
+
+    pub fn get_private_key(&self) -> Option<StaticSecret> {
+        self.private_key.clone()
+    }
+}
+
+#[derive(Clone)]
+pub struct NTorCertificate {
+    pub(crate) public_key: PublicKey,
+    pub(crate) server_id: String,
+}
+
+impl NTorCertificate {
+    pub fn new(public_key: Vec<u8>, server_id: String) -> Self {
+        let pub_key = TryInto::<[u8; 32]>::try_into(public_key).unwrap();
+        NTorCertificate {
+            public_key: PublicKey::from(pub_key),
+            server_id
+        }
+    }
+
+    pub fn public_key(&self) -> Vec<u8> {
+        self.public_key.to_bytes().to_vec()
+    }
+}
+
+// In the paper, the outgoing message is ("ntor", B_id, client_ephemeral_public_key).
+pub struct InitSessionMessage {
+    pub(crate) client_ephemeral_public_key: PublicKey,
+}
+
+impl InitSessionMessage {
+    pub fn from(bytes: Vec<u8>) -> Self {
+        let u8_array = vec_to_array32(bytes);
+        InitSessionMessage {
+            client_ephemeral_public_key: PublicKey::from(u8_array),
+        }
+    }
+
+    pub fn public_key(&self) -> Vec<u8> {
+        self.client_ephemeral_public_key.to_bytes().to_vec()
+    }
+}
+
+// In the paper, the return message is ("ntor", server_ephemeral_public_key, t_b_hash).
+pub struct InitSessionResponse {
+    pub(crate) server_ephemeral_public_key: PublicKey,
+    pub(crate) t_b_hash: Vec<u8>,
+}
+
+impl InitSessionResponse {
+    pub fn new(public_key: Vec<u8>, t_b_hash: Vec<u8>) -> Self {
+        let pub_key = TryInto::<[u8; 32]>::try_into(public_key).unwrap();
+        return InitSessionResponse {
+            server_ephemeral_public_key: PublicKey::from(pub_key),
+            t_b_hash,
+        }
+    }
+
+    pub fn public_key(&self) -> Vec<u8> {
+        self.server_ephemeral_public_key.to_bytes().to_vec()
+    }
+
+    pub fn t_b_hash(&self) -> Vec<u8> {
+        self.t_b_hash.clone()
+    }
+}
+
+pub struct EncryptedMessage {
+    pub nonce: [u8; 12],
+    pub data: Vec<u8>
+}
+
+pub trait NTorParty {
+    fn get_shared_secret(&self) -> Option<Vec<u8>>;
+
+    fn set_shared_secret(&mut self, shared_secret: Vec<u8>);
+
+    fn encrypt(&self, data: Vec<u8>) -> Result<EncryptedMessage, &'static str> {
+        if let Some(key) = self.get_shared_secret() {
+            let encrypt_key = key_derivation(&key);
+            return match helpers::encrypt(encrypt_key, data) {
+                Ok((nonce, encrypted_message)) => {
+                    Ok(EncryptedMessage {
+                        nonce,
+                        data: encrypted_message,
+                    })
+                }
+                Err(err) => Err(err)
+            }
+        }
+        Err("no encryption key found")
+    }
+
+    fn decrypt(&self, encrypted_message: EncryptedMessage) -> Result<Vec<u8>, &'static str> {
+        if let Some(key) = self.get_shared_secret() {
+            let decrypt_key = key_derivation(&key);
+            return helpers::decrypt(encrypted_message.nonce, decrypt_key, encrypted_message.data);
+        }
+        Err("no decryption key found")
+    }
+
+    fn wasm_encrypt(&self, data: Vec<u8>) -> Result<(Vec<u8>, Vec<u8>), &'static str> {
+        if let Some(key) = self.get_shared_secret() {
+            let encrypt_key = key_derivation(&key);
+            return match helpers::encrypt(encrypt_key, data) {
+                Ok((nonce, encrypted_message)) => {
+                    Ok((nonce.to_vec(), encrypted_message))
+                }
+                Err(err) => Err(err)
+            }
+        }
+        Err("no encryption key found")
+    }
+
+    fn wasm_decrypt(&self, nonce: Vec<u8>, data: Vec<u8>) -> Result<Vec<u8>, &'static str> {
+        if let Some(key) = self.get_shared_secret() {
+            let decrypt_key = key_derivation(&key);
+            // return helpers::wasm_decrypt(nonce, decrypt_key, encrypted_message.data);
+            return match TryInto::<[u8; 12]>::try_into(nonce) {
+                Ok(nonce12) => {
+                    return match helpers::decrypt(nonce12, decrypt_key, data) {
+                        Ok(decrypted) => Ok(decrypted),
+                        Err(err) => Err(err)
+                    }
+                },
+                Err(_err) => Err("invalid nonce")
+            }
+        }
+        Err("no decryption key found")
+    }
+}

src/helpers.rs

@@ -0,0 +1,71 @@
+use std::convert::TryInto;
+use aes_gcm::{Aes256Gcm, Key, KeyInit, Nonce};
+use aes_gcm::aead::Aead;
+use x25519_dalek::{PublicKey, StaticSecret};
+use crate::common::PrivatePublicKeyPair;
+
+/// Returns an array of zeros if conversion fails.
+pub fn vec_to_array32(vec: Vec<u8>) -> [u8; 32] {
+    if vec.len() == 32 {
+        vec.try_into().unwrap()
+    } else {
+        [0u8; 32]
+    }
+}
+
+pub fn generate_private_public_key_pair() -> PrivatePublicKeyPair {
+    let mut buf = [0u8; 32];
+    getrandom::getrandom(&mut buf).expect("generate random failed");
+    let private_key = StaticSecret::from(buf);
+    let public_key = PublicKey::from(&private_key);
+
+    PrivatePublicKeyPair {
+        private_key: Some(private_key),
+        public_key,
+    }
+}
+
+pub fn key_derivation(shared_secret: &Vec<u8>) -> Vec<u8> {
+    let mut encrypt_key = shared_secret.clone(); // fixme use a reliable kdf
+    encrypt_key.extend(shared_secret.clone());
+    encrypt_key
+}
+
+pub(crate) fn encrypt(key_bytes: Vec<u8>, data: Vec<u8>) -> Result<([u8; 12], Vec<u8>), &'static str> {
+    if key_bytes.len() != 32 {
+        return Err("Invalid key length for AES-256");
+    }
+
+    let key = Key::<Aes256Gcm>::from_slice(&key_bytes);
+    let cipher = Aes256Gcm::new(key);
+
+    let mut nonce_bytes = [0u8; 12];
+    getrandom::getrandom(&mut nonce_bytes).map_err(|_| "Random generation failed")?;
+    let nonce = Nonce::from_slice(&nonce_bytes); // 96-bits; unique per message
+
+    let ciphertext = cipher
+        .encrypt(nonce, data.as_ref())
+        .map_err(|_| "Encryption failed")?;
+
+    Ok((nonce_bytes, ciphertext))
+}
+
+pub(crate) fn decrypt(nonce_bytes: [u8; 12], key: Vec<u8>, ciphertext: Vec<u8>) -> Result<Vec<u8>, &'static str> {
+    return match TryInto::<[u8; 32]>::try_into(key) {
+        Ok(key_bytes) => {
+            let key = Key::<Aes256Gcm>::from_slice(&key_bytes);
+            let cipher = Aes256Gcm::new(key);
+            let nonce = Nonce::from_slice(&nonce_bytes);
+
+            let decrypted_data = cipher
+                .decrypt(nonce, ciphertext.as_ref())
+                .map_err(|_| "Decryption failed")?;
+
+            Ok(decrypted_data)
+        }
+        Err(_) => {
+            Err("Invalid key")
+        }
+    }
+}
+

src/lib.rs

@@ -0,0 +1,6 @@
+pub mod common;
+pub mod server;
+pub mod client;
+
+mod test;
+pub mod helpers;