chore: refactor code

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "age-plugin-xwing"
 description = "X-Wing plugin for age clients"
-version = "0.1.1"
+version = "0.2.0"
 edition = "2021"
 license = "MIT"
 repository = "https://github.com/Rixxc/age-plugin-xwing"
@@ -16,6 +16,5 @@ age-core = "0.11.0"
 age-plugin = "0.6.0"
 base64 = "0.22.1"
 clap = { version = "4.5.1", features = ["derive"] }
-x-wing = { version = "0.0.1-pre.1", features = ["getrandom", "zeroize"] }
+x-wing = { version = "0.1.0-pre.2", features = ["os_rng", "zeroize"] }
-rand_core = "0.6.4"
+rand_core = "0.9.3"
-kem = "0.3.0-pre.0"

README.md

@@ -41,6 +41,6 @@ echo 'It works!' | age -e -r age1xwing1jfdy5fhryzmfg6y89dka6xr2wup9favgprpq0x542
 Decryption:
 
 ```sh
-$ age -d -i x_wing_key.txt secret.enc
+$ age -d -i age_x_wing.key secret.enc
 It works!
 ```

src/main.rs

@@ -13,16 +13,14 @@ use age_plugin::{
 };
 use base64::prelude::*;
 use clap::Parser;
-use kem::{Decapsulate, Encapsulate};
 use rand_core::OsRng;
 use std::{
-    array::TryFromSliceError,
     collections::{HashMap, HashSet},
     io,
 };
 use x_wing::{
-    Ciphertext, DecapsulationKey, EncapsulationKey, CIPHERTEXT_SIZE, DECAPSULATION_KEY_SIZE,
+    Ciphertext, Decapsulate, DecapsulationKey, Encapsulate, EncapsulationKey, CIPHERTEXT_SIZE,
-    ENCAPSULATION_KEY_SIZE,
+    DECAPSULATION_KEY_SIZE, ENCAPSULATION_KEY_SIZE,
 };
 
 const PLUGIN_NAME: &str = "xwing";
@@ -52,11 +50,15 @@ impl RecipientPlugin {
         self.recipients
             .iter()
             .map(|recipient| {
-                let (ct, ss) = recipient.encapsulate(&mut OsRng).unwrap();
+                let (ct, shk) = recipient
-                let wrapped_key = aead_encrypt(&ss, file_key.expose_secret());
+                    .encapsulate(&mut OsRng)
+                    .expect("X-Wing encapsulation should not fail with a valid RNG");
+
+                let wrapped_key = aead_encrypt(&shk, file_key.expose_secret());
+
                 Stanza {
                     tag: PLUGIN_NAME.to_string(),
-                    args: vec![BASE64_STANDARD.encode(ct.as_bytes())],
+                    args: vec![BASE64_STANDARD_NO_PAD.encode(ct.to_bytes())],
                     body: wrapped_key,
                 }
             })
@@ -74,20 +76,16 @@ impl RecipientPluginV1 for RecipientPlugin {
         if plugin_name != PLUGIN_NAME {
             return Err(recipient::Error::Recipient {
                 index,
-                message: "This recipient should not be handeled by this plugin".to_string(),
+                message: "This recipient should not be handled by this plugin".to_string(),
             });
         }
 
-        let pk: Result<&[u8; ENCAPSULATION_KEY_SIZE], TryFromSliceError> = bytes.try_into();
+        let pk: &[u8; ENCAPSULATION_KEY_SIZE] =
-        let pk = match pk {
+            bytes.try_into().map_err(|_| recipient::Error::Recipient {
-            Ok(x) => EncapsulationKey::from(x),
+                index,
-            Err(_) => {
+                message: "Invalid recipient".to_string(),
-                return Err(recipient::Error::Recipient {
+            })?;
-                    index,
+        let pk = EncapsulationKey::from(pk);
-                    message: "Invalid recipient".to_string(),
-                })
-            }
-        };
 
         self.recipients.push(pk);
 
@@ -103,20 +101,16 @@ impl RecipientPluginV1 for RecipientPlugin {
         if plugin_name != PLUGIN_NAME {
             return Err(recipient::Error::Identity {
                 index,
-                message: "This Identity should not be handeled by this plugin".to_owned(),
+                message: "This identity should not be handled by this plugin".to_string(),
             });
         }
 
-        let sk: Result<&[u8; DECAPSULATION_KEY_SIZE], TryFromSliceError> = bytes.try_into();
+        let sk: [u8; DECAPSULATION_KEY_SIZE] =
-        let sk = match sk {
+            bytes.try_into().map_err(|_| recipient::Error::Identity {
-            Ok(x) => DecapsulationKey::from(x.to_owned()),
+                index,
-            Err(_) => {
+                message: "Invalid identity".to_string(),
-                return Err(recipient::Error::Identity {
+            })?;
-                    index,
+        let sk = DecapsulationKey::from(sk);
-                    message: "Invalid identity".to_string(),
-                })
-            }
-        };
 
         self.recipients.push(sk.encapsulation_key());
 
@@ -145,103 +139,78 @@ struct IdentityPlugin {
 }
 
 impl IdentityPlugin {
-    fn decrypt_stanzas(
+    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: Vec<(usize, Stanza)>,
+        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() {
+            // Try to decrypt this stanza
-                Some(arg) => arg,
+            let result = (|| {
-                None => {
+                let arg = stanza.args.first().ok_or("Stanza is missing arguments")?;
-                    errors.push(identity::Error::Stanza {
-                        file_index,
-                        stanza_index,
-                        message: "Stanza is missing arguments".to_string(),
-                    });
-                    continue;
-                }
-            };
 
-            let ct = match BASE64_STANDARD.decode(arg) {
+                let ct = Self::parse_ciphertext(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() {
+                let ss = Self::try_decapsulate(&self.identities, &ct)
-                Ok(ct) => ct,
+                    .ok_or("No identity found that can decrypt the file")?;
-                Err(_) => {
-                    errors.push(identity::Error::Stanza {
-                        file_index,
-                        stanza_index,
-                        message: "Malformed ciphertext".to_string(),
-                    });
-                    continue;
-                }
-            };
 
-            let ct = Ciphertext::from(&ct);
+                Self::unwrap_file_key(&ss, &stanza.body)
+            })();
 
-            let ss = match self
+            match result {
-                .identities
+                Ok(file_key) => return Ok(file_key),
-                .iter()
+                Err(message) => errors.push(identity::Error::Stanza {
-                .filter_map(|key| key.decapsulate(&ct).ok())
+                    file_index,
-                .next()
+                    stanza_index,
-            {
+                    message,
-                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() {
+        // If we get here, no stanza was successfully decrypted
-            return Err(errors);
+        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 handeled by this plugin".to_string(),
-        }])
     }
 }
 
@@ -255,19 +224,15 @@ impl IdentityPluginV1 for IdentityPlugin {
         if plugin_name != PLUGIN_NAME {
             return Err(identity::Error::Identity {
                 index,
-                message: "This Identity should not be handeled by this plugin".to_string(),
+                message: "This identity should not be handled by this plugin".to_string(),
             });
         }
 
-        let bytes: [u8; DECAPSULATION_KEY_SIZE] = match bytes.try_into() {
+        let bytes: [u8; DECAPSULATION_KEY_SIZE] =
-            Ok(x) => x,
+            bytes.try_into().map_err(|_| identity::Error::Identity {
-            Err(_) => {
+                index,
-                return Err(identity::Error::Identity {
+                message: "Invalid identity".to_string(),
-                    index,
+            })?;
-                    message: "Invalid identity".to_string(),
-                })
-            }
-        };
 
         self.identities.push(DecapsulationKey::from(bytes));
 
@@ -285,8 +250,7 @@ impl IdentityPluginV1 for IdentityPlugin {
             let x_wing_stanzas = stanzas
                 .into_iter()
                 .enumerate()
-                .filter(|(_, stanza)| stanza.tag == PLUGIN_NAME)
+                .filter(|(_, stanza)| stanza.tag == PLUGIN_NAME);
-                .collect();
 
             ret.insert(file_index, self.decrypt_stanzas(file_index, x_wing_stanzas));
         }
@@ -312,7 +276,7 @@ fn main() -> io::Result<()> {
 
     // Here you can assume the binary is being run directly by a user, 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());
+    print_new_identity(PLUGIN_NAME, sk.as_bytes(), &pk.to_bytes());
 
     Ok(())
 }