Refactors for reusability across supported recipients

src/builder.rs

@@ -11,9 +11,9 @@ use crate::{
     error::Error,
     fl,
     key::{self, Stub},
-    p256::Recipient,
+    piv_p256,
     util::{Metadata, POLICY_EXTENSION_OID},
-    BINARY_NAME, USABLE_SLOTS,
+    Recipient, BINARY_NAME, USABLE_SLOTS,
 };
 
 pub(crate) const DEFAULT_PIN_POLICY: PinPolicy = PinPolicy::Once;
@@ -104,7 +104,9 @@ impl IdentityBuilder {
             touch_policy,
         )?;
 
-        let recipient = Recipient::from_spki(&generated).expect("YubiKey generates a valid pubkey");
+        let recipient = Recipient::PivP256(
+            piv_p256::Recipient::from_spki(&generated).expect("YubiKey generates a valid pubkey"),
+        );
         let stub = Stub::new(yubikey.serial(), slot, &recipient);
 
         eprintln!();

src/key.rs

@@ -20,11 +20,10 @@ use yubikey::{
 
 use crate::{
     error::Error,
-    fl,
+    fl, piv_p256,
-    p256::{Recipient, TAG_BYTES},
+    recipient::TAG_BYTES,
-    piv_p256,
     util::{otp_serial_prefix, Metadata},
-    IDENTITY_PREFIX,
+    Recipient, IDENTITY_PREFIX,
 };
 
 const ONE_SECOND: Duration = Duration::from_secs(1);
@@ -394,7 +393,8 @@ pub(crate) fn list_slots(
         match key.slot() {
             SlotId::Retired(slot) => {
                 // Only P-256 keys are compatible with us.
-                let recipient = Recipient::from_certificate(key.certificate());
+                let recipient = piv_p256::Recipient::from_certificate(key.certificate())
+                    .map(Recipient::PivP256);
                 Some((key, slot, recipient))
             }
             _ => None,
@@ -449,7 +449,7 @@ impl Stub {
         Stub {
             serial,
             slot,
-            tag: recipient.tag(),
+            tag: recipient.static_tag(),
             identity_index: 0,
         }
     }
@@ -476,10 +476,6 @@ impl Stub {
         bytes
     }
 
-    pub(crate) fn matches(&self, line: &piv_p256::RecipientLine) -> bool {
-        self.tag == line.tag
-    }
-
     /// Returns:
     /// - `Ok(Ok(Some(connection)))` if we successfully connected to this YubiKey.
     /// - `Ok(Ok(None))` if the user told us to skip this YubiKey.
@@ -601,9 +597,9 @@ impl Stub {
         let (cert, pk) = match Certificate::read(&mut yubikey, SlotId::Retired(self.slot))
             .ok()
             .and_then(|cert| {
-                Recipient::from_certificate(&cert)
+                piv_p256::Recipient::from_certificate(&cert)
                     .filter(|pk| pk.tag() == self.tag)
-                    .map(|pk| (cert, pk))
+                    .map(|pk| (cert, Recipient::PivP256(pk)))
             }) {
             Some(pk) => pk,
             None => {

src/main.rs

@@ -17,16 +17,17 @@ use yubikey::{piv::RetiredSlotId, reader::Context, PinPolicy, Serial, TouchPolic
 mod builder;
 mod error;
 mod key;
-mod p256;
 mod piv_p256;
 mod plugin;
 mod util;
 
+mod recipient;
+use recipient::Recipient;
+
 use error::Error;
 
 const PLUGIN_NAME: &str = "yubikey";
 const BINARY_NAME: &str = "age-plugin-yubikey";
-const RECIPIENT_PREFIX: &str = "age1yubikey";
 const IDENTITY_PREFIX: &str = "age-plugin-yubikey-";
 
 const USABLE_SLOTS: [RetiredSlotId; 20] = [
@@ -193,7 +194,7 @@ fn generate(flags: PluginFlags) -> Result<(), Error> {
 fn print_single(
     serial: Option<Serial>,
     slot: RetiredSlotId,
-    printer: impl Fn(key::Stub, p256::Recipient, util::Metadata),
+    printer: impl Fn(key::Stub, Recipient, util::Metadata),
 ) -> Result<(), Error> {
     let mut yubikey = key::open(serial)?;
 
@@ -215,7 +216,7 @@ fn print_multiple(
     kind: &str,
     serial: Option<Serial>,
     all: bool,
-    printer: impl Fn(key::Stub, p256::Recipient, util::Metadata),
+    printer: impl Fn(key::Stub, Recipient, util::Metadata),
 ) -> Result<(), Error> {
     let mut readers = Context::open()?;
 
@@ -255,7 +256,7 @@ fn print_details(
     kind: &str,
     flags: PluginFlags,
     all: bool,
-    printer: impl Fn(key::Stub, p256::Recipient, util::Metadata),
+    printer: impl Fn(key::Stub, Recipient, util::Metadata),
 ) -> Result<(), Error> {
     if let Some(slot) = flags.slot {
         print_single(flags.serial, slot, printer)

src/piv_p256.rs

@@ -11,12 +11,14 @@ use p256::{
 use rand::rngs::OsRng;
 use sha2::Sha256;
 
-use crate::{key::Connection, p256::Recipient};
+use crate::{key::Connection, recipient::TAG_BYTES, util::base64_arg};
+
+mod recipient;
+pub(crate) use recipient::Recipient;
 
 const STANZA_TAG: &str = "piv-p256";
 pub(crate) const STANZA_KEY_LABEL: &[u8] = b"piv-p256";
 
-const TAG_BYTES: usize = 4;
 const EPK_BYTES: usize = 33;
 const ENCRYPTED_FILE_KEY_BYTES: usize = 32;
 
@@ -81,17 +83,6 @@ impl RecipientLine {
             return None;
         }
 
-        fn base64_arg<A: AsRef<[u8]>, B: AsMut<[u8]>>(arg: &A, mut buf: B) -> Option<B> {
-            if arg.as_ref().len() != ((4 * buf.as_mut().len()) + 2) / 3 {
-                return None;
-            }
-
-            BASE64_STANDARD_NO_PAD
-                .decode_slice_unchecked(arg, buf.as_mut())
-                .ok()
-                .and_then(|len| (len == buf.as_mut().len()).then_some(buf))
-        }
-
         let (tag, epk_bytes) = match &s.args[..] {
             [tag, epk_bytes] => (
                 base64_arg(tag, [0; TAG_BYTES]),
@@ -110,15 +101,17 @@ impl RecipientLine {
             _ => Err(()),
         })
     }
+}
 
-    pub(crate) fn wrap_file_key(file_key: &FileKey, pk: &Recipient) -> Self {
+impl Recipient {
+    pub(crate) fn wrap_file_key(&self, file_key: &FileKey) -> RecipientLine {
         let esk = EphemeralSecret::random(&mut OsRng);
         let epk = esk.public_key();
         let epk_bytes = EphemeralKeyBytes::from_public_key(&epk);
 
-        let shared_secret = esk.diffie_hellman(pk.public_key());
+        let shared_secret = esk.diffie_hellman(self.public_key());
 
-        let salt = salt(&epk_bytes, pk);
+        let salt = salt(&epk_bytes, self);
 
         let enc_key = {
             let mut okm = [0; 32];
@@ -136,21 +129,24 @@ impl RecipientLine {
         };
 
         RecipientLine {
-            tag: pk.tag(),
+            tag: self.tag(),
             epk_bytes,
             encrypted_file_key,
         }
     }
+}
 
+impl RecipientLine {
     pub(crate) fn unwrap_file_key(&self, conn: &mut Connection) -> Result<FileKey, ()> {
-        assert_eq!(self.tag, conn.recipient().tag());
+        let crate::recipient::Recipient::PivP256(recipient) = conn.recipient();
+        assert_eq!(self.tag, recipient.tag());
+
+        let salt = salt(&self.epk_bytes, recipient);
 
         // The YubiKey API for performing scalar multiplication takes the point in its
         // uncompressed SEC-1 encoding.
         let shared_secret = conn.p256_ecdh(self.epk_bytes.decompress().as_bytes())?;
 
-        let salt = salt(&self.epk_bytes, conn.recipient());
-
         let enc_key = hkdf(&salt, STANZA_KEY_LABEL, shared_secret.as_ref());
 
         // A failure to decrypt is fatal, because we assume that we won't

src/piv_p256/recipient.rs

@@ -1,13 +1,12 @@
 use bech32::{ToBase32, Variant};
 use p256::elliptic_curve::sec1::{FromEncodedPoint, ToEncodedPoint};
-use sha2::{Digest, Sha256};
 use yubikey::{certificate::PublicKeyInfo, Certificate};
 
 use std::fmt;
 
-use crate::RECIPIENT_PREFIX;
+use crate::recipient::{static_tag, TAG_BYTES};
 
-pub(crate) const TAG_BYTES: usize = 4;
+const RECIPIENT_PREFIX: &str = "age1yubikey";
 
 /// Wrapper around a compressed secp256r1 curve point.
 #[derive(Clone)]
@@ -69,8 +68,7 @@ impl Recipient {
     }
 
     pub(crate) fn tag(&self) -> [u8; TAG_BYTES] {
-        let tag = Sha256::digest(self.to_encoded().as_bytes());
+        static_tag(self.to_encoded().as_bytes())
-        (&tag[0..TAG_BYTES]).try_into().expect("length is correct")
     }
 
     /// Exposes the wrapped public key.

src/plugin.rs

@@ -7,7 +7,7 @@ use age_plugin::{
 use std::collections::{HashMap, HashSet};
 use std::io;
 
-use crate::{fl, key, p256::Recipient, piv_p256, PLUGIN_NAME};
+use crate::{fl, key, piv_p256, Recipient, PLUGIN_NAME};
 
 pub(crate) struct Handler;
 
@@ -37,11 +37,7 @@ impl RecipientPluginV1 for RecipientPlugin {
         plugin_name: &str,
         bytes: &[u8],
     ) -> Result<(), recipient::Error> {
-        if let Some(pk) = if plugin_name == PLUGIN_NAME {
+        if let Some(pk) = Recipient::from_bytes(plugin_name, bytes) {
-            Recipient::from_bytes(bytes)
-        } else {
-            None
-        } {
             self.recipients.push(pk);
             Ok(())
         } else {
@@ -114,7 +110,7 @@ impl RecipientPluginV1 for RecipientPlugin {
                     self.recipients
                         .iter()
                         .chain(yk_recipients.iter())
-                        .map(|pk| piv_p256::RecipientLine::wrap_file_key(&file_key, pk).into())
+                        .map(|pk| pk.wrap_file_key(&file_key))
                         .collect()
                 })
                 .collect())
@@ -159,16 +155,16 @@ impl IdentityPluginV1 for IdentityPlugin {
         let mut file_keys = HashMap::with_capacity(files.len());
 
         // Filter to files / stanzas for which we have matching YubiKeys
-        let mut candidate_stanzas: Vec<(&key::Stub, HashMap<usize, Vec<piv_p256::RecipientLine>>)> =
+        let mut candidate_stanzas: Vec<(&key::Stub, HashMap<usize, Vec<SupportedStanza>>)> = self
-            self.yubikeys
+            .yubikeys
             .iter()
             .map(|stub| (stub, HashMap::new()))
             .collect();
 
-        for (file, stanzas) in files.iter().enumerate() {
+        for (file, stanzas) in files.into_iter().enumerate() {
-            for (stanza_index, stanza) in stanzas.iter().enumerate() {
+            for (stanza_index, stanza) in stanzas.into_iter().enumerate() {
                 match (
-                    piv_p256::RecipientLine::from_stanza(stanza).map(|res| {
+                    SupportedStanza::parse(stanza).map(|res| {
                         res.map_err(|_| identity::Error::Stanza {
                             file_index: file,
                             stanza_index,
@@ -182,7 +178,7 @@ impl IdentityPluginV1 for IdentityPlugin {
                         // A line will match at most one YubiKey.
                         if let Some(files) =
                             candidate_stanzas.iter_mut().find_map(|(stub, files)| {
-                                if stub.matches(&line) {
+                                if line.matches_stub(stub) {
                                     Some(files)
                                 } else {
                                     None
@@ -274,3 +270,25 @@ impl IdentityPluginV1 for IdentityPlugin {
         Ok(file_keys)
     }
 }
+
+enum SupportedStanza {
+    PivP256(piv_p256::RecipientLine),
+}
+
+impl SupportedStanza {
+    fn parse(stanza: Stanza) -> Option<Result<Self, ()>> {
+        piv_p256::RecipientLine::from_stanza(&stanza).map(|res| res.map(Self::PivP256))
+    }
+
+    pub(crate) fn matches_stub(&self, stub: &key::Stub) -> bool {
+        match self {
+            SupportedStanza::PivP256(line) => stub.tag == line.tag,
+        }
+    }
+
+    pub(crate) fn unwrap_file_key(&self, conn: &mut key::Connection) -> Result<FileKey, ()> {
+        match self {
+            SupportedStanza::PivP256(line) => line.unwrap_file_key(conn),
+        }
+    }
+}

src/recipient.rs

@@ -0,0 +1,50 @@
+use std::fmt;
+
+use age_core::format::{FileKey, Stanza};
+use sha2::{Digest, Sha256};
+
+use crate::{piv_p256, PLUGIN_NAME};
+
+pub(crate) const TAG_BYTES: usize = 4;
+
+#[derive(Clone, Debug)]
+pub(crate) enum Recipient {
+    PivP256(piv_p256::Recipient),
+}
+
+impl fmt::Display for Recipient {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Recipient::PivP256(recipient) => recipient.fmt(f),
+        }
+    }
+}
+
+impl Recipient {
+    /// Attempts to parse a supported YubiKey recipient.
+    pub(crate) fn from_bytes(plugin_name: &str, bytes: &[u8]) -> Option<Self> {
+        match plugin_name {
+            PLUGIN_NAME => piv_p256::Recipient::from_bytes(bytes).map(Self::PivP256),
+            _ => None,
+        }
+    }
+
+    /// Returns the static tag for this recipient.
+    pub(crate) fn static_tag(&self) -> [u8; TAG_BYTES] {
+        match self {
+            Recipient::PivP256(recipient) => recipient.tag(),
+        }
+    }
+
+    pub(crate) fn wrap_file_key(&self, file_key: &FileKey) -> Stanza {
+        match self {
+            Recipient::PivP256(recipient) => recipient.wrap_file_key(file_key).into(),
+        }
+    }
+}
+
+pub(crate) fn static_tag(pk: &[u8]) -> [u8; TAG_BYTES] {
+    Sha256::digest(pk)[0..TAG_BYTES]
+        .try_into()
+        .expect("length is correct")
+}

src/util.rs

@@ -1,6 +1,7 @@
 use std::fmt;
 use std::iter;
 
+use base64::{prelude::BASE64_STANDARD_NO_PAD, Engine};
 use x509_parser::{certificate::X509Certificate, der_parser::oid::Oid};
 use yubikey::{
     piv::{RetiredSlotId, SlotId},
@@ -8,7 +9,7 @@ use yubikey::{
 };
 
 use crate::fl;
-use crate::{error::Error, key::Stub, p256::Recipient, BINARY_NAME, USABLE_SLOTS};
+use crate::{error::Error, key::Stub, Recipient, BINARY_NAME, USABLE_SLOTS};
 
 pub(crate) const POLICY_EXTENSION_OID: &[u64] = &[1, 3, 6, 1, 4, 1, 41482, 3, 8];
 
@@ -219,3 +220,14 @@ pub(crate) fn print_identity(stub: Stub, recipient: Recipient, metadata: Metadat
         )
     );
 }
+
+pub(crate) fn base64_arg<A: AsRef<[u8]>, B: AsMut<[u8]>>(arg: &A, mut buf: B) -> Option<B> {
+    if arg.as_ref().len() != ((4 * buf.as_mut().len()) + 2) / 3 {
+        return None;
+    }
+
+    BASE64_STANDARD_NO_PAD
+        .decode_slice_unchecked(arg, buf.as_mut())
+        .ok()
+        .and_then(|len| (len == buf.as_mut().len()).then_some(buf))
+}