From 265760e24d8ed120f7f978490961cbdba2eb2d0d Mon Sep 17 00:00:00 2001 From: Aaron Kaiser Date: Tue, 25 Nov 2025 23:31:37 +0000 Subject: [PATCH] chore: refactor code --- src/main.rs | 173 ++++++++++++++++++++-------------------------------- 1 file changed, 67 insertions(+), 106 deletions(-) diff --git a/src/main.rs b/src/main.rs index a5d0869..719b123 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,7 +15,6 @@ use base64::prelude::*; use clap::Parser; use rand_core::OsRng; use std::{ - array::TryFromSliceError, collections::{HashMap, HashSet}, io, }; @@ -51,10 +50,12 @@ impl RecipientPlugin { self.recipients .iter() .map(|recipient| { - let (ct, ss) = recipient + let (ct, shk) = recipient .encapsulate(&mut OsRng) .expect("X-Wing encapsulation should not fail with a valid RNG"); - let wrapped_key = aead_encrypt(&ss, file_key.expose_secret()); + + let wrapped_key = aead_encrypt(&shk, file_key.expose_secret()); + Stanza { tag: PLUGIN_NAME.to_string(), args: vec![BASE64_STANDARD_NO_PAD.encode(ct.to_bytes())], @@ -79,14 +80,12 @@ impl RecipientPluginV1 for RecipientPlugin { }); } - let pk: Result<&[u8; ENCAPSULATION_KEY_SIZE], TryFromSliceError> = bytes.try_into(); - let pk = match pk { - Ok(x) => Ok(EncapsulationKey::from(x)), - Err(_) => Err(recipient::Error::Recipient { + let pk: &[u8; ENCAPSULATION_KEY_SIZE] = + bytes.try_into().map_err(|_| recipient::Error::Recipient { index, message: "Invalid recipient".to_string(), - }), - }?; + })?; + let pk = EncapsulationKey::from(pk); self.recipients.push(pk); @@ -102,7 +101,7 @@ impl RecipientPluginV1 for RecipientPlugin { if plugin_name != PLUGIN_NAME { return Err(recipient::Error::Identity { index, - message: "This Identity should not be handled by this plugin".to_owned(), + message: "This identity should not be handled by this plugin".to_string(), }); } @@ -140,116 +139,78 @@ struct IdentityPlugin { } impl IdentityPlugin { + fn parse_ciphertext(arg: &str) -> Result { + // age-plugin-xwing up to version 0.1.1 encoded its ciphertext using BASE64_STANDARD. + // We still want to be able to support decrypting those. + let decoded = BASE64_STANDARD_NO_PAD + .decode(arg) + .or_else(|_| BASE64_STANDARD.decode(arg)) + .map_err(|_| "Malformed base64".to_string())?; + + let bytes: [u8; CIPHERTEXT_SIZE] = decoded + .try_into() + .map_err(|_| "Malformed ciphertext".to_string())?; + + Ok(Ciphertext::from(&bytes)) + } + + fn try_decapsulate( + identities: &[DecapsulationKey], + ct: &Ciphertext, + ) -> Option { + identities.iter().find_map(|key| key.decapsulate(ct).ok()) + } + + fn unwrap_file_key(ss: &x_wing::SharedSecret, body: &[u8]) -> Result { + let plaintext = aead_decrypt(ss, FILE_KEY_BYTES, body).map_err(|e| e.to_string())?; + + let key_bytes: [u8; FILE_KEY_BYTES] = plaintext.try_into().map_err(|_| { + format!("aead_decrypt returned a plaintext with a different size than {FILE_KEY_BYTES}") + })?; + + Ok(FileKey::new(Box::new(key_bytes))) + } + fn decrypt_stanzas>( &self, file_index: usize, stanzas: I, ) -> Result> { - let mut file_key = None; let mut errors = Vec::new(); for (stanza_index, stanza) in stanzas { - let arg = match stanza.args.first() { - Some(arg) => arg, - None => { - errors.push(identity::Error::Stanza { - file_index, - stanza_index, - message: "Stanza is missing arguments".to_string(), - }); - continue; - } - }; + // Try to decrypt this stanza + let result = (|| { + let arg = stanza.args.first().ok_or("Stanza is missing arguments")?; - // age-plugin-xwing up to version 0.1.1 encoded its ciphertext using BASE64_STANDARD. - // We still want to be able to support decrypting those. - let ct = match ( - BASE64_STANDARD_NO_PAD.decode(arg), - BASE64_STANDARD.decode(arg), - ) { - (Ok(ct), _) => ct, - (_, Ok(ct)) => ct, - _ => { - errors.push(identity::Error::Stanza { - file_index, - stanza_index, - message: "Malformed base64".to_string(), - }); - continue; - } - }; + let ct = Self::parse_ciphertext(arg)?; - let ct: [u8; CIPHERTEXT_SIZE] = match ct.try_into() { - Ok(ct) => ct, - Err(_) => { - errors.push(identity::Error::Stanza { - file_index, - stanza_index, - message: "Malformed ciphertext".to_string(), - }); - continue; - } - }; + let ss = Self::try_decapsulate(&self.identities, &ct) + .ok_or("No identity found that can decrypt the file")?; - let ct = Ciphertext::from(&ct); + Self::unwrap_file_key(&ss, &stanza.body) + })(); - let ss = match self - .identities - .iter() - .filter_map(|key| key.decapsulate(&ct).ok()) - .next() - { - Some(ss) => ss, - None => { - errors.push(identity::Error::Stanza { - file_index, - stanza_index, - message: "No identity found that can decrypt the file".to_string(), - }); - continue; - } - }; - - let unwrapped_file_key = match aead_decrypt(&ss, FILE_KEY_BYTES, &stanza.body) { - Ok(file_key) => match file_key.try_into() { - Ok(key_bytes) => FileKey::new(Box::new(key_bytes)), - Err(_) => { - errors.push(identity::Error::Stanza { - file_index, - stanza_index, - message: format!( - "aead_decrypt returned a plaintext with a different size than {FILE_KEY_BYTES}" - ), - }); - continue; - } - }, - Err(e) => { - errors.push(identity::Error::Stanza { - file_index, - stanza_index, - message: e.to_string(), - }); - continue; - } - }; - - file_key = Some(unwrapped_file_key); + match result { + Ok(file_key) => return Ok(file_key), + Err(message) => errors.push(identity::Error::Stanza { + file_index, + stanza_index, + message, + }), + } } - if !errors.is_empty() { - return Err(errors); + // If we get here, no stanza was successfully decrypted + if errors.is_empty() { + Err(vec![identity::Error::Stanza { + file_index, + stanza_index: 0, + message: "No stanzas found to be handled by this plugin".to_string(), + }]) + } else { + Err(errors) } - - if let Some(file_key) = file_key { - return Ok(file_key); - } - - Err(vec![identity::Error::Stanza { - file_index, - stanza_index: 0, - message: "No stanzas found to be handled by this plugin".to_string(), - }]) } } @@ -263,7 +224,7 @@ impl IdentityPluginV1 for IdentityPlugin { if plugin_name != PLUGIN_NAME { return Err(identity::Error::Identity { index, - message: "This Identity should not be handled by this plugin".to_string(), + message: "This identity should not be handled by this plugin".to_string(), }); }