chore: refactor code

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<Ciphertext, String> {
+        // 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<x_wing::SharedSecret> {
+        identities.iter().find_map(|key| key.decapsulate(ct).ok())
+    }
+
+    fn unwrap_file_key(ss: &x_wing::SharedSecret, body: &[u8]) -> Result<FileKey, String> {
+        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<I: Iterator<Item = (usize, Stanza)>>(
         &self,
         file_index: usize,
         stanzas: I,
     ) -> Result<FileKey, Vec<identity::Error>> {
-        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(),
             });
         }