refactor: make the code more organzied

src/main.rs

@@ -19,7 +19,8 @@ use std::{
     io,
 };
 use x_wing::{
-    Ciphertext, DecapsulationKey, EncapsulationKey, DECAPSULATION_KEY_SIZE, ENCAPSULATION_KEY_SIZE,
+    Ciphertext, DecapsulationKey, EncapsulationKey, CIPHERTEXT_SIZE, DECAPSULATION_KEY_SIZE,
+    ENCAPSULATION_KEY_SIZE,
 };
 
 const PLUGIN_NAME: &str = "xwing";
@@ -44,6 +45,23 @@ struct RecipientPlugin {
     recipients: Vec<EncapsulationKey>,
 }
 
+impl RecipientPlugin {
+    fn wrap_file_key(&self, file_key: FileKey) -> Vec<Stanza> {
+        self.recipients
+            .iter()
+            .map(|recipient| {
+                let (ct, ss) = recipient.encapsulate(&mut OsRng).unwrap();
+                let wrapped_key = aead_encrypt(&ss, file_key.expose_secret());
+                Stanza {
+                    tag: PLUGIN_NAME.to_string(),
+                    args: vec![BASE64_STANDARD.encode(ct.as_bytes())],
+                    body: wrapped_key,
+                }
+            })
+            .collect()
+    }
+}
+
 impl RecipientPluginV1 for RecipientPlugin {
     fn add_recipient(
         &mut self,
@@ -110,20 +128,7 @@ impl RecipientPluginV1 for RecipientPlugin {
     ) -> io::Result<Result<Vec<Vec<Stanza>>, Vec<recipient::Error>>> {
         Ok(Ok(file_keys
             .into_iter()
-            .map(|file_key| {
+            .map(|file_key| self.wrap_file_key(file_key))
-                self.recipients
-                    .iter()
-                    .map(|recipient| {
-                        let (ct, ss) = recipient.encapsulate(&mut OsRng).unwrap();
-                        let wrapped_key = aead_encrypt(&ss, file_key.expose_secret());
-                        Stanza {
-                            tag: PLUGIN_NAME.to_string(),
-                            args: vec![BASE64_STANDARD.encode(ct.as_bytes())],
-                            body: wrapped_key,
-                        }
-                    })
-                    .collect()
-            })
             .collect()))
     }
 
@@ -137,6 +142,107 @@ struct IdentityPlugin {
     identities: Vec<DecapsulationKey>,
 }
 
+impl IdentityPlugin {
+    fn decrypt_stanzas(
+        &self,
+        file_index: usize,
+        stanzas: Vec<(usize, Stanza)>,
+    ) -> 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;
+                }
+            };
+
+            let ct = match BASE64_STANDARD.decode(arg) {
+                Ok(ct) => ct,
+                Err(_) => {
+                    errors.push(identity::Error::Stanza {
+                        file_index,
+                        stanza_index,
+                        message: "Malformed base64".to_string(),
+                    });
+                    continue;
+                }
+            };
+
+            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 ct = Ciphertext::from(&ct);
+
+            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) => FileKey::new(Box::new(file_key.try_into().unwrap_or_else(|_| {
+                    panic!(
+                        "aead_decrypt returned a plaintext with a different size as {}",
+                        FILE_KEY_BYTES
+                    )
+                }))),
+                Err(e) => {
+                    errors.push(identity::Error::Stanza {
+                        file_index,
+                        stanza_index,
+                        message: e.to_string(),
+                    });
+                    continue;
+                }
+            };
+
+            file_key = Some(unwrapped_file_key);
+        }
+
+        if !errors.is_empty() {
+            return 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 handeled by this plugin".to_string(),
+        }])
+    }
+}
+
 impl IdentityPluginV1 for IdentityPlugin {
     fn add_identity(
         &mut self,
@@ -153,7 +259,7 @@ impl IdentityPluginV1 for IdentityPlugin {
 
         let bytes: [u8; DECAPSULATION_KEY_SIZE] = match bytes.try_into() {
             Ok(x) => x,
-            _ => {
+            Err(_) => {
                 return Err(identity::Error::Identity {
                     index,
                     message: "Invalid identity".to_string(),
@@ -171,43 +277,20 @@ impl IdentityPluginV1 for IdentityPlugin {
         files: Vec<Vec<Stanza>>,
         mut _callbacks: impl Callbacks<identity::Error>,
     ) -> io::Result<HashMap<usize, Result<FileKey, Vec<identity::Error>>>> {
-        Ok(files
+        let mut ret = HashMap::default();
-            .into_iter()
-            .map(|file| try_decrypt_file_key(&self.identities, file))
-            .enumerate()
-            .map(|(file_index, file_key)| {
-                file_key.ok_or(vec![identity::Error::Stanza {
-                    file_index,
-                    stanza_index: 1,
-                    message: "Invalid stanzas".to_string(),
-                }])
-            })
-            .enumerate()
-            .collect())
-    }
-}
 
-fn try_decrypt_file_key(keys: &Vec<DecapsulationKey>, stanzas: Vec<Stanza>) -> Option<FileKey> {
+        for (file_index, stanzas) in files.into_iter().enumerate() {
-    stanzas
+            let x_wing_stanzas = stanzas
-        .iter()
+                .into_iter()
-        .filter_map(|stanza| try_decrypt_stanza(keys, stanza))
+                .enumerate()
-        .next()
+                .filter(|(_, stanza)| stanza.tag == PLUGIN_NAME)
-}
+                .collect();
 
-fn try_decrypt_stanza(keys: &Vec<DecapsulationKey>, stanza: &Stanza) -> Option<FileKey> {
+            ret.insert(file_index, self.decrypt_stanzas(file_index, x_wing_stanzas));
-    let ct = BASE64_STANDARD.decode(stanza.args.first()?).ok()?;
-    let ct = Ciphertext::from(&ct.try_into().ok()?);
-
-    for key in keys {
-        let ss = key.decapsulate(&ct).unwrap();
-        let file_key = aead_decrypt(&ss, FILE_KEY_BYTES, &stanza.body);
-
-        if let Ok(file_key) = file_key {
-            return Some(FileKey::new(Box::new(file_key.try_into().ok()?)));
         }
-    }
 
-    None
+        Ok(ret)
+    }
 }
 
 #[derive(Debug, Parser)]
@@ -225,8 +308,7 @@ fn main() -> io::Result<()> {
         return Ok(());
     }
 
-    // Here you can assume the binary is being run directly by a user,
+    // Here you can assume the binary is being run directly by a user, and perform administrative tasks like generating keys.
-    // and perform administrative tasks like generating keys.
     let (sk, pk) = x_wing::generate_key_pair_from_os_rng();
     print_new_identity(PLUGIN_NAME, sk.as_bytes(), &pk.as_bytes());