Skip to content

Commit

Permalink
Encryption and simple key management
Browse files Browse the repository at this point in the history
  • Loading branch information
ktdlr committed Oct 23, 2024
1 parent 399813d commit ef311f7
Show file tree
Hide file tree
Showing 10 changed files with 218 additions and 14 deletions.
29 changes: 29 additions & 0 deletions packages/ciphernode/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions packages/ciphernode/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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"
Expand All @@ -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"] }
Expand All @@ -62,3 +65,4 @@ libp2p = { version = "0.53.2", features = [
"gossipsub",
"quic",
] }
zeroize = "1.8.1"
2 changes: 1 addition & 1 deletion packages/ciphernode/data/src/snapshot.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::{DataStore, Repository};
use crate::Repository;
use anyhow::Result;
use async_trait::async_trait;
use serde::{de::DeserializeOwned, Serialize};
Expand Down
9 changes: 5 additions & 4 deletions packages/ciphernode/enclave/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
5 changes: 4 additions & 1 deletion packages/ciphernode/enclave/src/main.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use std::env;

use alloy::primitives::Address;
use clap::Parser;
use enclave::load_config;
use enclave::{ensure_req_env, load_config};
use enclave_node::{listen_for_shutdown, MainCiphernode};
use tracing::info;

Expand Down Expand Up @@ -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?;

Expand Down
13 changes: 10 additions & 3 deletions packages/ciphernode/keyshare/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,18 @@ version = "0.1.0"
edition = "2021"

[dependencies]
actix = { workspace = true }
aes-gcm = { workspace = true }
anyhow = { workspace = true }
argon2 = { workspace = true }
async-trait = { workspace = true }
data = { path = "../data" }
enclave-core = { path = "../core" }
fhe = { path = "../fhe" }
actix = { workspace = true }
anyhow = { workspace = true }
rand = { workspace = true }
serde = { workspace = true }
async-trait = { workspace = true }
tracing = { workspace = true }
zeroize = { workspace = true }

[dev-dependencies]
proptest = { workspace = true }
132 changes: 132 additions & 0 deletions packages/ciphernode/keyshare/src/encryption.rs
Original file line number Diff line number Diff line change
@@ -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[..]);
}
}
32 changes: 28 additions & 4 deletions packages/ciphernode/keyshare/src/keyshare.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use crate::{decrypt_data, encrypt_data};
use actix::prelude::*;
use anyhow::{anyhow, Result};
use async_trait::async_trait;
Expand All @@ -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>,
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -134,10 +158,10 @@ 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;
};
Expand Down
2 changes: 2 additions & 0 deletions packages/ciphernode/keyshare/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,2 +1,4 @@
mod encryption;
mod keyshare;
pub use encryption::*;
pub use keyshare::*;
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit ef311f7

Please sign in to comment.