Encryption and key management (ready for review)

packages/ciphernode/Cargo.toml

@@ -19,6 +19,7 @@ members = [
 [workspace.dependencies]
 actix = "0.13.5"
 actix-rt = "2.10.0"
+aes-gcm = "0.10.3"
 alloy = { version = "0.3.3", features = ["full"] }
 alloy-primitives = { version = "0.6", default-features = false, features = [
   "rlp",
@@ -27,6 +28,7 @@ alloy-primitives = { version = "0.6", default-features = false, features = [
 ] }
 alloy-sol-types = { version = "0.6" }
 anyhow = "1.0.86"
+argon2 = "0.5.3"
 async-std = { version = "1.12", features = ["attributes"] }
 async-trait = "0.1"
 bincode = "1.3.3"
@@ -40,6 +42,7 @@ fhe-util = { git = "https://github.com/gnosisguild/fhe.rs", version = "0.1.0-bet
 futures = "0.3.30"
 futures-util = "0.3"
 num = "0.4.3"
+proptest = "1.5.0"
 rand_chacha = "0.3.1"
 rand = "0.8.5"
 serde = { version = "1.0.208", features = ["derive"] }
@@ -62,3 +65,4 @@ libp2p = { version = "0.53.2", features = [
   "gossipsub",
   "quic",
 ] }
+zeroize = "1.8.1"

packages/ciphernode/data/src/snapshot.rs

@@ -1,4 +1,4 @@
-use crate::{DataStore, Repository};
+use crate::Repository;
 use anyhow::Result;
 use async_trait::async_trait;
 use serde::{de::DeserializeOwned, Serialize};

packages/ciphernode/enclave/Cargo.toml

@@ -8,12 +8,13 @@ repository = "https://github.com/gnosisguild/enclave/packages/ciphernode"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-enclave_node = { path = "../enclave_node" }
-alloy = { workspace = true }
-clap = { workspace = true }
 actix = { workspace = true }
 actix-rt = { workspace = true }
-tokio = { workspace = true }
+alloy = { workspace = true }
+anyhow = { workspace = true }
+clap = { workspace = true }
 config = "0.14.0"
+enclave_node = { path = "../enclave_node" }
+tokio = { workspace = true }
 tracing-subscriber = { workspace = true } 
 tracing = { workspace = true }

packages/ciphernode/enclave/src/main.rs

@@ -1,3 +1,5 @@
+use std::env;
+
 use alloy::primitives::Address;
 use clap::Parser;
 use enclave::load_config;
@@ -39,6 +41,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let address = Address::parse_checksummed(&args.address, None).expect("Invalid address");
     info!("LAUNCHING CIPHERNODE: ({})", address);
     let config = load_config(&args.config)?;
+
     let (bus, handle) =
         MainCiphernode::attach(config, address, args.data_location.as_deref()).await?;
 

packages/ciphernode/keyshare/Cargo.toml

@@ -4,11 +4,15 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+actix = { workspace = true }
+aes-gcm = { workspace = true, version = "=0.10.3" }
+anyhow = { workspace = true }
+argon2  = { workspace = true, version = "=0.5.2" }
+async-trait = { workspace = true }
 data = { path = "../data" }
 enclave-core = { path = "../core" }
 fhe = { path = "../fhe" }
-actix = { workspace = true }
-anyhow = { workspace = true }
+rand = { workspace = true, version = "=0.8.5" }
 serde = { workspace = true }
-async-trait = { workspace = true }
 tracing = { workspace = true }
+zeroize = { workspace = true, version = "=1.6.0" }

packages/ciphernode/keyshare/src/encryption.rs

@@ -0,0 +1,132 @@
+use std::{env, ops::Deref};
+
+use aes_gcm::{
+    aead::{Aead, KeyInit},
+    Aes256Gcm, Nonce,
+};
+use anyhow::{anyhow, Result};
+use argon2::{Algorithm, Argon2, Params, Version};
+use rand::{rngs::OsRng, RngCore};
+use zeroize::{Zeroize, Zeroizing};
+
+// ARGON2 PARAMS
+const ARGON2_M_COST: u32 = 32 * 1024; // 32 MiB
+const ARGON2_T_COST: u32 = 2;
+const ARGON2_P_COST: u32 = 2;
+const ARGON2_OUTPUT_LEN: usize = 32;
+const ARGON2_ALGORITHM: Algorithm = Algorithm::Argon2id;
+const ARGON2_VERSION: Version = Version::V0x13;
+
+// AES PARAMS
+const AES_SALT_LEN: usize = 32;
+const AES_NONCE_LEN: usize = 12;
+
+fn argon2_derive_key(
+    password_bytes: Zeroizing<Vec<u8>>,
+    salt: &[u8],
+) -> Result<Zeroizing<Vec<u8>>> {
+    let mut derived_key = Zeroizing::new(vec![0u8; ARGON2_OUTPUT_LEN]);
+
+    let params = Params::new(
+        ARGON2_M_COST,
+        ARGON2_T_COST,
+        ARGON2_P_COST,
+        Some(ARGON2_OUTPUT_LEN),
+    )
+    .map_err(|_| anyhow!("Could not create params"))?;
+    Argon2::new(ARGON2_ALGORITHM, ARGON2_VERSION, params)
+        .hash_password_into(&password_bytes, &salt, &mut derived_key)
+        .map_err(|_| anyhow!("Key derivation error"))?;
+    Ok(derived_key)
+}
+
+pub fn encrypt_data(data: &mut Vec<u8>) -> Result<Vec<u8>> {
+    // Convert password to bytes in a zeroizing buffer
+    let password_bytes = Zeroizing::new(env::var("CIPHERNODE_SECRET")?.as_bytes().to_vec());
+
+    // Generate a random salt for Argon2
+    let mut salt = [0u8; AES_SALT_LEN];
+    OsRng.fill_bytes(&mut salt);
+
+    // Generate a random nonce for AES-GCM
+    let mut nonce_bytes = [0u8; AES_NONCE_LEN];
+    OsRng.fill_bytes(&mut nonce_bytes);
+    let nonce = Nonce::from_slice(&nonce_bytes);
+
+    // Derive key using Argon2
+    let derived_key = argon2_derive_key(password_bytes, &salt)?;
+
+    // Create AES-GCM cipher
+    let cipher = Aes256Gcm::new_from_slice(&derived_key)?;
+
+    // Encrypt the data
+    let ciphertext = cipher
+        .encrypt(nonce, data.as_ref())
+        .map_err(|_| anyhow!("Could not AES Encrypt given plaintext."))?;
+
+    data.zeroize(); // Zeroize sensitive input data
+
+    // NOTE: password_bytes and derived_key will be automatically zeroized when dropped
+
+    // Pack data
+    let mut output = Vec::with_capacity(salt.len() + nonce_bytes.len() + ciphertext.len());
+    output.extend_from_slice(&salt);
+    output.extend_from_slice(&nonce_bytes);
+    output.extend_from_slice(&ciphertext);
+
+    Ok(output)
+}
+
+pub fn decrypt_data(encrypted_data: &[u8]) -> Result<Vec<u8>> {
+    const AES_HEADER_LEN: usize = AES_SALT_LEN + AES_NONCE_LEN;
+    if encrypted_data.len() < AES_HEADER_LEN {
+        return Err(anyhow!("Invalid encrypted data length"));
+    }
+
+    let password_bytes = Zeroizing::new(env::var("CIPHERNODE_SECRET")?.as_bytes().to_vec());
+
+    // Extract salt and nonce
+    let salt = &encrypted_data[..AES_SALT_LEN];
+    let nonce = Nonce::from_slice(&encrypted_data[AES_SALT_LEN..AES_HEADER_LEN]);
+    let ciphertext = &encrypted_data[AES_HEADER_LEN..];
+
+    // Derive key using Argon2
+    let derived_key = argon2_derive_key(password_bytes, &salt)?;
+
+    // Create cipher and decrypt
+    let cipher = Aes256Gcm::new_from_slice(&derived_key)?;
+    let plaintext = cipher
+        .decrypt(nonce, ciphertext)
+        .map_err(|_| anyhow!("Could not decrypt data"))?;
+
+    // NOTE: password_bytes and derived_key will be automatically zeroized when dropped
+
+    Ok(plaintext)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_encryption_decryption() {
+        use std::time::Instant;
+        println!("TESTING");
+        env::set_var("CIPHERNODE_SECRET", "my_secure_password");
+        let data = b"Hello, world!";
+
+        let start = Instant::now();
+        let encrypted = encrypt_data(&mut data.to_vec()).unwrap();
+        let encryption_time = start.elapsed();
+
+        let start = Instant::now();
+        let decrypted = decrypt_data(&encrypted).unwrap();
+        let decryption_time = start.elapsed();
+
+        println!("Encryption took: {:?}", encryption_time);
+        println!("Decryption took: {:?}", decryption_time);
+        println!("Total time: {:?}", encryption_time + decryption_time);
+
+        assert_eq!(data, &decrypted[..]);
+    }
+}

packages/ciphernode/keyshare/src/keyshare.rs

@@ -1,3 +1,4 @@
+use crate::{decrypt_data, encrypt_data};
 use actix::prelude::*;
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
@@ -8,8 +9,10 @@ use enclave_core::{
 };
 use fhe::{DecryptCiphertext, Fhe};
 use serde::{Deserialize, Serialize};
-use std::sync::Arc;
+use std::env;
+use std::{process, sync::Arc};
 use tracing::warn;
+use zeroize::Zeroizing;
 
 pub struct Keyshare {
     fhe: Arc<Fhe>,
@@ -45,6 +48,21 @@ impl Keyshare {
             address: params.address,
         }
     }
+
+    fn set_secret(&mut self, mut data: Vec<u8>) -> Result<()> {
+        let encrypted = encrypt_data(&mut data)?;
+        self.secret = Some(encrypted);
+        Ok(())
+    }
+
+    fn get_secret(&self) -> Result<Vec<u8>> {
+        let encrypted = self
+            .secret
+            .clone()
+            .ok_or(anyhow!("No secret share available on Keyshare"))?;
+        let decrypted = decrypt_data(&encrypted)?;
+        Ok(decrypted)
+    }
 }
 
 impl Snapshot for Keyshare {
@@ -107,7 +125,13 @@ impl Handler<CiphernodeSelected> for Keyshare {
         };
 
         // Save secret on state
-        self.secret = Some(secret);
+        let Ok(()) = self.set_secret(secret) else {
+            self.bus.do_send(EnclaveEvent::from_error(
+                EnclaveErrorType::KeyGeneration,
+                anyhow!("Error encrypting Keyshare for {e3_id}"),
+            ));
+            return;
+        };
 
         // Broadcast the KeyshareCreated message
         self.bus.do_send(EnclaveEvent::from(KeyshareCreated {
@@ -134,17 +158,17 @@ impl Handler<CiphertextOutputPublished> for Keyshare {
             ciphertext_output,
         } = event;
 
-        let Some(secret) = &self.secret else {
+        let Ok(secret) = self.get_secret() else {
             self.bus.err(
                 EnclaveErrorType::Decryption,
-                anyhow!("secret not found on Keyshare for e3_id {e3_id}"),
+                anyhow!("Secret not available for Keyshare for e3_id {e3_id}"),
             );
             return;
         };
 
         let Ok(decryption_share) = self.fhe.decrypt_ciphertext(DecryptCiphertext {
             ciphertext: ciphertext_output.clone(),
-            unsafe_secret: secret.to_vec(),
+            unsafe_secret: secret,
         }) else {
             self.bus.err(
                 EnclaveErrorType::Decryption,

packages/ciphernode/keyshare/src/lib.rs

@@ -1,2 +1,4 @@
+mod encryption;
 mod keyshare;
+pub use encryption::*;
 pub use keyshare::*;

packages/ciphernode/tests/tests/test_aggregation_and_decryption.rs

@@ -24,7 +24,7 @@ use fhe_traits::{FheEncoder, FheEncrypter, Serialize};
 use rand::Rng;
 use rand::SeedableRng;
 use rand_chacha::ChaCha20Rng;
-use std::{sync::Arc, time::Duration};
+use std::{env, sync::Arc, time::Duration};
 use tokio::sync::Mutex;
 use tokio::{sync::mpsc::channel, time::sleep};
 
@@ -271,6 +271,7 @@ fn get_common_setup() -> Result<(
 #[actix::test]
 async fn test_public_key_aggregation_and_decryption() -> Result<()> {
     // Setup
+    env::set_var("CIPHERNODE_SECRET", "Don't tell anyone my secret");
     let (bus, rng, seed, params, crpoly, e3_id) = get_common_setup()?;
 
     // Setup actual ciphernodes and dispatch add events
@@ -376,6 +377,7 @@ async fn test_public_key_aggregation_and_decryption() -> Result<()> {
 
 #[actix::test]
 async fn test_stopped_keyshares_retain_state() -> Result<()> {
+    env::set_var("CIPHERNODE_SECRET", "Don't tell anyone my secret");
     let (bus, rng, seed, params, crpoly, e3_id) = get_common_setup()?;
 
     let eth_addrs = create_random_eth_addrs(2);