From 85d0ee7e319b01cd728d39c77c6b288d972260f9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rodolphe=20Br=C3=A9ard?= <v.code.ngyvz7vf@what.tf>
Date: Sun, 25 Feb 2024 13:40:19 +0100
Subject: [PATCH] Implement decryption

---
 src/encryption.rs | 62 +++++++++++++++++++++++++++++++++++++++++------
 src/scheme.rs     | 10 +++++++-
 2 files changed, 64 insertions(+), 8 deletions(-)

diff --git a/src/encryption.rs b/src/encryption.rs
index 24d18be..4e60f78 100644
--- a/src/encryption.rs
+++ b/src/encryption.rs
@@ -5,6 +5,7 @@ use crate::{storage, InputKeyMaterialList};
 use chacha20poly1305::aead::{Aead, KeyInit, Payload};
 use chacha20poly1305::{Key, XChaCha20Poly1305, XNonce};
 
+pub(crate) type DecryptionFunction = dyn Fn(&[u8], &EncryptedData, &str) -> Result<Vec<u8>>;
 pub(crate) type EncryptionFunction = dyn Fn(&[u8], &[u8], &str) -> Result<EncryptedData>;
 
 pub(crate) struct EncryptedData {
@@ -12,6 +13,13 @@ pub(crate) struct EncryptedData {
 	pub(crate) ciphertext: Vec<u8>,
 }
 
+#[inline]
+fn generate_aad(key_context: &[&str], data_context: &[impl AsRef<[u8]>]) -> String {
+	let key_context_canon = canonicalize(key_context);
+	let data_context_canon = canonicalize(data_context);
+	join_canonicalized_str(&key_context_canon, &data_context_canon)
+}
+
 pub fn encrypt(
 	ikml: &InputKeyMaterialList,
 	key_context: &[&str],
@@ -23,9 +31,7 @@ pub fn encrypt(
 	let key = derive_key(ikm, key_context);
 
 	// Generate the AAD
-	let key_context_canon = canonicalize(key_context);
-	let data_context_canon = canonicalize(data_context);
-	let aad = join_canonicalized_str(&key_context_canon, &data_context_canon);
+	let aad = generate_aad(key_context, data_context);
 
 	// Encrypt
 	let encryption_function = ikm.scheme.get_encryption();
@@ -68,10 +74,42 @@ pub(crate) fn xchacha20poly1305_encrypt(
 pub fn decrypt(
 	ikml: &InputKeyMaterialList,
 	key_context: &[&str],
-	data: impl AsRef<[u8]>,
+	stored_data: &str,
 	data_context: &[impl AsRef<[u8]>],
 ) -> Result<Vec<u8>> {
-	unimplemented!("decrypt");
+	// Retreive the IKM and the encrypted data
+	let (ikm_id, encrypted_data) = storage::decode(stored_data)?;
+	let ikm = ikml.get_ikm_by_id(ikm_id)?;
+
+	// Derive the key
+	let key = derive_key(ikm, key_context);
+
+	// Generate the AAD
+	let aad = generate_aad(key_context, data_context);
+
+	// Decrypt
+	let decryption_function = ikm.scheme.get_decryption();
+	decryption_function(&key, &encrypted_data, &aad)
+}
+
+pub(crate) fn xchacha20poly1305_decrypt(
+	key: &[u8],
+	encrypted_data: &EncryptedData,
+	aad: &str,
+) -> Result<Vec<u8>> {
+	// Adapt the key and the nonce
+	let key = Key::from_slice(key);
+	let nonce = XNonce::from_slice(&encrypted_data.nonce);
+
+	// Prepare the payload
+	let payload = Payload {
+		msg: &encrypted_data.ciphertext,
+		aad: aad.as_bytes(),
+	};
+
+	// Decrypt the payload
+	let cipher = XChaCha20Poly1305::new(key);
+	Ok(cipher.decrypt(nonce, payload)?)
 }
 
 #[cfg(test)]
@@ -92,6 +130,7 @@ mod tests {
 
 	#[test]
 	fn encrypt_decrypt_no_context() {
+		// Encrypt
 		let lst = get_ikm_lst();
 		let res = encrypt(&lst, &[], TEST_DATA, EMPTY_DATA_CTX);
 		assert!(res.is_ok());
@@ -99,11 +138,16 @@ mod tests {
 		assert!(ciphertext.starts_with("AQAAAA:"));
 		assert_eq!(ciphertext.len(), 98);
 
-		// TODO: decrypt
+		// Decrypt
+		let res = decrypt(&lst, &[], &ciphertext, EMPTY_DATA_CTX);
+		assert!(res.is_ok());
+		let plaintext = res.unwrap();
+		assert_eq!(plaintext, TEST_DATA);
 	}
 
 	#[test]
 	fn encrypt_decrypt_with_context() {
+		// Encrypt
 		let lst = get_ikm_lst();
 		let res = encrypt(&lst, TEST_KEY_CTX, TEST_DATA, TEST_DATA_CTX);
 		assert!(res.is_ok());
@@ -111,6 +155,10 @@ mod tests {
 		assert!(ciphertext.starts_with("AQAAAA:"));
 		assert_eq!(ciphertext.len(), 98);
 
-		// TODO: decrypt
+		// Decrypt
+		let res = decrypt(&lst, TEST_KEY_CTX, &ciphertext, TEST_DATA_CTX);
+		assert!(res.is_ok());
+		let plaintext = res.unwrap();
+		assert_eq!(plaintext, TEST_DATA);
 	}
 }
diff --git a/src/scheme.rs b/src/scheme.rs
index da67e0c..df0d34e 100644
--- a/src/scheme.rs
+++ b/src/scheme.rs
@@ -1,4 +1,4 @@
-use crate::encryption::EncryptionFunction;
+use crate::encryption::{DecryptionFunction, EncryptionFunction};
 use crate::kdf::KdfFunction;
 use crate::Error;
 
@@ -16,6 +16,14 @@ impl Scheme {
 		}
 	}
 
+	pub(crate) fn get_decryption(&self) -> Box<DecryptionFunction> {
+		match self {
+			Scheme::XChaCha20Poly1305WithBlake3 => {
+				Box::new(crate::encryption::xchacha20poly1305_decrypt)
+			}
+		}
+	}
+
 	pub(crate) fn get_encryption(&self) -> Box<EncryptionFunction> {
 		match self {
 			Scheme::XChaCha20Poly1305WithBlake3 => {