feat(rust): Add struct for XChaCha20Poly1305 key (#5365)

Author: rohanjadvani

Cargo.lock

@@ -53,6 +53,8 @@ rand = "0.8.5"
 tiny-bip39 = "1.0"
 xxhash-rust = { version = "0.8.5", features = ["xxh3"] }
 argon2 = { version = "0.5.3", features = ["password-hash"] }
+hkdf = "0.12.4"
+sha2 = "0.10"
 
 [dev-dependencies]
 hex-literal = "0.4"

ironfish-rust/Cargo.toml

@@ -32,6 +32,7 @@ pub enum IronfishErrorKind {
     FailedSignatureVerification,
     FailedXChaCha20Poly1305Decryption,
     FailedXChaCha20Poly1305Encryption,
+    FailedHkdfExpansion,
     IllegalValue,
     InconsistentWitness,
     InvalidAssetIdentifier,

ironfish-rust/src/errors.rs

@@ -4,21 +4,140 @@
 
 use std::io;
 
+use argon2::RECOMMENDED_SALT_LEN;
 use argon2::{password_hash::SaltString, Argon2};
 use chacha20poly1305::aead::Aead;
 use chacha20poly1305::{Key, KeyInit, XChaCha20Poly1305, XNonce};
+use hkdf::Hkdf;
 use rand::{thread_rng, RngCore};
+use sha2::Sha256;
 
 use crate::errors::{IronfishError, IronfishErrorKind};
 
-const KEY_LENGTH: usize = 32;
-const NONCE_LENGTH: usize = 24;
+pub const KEY_LENGTH: usize = 32;
+pub const SALT_LENGTH: usize = RECOMMENDED_SALT_LEN;
+pub const XNONCE_LENGTH: usize = 24;
+
+#[derive(Debug)]
+pub struct XChaCha20Poly1305Key {
+    pub key: [u8; KEY_LENGTH],
+
+    pub nonce: [u8; XNONCE_LENGTH],
+
+    pub salt: [u8; SALT_LENGTH],
+}
+
+impl XChaCha20Poly1305Key {
+    pub fn generate(passphrase: &[u8]) -> Result<XChaCha20Poly1305Key, IronfishError> {
+        let mut nonce = [0u8; XNONCE_LENGTH];
+        thread_rng().fill_bytes(&mut nonce);
+
+        let mut salt = [0u8; SALT_LENGTH];
+        thread_rng().fill_bytes(&mut salt);
+
+        XChaCha20Poly1305Key::from_parts(passphrase, salt, nonce)
+    }
+
+    pub fn from_parts(
+        passphrase: &[u8],
+        salt: [u8; SALT_LENGTH],
+        nonce: [u8; XNONCE_LENGTH],
+    ) -> Result<XChaCha20Poly1305Key, IronfishError> {
+        let mut key = [0u8; KEY_LENGTH];
+        let argon2 = Argon2::default();
+
+        argon2
+            .hash_password_into(passphrase, &salt, &mut key)
+            .map_err(|_| IronfishError::new(IronfishErrorKind::FailedArgon2Hash))?;
+
+        Ok(XChaCha20Poly1305Key { key, salt, nonce })
+    }
+
+    pub fn derive_key(&self, salt: [u8; SALT_LENGTH]) -> Result<[u8; KEY_LENGTH], IronfishError> {
+        let hkdf = Hkdf::<Sha256>::new(None, &self.key);
+
+        let mut okm = [0u8; KEY_LENGTH];
+        hkdf.expand(&salt, &mut okm)
+            .map_err(|_| IronfishError::new(IronfishErrorKind::FailedHkdfExpansion))?;
+
+        Ok(okm)
+    }
+
+    pub fn derive_new_key(&self) -> Result<XChaCha20Poly1305Key, IronfishError> {
+        let mut nonce = [0u8; XNONCE_LENGTH];
+        thread_rng().fill_bytes(&mut nonce);
+
+        let mut salt = [0u8; SALT_LENGTH];
+        thread_rng().fill_bytes(&mut salt);
+
+        let hkdf = Hkdf::<Sha256>::new(None, &self.key);
+
+        let mut okm = [0u8; KEY_LENGTH];
+        hkdf.expand(&salt, &mut okm)
+            .map_err(|_| IronfishError::new(IronfishErrorKind::FailedHkdfExpansion))?;
+
+        Ok(XChaCha20Poly1305Key {
+            key: okm,
+            salt,
+            nonce,
+        })
+    }
+
+    pub fn read<R: io::Read>(mut reader: R) -> Result<Self, IronfishError> {
+        let mut salt = [0u8; SALT_LENGTH];
+        reader.read_exact(&mut salt)?;
+
+        let mut nonce = [0u8; XNONCE_LENGTH];
+        reader.read_exact(&mut nonce)?;
+
+        let mut key = [0u8; KEY_LENGTH];
+        reader.read_exact(&mut key)?;
+
+        Ok(XChaCha20Poly1305Key { salt, nonce, key })
+    }
+
+    pub fn write<W: io::Write>(&self, mut writer: W) -> Result<(), IronfishError> {
+        writer.write_all(&self.salt)?;
+        writer.write_all(&self.nonce)?;
+        writer.write_all(&self.key)?;
+
+        Ok(())
+    }
+
+    pub fn destroy(&mut self) {
+        self.key.fill(0);
+        self.nonce.fill(0);
+        self.salt.fill(0);
+    }
+
+    pub fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, IronfishError> {
+        let nonce = XNonce::from_slice(&self.nonce);
+        let key = Key::from(self.key);
+        let cipher = XChaCha20Poly1305::new(&key);
+
+        let ciphertext = cipher.encrypt(nonce, plaintext).map_err(|_| {
+            IronfishError::new(IronfishErrorKind::FailedXChaCha20Poly1305Encryption)
+        })?;
+
+        Ok(ciphertext)
+    }
+
+    pub fn decrypt(&self, ciphertext: Vec<u8>) -> Result<Vec<u8>, IronfishError> {
+        let nonce = XNonce::from_slice(&self.nonce);
+        let key = Key::from(self.key);
+        let cipher = XChaCha20Poly1305::new(&key);
+
+        cipher
+            .decrypt(nonce, ciphertext.as_ref())
+            .map_err(|_| IronfishError::new(IronfishErrorKind::FailedXChaCha20Poly1305Decryption))
+    }
+}
 
 #[derive(Debug)]
 pub struct EncryptOutput {
     pub salt: Vec<u8>,
 
-    pub nonce: [u8; NONCE_LENGTH],
+    pub nonce: [u8; XNONCE_LENGTH],
 
     pub ciphertext: Vec<u8>,
 }
@@ -46,7 +165,7 @@ impl EncryptOutput {
         let mut salt = vec![0u8; salt_len];
         reader.read_exact(&mut salt)?;
 
-        let mut nonce = [0u8; NONCE_LENGTH];
+        let mut nonce = [0u8; XNONCE_LENGTH];
         reader.read_exact(&mut nonce)?;
 
         let mut ciphertext_len = [0u8; 4];
@@ -88,7 +207,7 @@ pub fn encrypt(plaintext: &[u8], passphrase: &[u8]) -> Result<EncryptOutput, Iro
     let key = derive_key(passphrase, salt_bytes)?;
 
     let cipher = XChaCha20Poly1305::new(&key);
-    let mut nonce_bytes = [0u8; NONCE_LENGTH];
+    let mut nonce_bytes = [0u8; XNONCE_LENGTH];
     thread_rng().fill_bytes(&mut nonce_bytes);
     let nonce = XNonce::from_slice(&nonce_bytes);
 
@@ -119,19 +238,22 @@ pub fn decrypt(
 
 #[cfg(test)]
 mod test {
-    use crate::xchacha20poly1305::{decrypt, encrypt};
-
-    use super::EncryptOutput;
+    use crate::xchacha20poly1305::XChaCha20Poly1305Key;
 
     #[test]
     fn test_valid_passphrase() {
         let plaintext = "thisissensitivedata";
         let passphrase = "supersecretpassword";
 
-        let encrypted_output = encrypt(plaintext.as_bytes(), passphrase.as_bytes())
+        let key =
+            XChaCha20Poly1305Key::generate(passphrase.as_bytes()).expect("should generate key");
+
+        let encrypted_output = key
+            .encrypt(plaintext.as_bytes())
             .expect("should successfully encrypt");
-        let decrypted =
-            decrypt(encrypted_output, passphrase.as_bytes()).expect("should decrypt successfully");
+        let decrypted = key
+            .decrypt(encrypted_output)
+            .expect("should decrypt successfully");
 
         assert_eq!(decrypted, plaintext.as_bytes());
     }
@@ -142,28 +264,33 @@ mod test {
         let passphrase = "supersecretpassword";
         let incorrect_passphrase = "foobar";
 
-        let encrypted_output = encrypt(plaintext.as_bytes(), passphrase.as_bytes())
+        let key =
+            XChaCha20Poly1305Key::generate(passphrase.as_bytes()).expect("should generate key");
+
+        let encrypted_output = key
+            .encrypt(plaintext.as_bytes())
             .expect("should successfully encrypt");
 
-        decrypt(encrypted_output, incorrect_passphrase.as_bytes())
+        let incorrect_key = XChaCha20Poly1305Key::generate(incorrect_passphrase.as_bytes())
+            .expect("should generate key");
+
+        incorrect_key
+            .decrypt(encrypted_output)
             .expect_err("should fail decryption");
     }
 
     #[test]
-    fn test_encrypt_output_serialization() {
-        let plaintext = "thisissensitivedata";
+    fn test_derive_key() {
         let passphrase = "supersecretpassword";
 
-        let encrypted_output = encrypt(plaintext.as_bytes(), passphrase.as_bytes())
-            .expect("should successfully encrypt");
-
-        let mut vec: Vec<u8> = vec![];
-        encrypted_output
-            .write(&mut vec)
-            .expect("should serialize successfully");
+        let encryption_key = XChaCha20Poly1305Key::generate(passphrase.as_bytes())
+            .expect("should successfully generate key");
 
-        let deserialized = EncryptOutput::read(&vec[..]).expect("should deserialize successfully");
+        let key = encryption_key.derive_new_key().expect("should derive key");
+        let derived_key = encryption_key
+            .derive_key(key.salt)
+            .expect("should derive key");
 
-        assert_eq!(encrypted_output, deserialized);
+        assert_eq!(key.key, derived_key);
     }
 }

ironfish-rust/src/xchacha20poly1305.rs

@@ -390,6 +390,10 @@ criteria = "safe-to-deploy"
 version = "0.1.19"
 criteria = "safe-to-deploy"
 
+[[exemptions.hkdf]]
+version = "0.12.4"
+criteria = "safe-to-deploy"
+
 [[exemptions.hmac]]
 version = "0.11.0"
 criteria = "safe-to-deploy"