age-plugin: Add labels extension to recipient-v1

age-core/CHANGELOG.md

@@ -10,6 +10,10 @@ to 1.0.0 are beta releases.
 ### Added
 - `age_core::format::is_arbitrary_string`
 
+### Changed
+- `age::plugin::Connection::unidir_receive` now takes an additional argument to
+  enable handling an optional fourth command.
+
 ## [0.10.0] - 2024-02-04
 ### Added
 - `impl Eq for age_core::format::Stanza`

age-core/src/plugin.rs

@@ -51,10 +51,11 @@ impl std::error::Error for Error {}
 ///   should explicitly handle.
 pub type Result<T> = io::Result<std::result::Result<T, Error>>;
 
-type UnidirResult<A, B, C, E> = io::Result<(
+type UnidirResult<A, B, C, D, E> = io::Result<(
     std::result::Result<Vec<A>, Vec<E>>,
     std::result::Result<Vec<B>, Vec<E>>,
     Option<std::result::Result<Vec<C>, Vec<E>>>,
+    Option<std::result::Result<Vec<D>, Vec<E>>>,
 )>;
 
 /// A connection to a plugin binary.
@@ -205,23 +206,26 @@ impl<R: Read, W: Write> Connection<R, W> {
     ///
     /// # Arguments
     ///
-    /// `command_a`, `command_b`, and (optionally) `command_c` are the known commands that
-    /// are expected to be received. All other received commands (including grease) will
-    /// be ignored.
-    pub fn unidir_receive<A, B, C, E, F, G, H>(
+    /// `command_a`, `command_b`, and (optionally) `command_c` and `command_d` are the
+    /// known commands that are expected to be received. All other received commands
+    /// (including grease) will be ignored.
+    pub fn unidir_receive<A, B, C, D, E, F, G, H, I>(
         &mut self,
         command_a: (&str, F),
         command_b: (&str, G),
         command_c: (Option<&str>, H),
-    ) -> UnidirResult<A, B, C, E>
+        command_d: (Option<&str>, I),
+    ) -> UnidirResult<A, B, C, D, E>
     where
         F: Fn(Stanza) -> std::result::Result<A, E>,
         G: Fn(Stanza) -> std::result::Result<B, E>,
         H: Fn(Stanza) -> std::result::Result<C, E>,
+        I: Fn(Stanza) -> std::result::Result<D, E>,
     {
         let mut res_a = Ok(vec![]);
         let mut res_b = Ok(vec![]);
         let mut res_c = Ok(vec![]);
+        let mut res_d = Ok(vec![]);
 
         for stanza in iter::repeat_with(|| self.receive()).take_while(|res| match res {
             Ok(stanza) => stanza.tag != COMMAND_DONE,
@@ -251,14 +255,28 @@ impl<R: Read, W: Write> Connection<R, W> {
                 validate(command_a.1(stanza), &mut res_a)
             } else if stanza.tag.as_str() == command_b.0 {
                 validate(command_b.1(stanza), &mut res_b)
-            } else if let Some(tag) = command_c.0 {
-                if stanza.tag.as_str() == tag {
-                    validate(command_c.1(stanza), &mut res_c)
+            } else {
+                if let Some(tag) = command_c.0 {
+                    if stanza.tag.as_str() == tag {
+                        validate(command_c.1(stanza), &mut res_c);
+                        continue;
+                    }
+                }
+                if let Some(tag) = command_d.0 {
+                    if stanza.tag.as_str() == tag {
+                        validate(command_d.1(stanza), &mut res_d);
+                        continue;
+                    }
                 }
             }
         }
 
-        Ok((res_a, res_b, command_c.0.map(|_| res_c)))
+        Ok((
+            res_a,
+            res_b,
+            command_c.0.map(|_| res_c),
+            command_d.0.map(|_| res_d),
+        ))
     }
 
     /// Runs a bidirectional phase as the controller.
@@ -481,10 +499,11 @@ mod tests {
             .unidir_send(|mut phase| phase.send("test", &["foo"], b"bar"))
             .unwrap();
         let stanza = plugin_conn
-            .unidir_receive::<_, (), (), _, _, _, _>(
+            .unidir_receive::<_, (), (), (), _, _, _, _, _>(
                 ("test", Ok),
                 ("other", |_| Err(())),
                 (None, |_| Ok(())),
+                (None, |_| Ok(())),
             )
             .unwrap();
         assert_eq!(
@@ -496,7 +515,8 @@ mod tests {
                     body: b"bar"[..].to_owned()
                 }]),
                 Ok(vec![]),
-                None
+                None,
+                None,
             )
         );
     }

age-plugin/CHANGELOG.md

@@ -14,6 +14,13 @@ to 1.0.0 are beta releases.
 - `impl age_plugin::identity::IdentityPluginV1 for std::convert::Infallible`
 - `impl age_plugin::recipient::RecipientPluginV1 for std::convert::Infallible`
 
+### Changed
+- `age_plugin::recipient::RecipientPluginV1` has a new `labels` method. Existing
+  implementations of the trait should either return `HashSet::new()` to maintain
+  existing compatibility, or return labels that apply the desired constraints.
+- `age_plugin::run_state_machine` now supports the `recipient-v1` labels
+  extension.
+
 ### Fixed
 - `age_plugin::run_state_machine` now takes an `impl age_plugin::PluginHandler`
   argument, instead of its previous arguments.

age-plugin/examples/age-plugin-unencrypted.rs

@@ -10,7 +10,7 @@ use age_plugin::{
 };
 use clap::Parser;
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::convert::Infallible;
 use std::env;
 use std::io;
@@ -104,6 +104,16 @@ impl RecipientPluginV1 for RecipientPlugin {
         }
     }
 
+    fn labels(&mut self) -> HashSet<String> {
+        let mut labels = HashSet::new();
+        if let Ok(s) = env::var("AGE_PLUGIN_LABELS") {
+            for label in s.split(',') {
+                labels.insert(label.into());
+            }
+        }
+        labels
+    }
+
     fn wrap_file_keys(
         &mut self,
         file_keys: Vec<FileKey>,

age-plugin/src/identity.rs

@@ -222,7 +222,7 @@ pub(crate) fn run_v1<P: IdentityPluginV1>(mut plugin: P) -> io::Result<()> {
 
     // Phase 1: receive identities and stanzas
     let (identities, recipient_stanzas) = {
-        let (identities, stanzas, _) = conn.unidir_receive(
+        let (identities, stanzas, _, _) = conn.unidir_receive(
             (ADD_IDENTITY, |s| match (&s.args[..], &s.body[..]) {
                 ([identity], []) => Ok(identity.clone()),
                 _ => Err(Error::Internal {
@@ -255,6 +255,7 @@ pub(crate) fn run_v1<P: IdentityPluginV1>(mut plugin: P) -> io::Result<()> {
                 }
             }),
             (None, |_| Ok(())),
+            (None, |_| Ok(())),
         )?;
 
         // Now that we have the full list of identities, parse them as Bech32 and add them

age-plugin/src/lib.rs

@@ -78,7 +78,7 @@
 //! };
 //! use clap::Parser;
 //!
-//! use std::collections::HashMap;
+//! use std::collections::{HashMap, HashSet};
 //! use std::io;
 //!
 //! struct Handler;
@@ -117,6 +117,10 @@
 //!         todo!()
 //!     }
 //!
+//!     fn labels(&mut self) -> HashSet<String> {
+//!         todo!()
+//!     }
+//!
 //!     fn wrap_file_keys(
 //!         &mut self,
 //!         file_keys: Vec<FileKey>,

age-plugin/src/recipient.rs

@@ -1,13 +1,14 @@
 //! Recipient plugin helpers.
 
 use age_core::{
-    format::{FileKey, Stanza, FILE_KEY_BYTES},
+    format::{is_arbitrary_string, FileKey, Stanza, FILE_KEY_BYTES},
     plugin::{self, BidirSend, Connection},
     secrecy::SecretString,
 };
 use base64::{prelude::BASE64_STANDARD_NO_PAD, Engine};
 use bech32::FromBase32;
 
+use std::collections::HashSet;
 use std::convert::Infallible;
 use std::io;
 
@@ -16,7 +17,9 @@ use crate::{Callbacks, PLUGIN_IDENTITY_PREFIX, PLUGIN_RECIPIENT_PREFIX};
 const ADD_RECIPIENT: &str = "add-recipient";
 const ADD_IDENTITY: &str = "add-identity";
 const WRAP_FILE_KEY: &str = "wrap-file-key";
+const EXTENSION_LABELS: &str = "extension-labels";
 const RECIPIENT_STANZA: &str = "recipient-stanza";
+const LABELS: &str = "labels";
 
 /// The interface that age implementations will use to interact with an age plugin.
 ///
@@ -39,6 +42,36 @@ pub trait RecipientPluginV1 {
     /// Returns an error if the identity is unknown or invalid.
     fn add_identity(&mut self, index: usize, plugin_name: &str, bytes: &[u8]) -> Result<(), Error>;
 
+    /// Returns labels that constrain how the stanzas produced by [`Self::wrap_file_keys`]
+    /// may be combined with those from other recipients.
+    ///
+    /// Encryption will succeed only if every recipient returns the same set of labels.
+    /// Subsets or partial overlapping sets are not allowed; all sets must be identical.
+    /// Labels are compared exactly, and are case-sensitive.
+    ///
+    /// Label sets can be used to ensure a recipient is only encrypted to alongside other
+    /// recipients with equivalent properties, or to ensure a recipient is always used
+    /// alone. A recipient with no particular properties to enforce should return an empty
+    /// label set.
+    ///
+    /// Labels can have any value that is a valid arbitrary string (`1*VCHAR` in ABNF),
+    /// but usually take one of several forms:
+    ///   - *Common public label* - used by multiple recipients to permit their stanzas to
+    ///     be used only together. Examples include:
+    ///     - `postquantum` - indicates that the recipient stanzas being generated are
+    ///       postquantum-secure, and that they can only be combined with other stanzas
+    ///       that are also postquantum-secure.
+    ///   - *Common private label* - used by recipients created by the same private entity
+    ///     to permit their recipient stanzas to be used only together. For example,
+    ///     private recipients used in a corporate environment could all send the same
+    ///     private label in order to prevent compliant age clients from simultaneously
+    ///     wrapping file keys with other recipients.
+    ///   - *Random label* - used by recipients that want to ensure their stanzas are not
+    ///     used with any other recipient stanzas. This can be used to produce a file key
+    ///     that is only encrypted to a single recipient stanza, for example to preserve
+    ///     its authentication properties.
+    fn labels(&mut self) -> HashSet<String>;
+
     /// Wraps each `file_key` to all recipients and identities previously added via
     /// `add_recipient` and `add_identity`.
     ///
@@ -65,6 +98,11 @@ impl RecipientPluginV1 for Infallible {
         Ok(())
     }
 
+    fn labels(&mut self) -> HashSet<String> {
+        // This is never executed.
+        HashSet::new()
+    }
+
     fn wrap_file_keys(
         &mut self,
         _: Vec<FileKey>,
@@ -215,8 +253,8 @@ pub(crate) fn run_v1<P: RecipientPluginV1>(mut plugin: P) -> io::Result<()> {
     let mut conn = Connection::accept();
 
     // Phase 1: collect recipients, and file keys to be wrapped
-    let ((recipients, identities), file_keys) = {
-        let (recipients, identities, file_keys) = conn.unidir_receive(
+    let ((recipients, identities), file_keys, labels_supported) = {
+        let (recipients, identities, file_keys, labels_supported) = conn.unidir_receive(
             (ADD_RECIPIENT, |s| match (&s.args[..], &s.body[..]) {
                 ([recipient], []) => Ok(recipient.clone()),
                 _ => Err(Error::Internal {
@@ -243,6 +281,7 @@ pub(crate) fn run_v1<P: RecipientPluginV1>(mut plugin: P) -> io::Result<()> {
                     })
                     .map(FileKey::from)
             }),
+            (Some(EXTENSION_LABELS), |_| Ok(())),
         )?;
         (
             match (recipients, identities) {
@@ -263,6 +302,13 @@ pub(crate) fn run_v1<P: RecipientPluginV1>(mut plugin: P) -> io::Result<()> {
                 }]),
                 r => r,
             },
+            match &labels_supported.unwrap() {
+                Ok(v) if v.is_empty() => Ok(false),
+                Ok(v) if v.len() == 1 => Ok(true),
+                _ => Err(vec![Error::Internal {
+                    message: format!("Received more than one {} command", EXTENSION_LABELS),
+                }]),
+            },
         )
     };
 
@@ -327,23 +373,58 @@ pub(crate) fn run_v1<P: RecipientPluginV1>(mut plugin: P) -> io::Result<()> {
         |index, plugin_name, bytes| plugin.add_identity(index, plugin_name, &bytes),
     );
 
+    let required_labels = plugin.labels();
+
+    let labels = match (labels_supported, required_labels.is_empty()) {
+        (Ok(true), _) | (Ok(false), true) => {
+            if required_labels.contains("") {
+                Err(vec![Error::Internal {
+                    message: "Plugin tried to use the empty string as a label".into(),
+                }])
+            } else if required_labels.iter().all(is_arbitrary_string) {
+                Ok(required_labels)
+            } else {
+                Err(vec![Error::Internal {
+                    message: "Plugin tried to use a label containing an invalid character".into(),
+                }])
+            }
+        }
+        (Ok(false), false) => Err(vec![Error::Internal {
+            message: "Plugin requires labels but client does not support them".into(),
+        }]),
+        (Err(errors), true) => Err(errors),
+        (Err(mut errors), false) => {
+            errors.push(Error::Internal {
+                message: "Plugin requires labels but client does not support them".into(),
+            });
+            Err(errors)
+        }
+    };
+
     // Phase 2: wrap the file keys or return errors
     conn.bidir_send(|mut phase| {
-        let (expected_stanzas, file_keys) = match (recipients, identities, file_keys) {
-            (Ok(recipients), Ok(identities), Ok(file_keys)) => (recipients + identities, file_keys),
-            (recipients, identities, file_keys) => {
-                for error in recipients
-                    .err()
-                    .into_iter()
-                    .chain(identities.err())
-                    .chain(file_keys.err())
-                    .flatten()
-                {
-                    error.send(&mut phase)?;
+        let (expected_stanzas, file_keys, labels) =
+            match (recipients, identities, file_keys, labels) {
+                (Ok(recipients), Ok(identities), Ok(file_keys), Ok(labels)) => {
+                    (recipients + identities, file_keys, labels)
                 }
-                return Ok(());
-            }
-        };
+                (recipients, identities, file_keys, labels) => {
+                    for error in recipients
+                        .err()
+                        .into_iter()
+                        .chain(identities.err())
+                        .chain(file_keys.err())
+                        .chain(labels.err())
+                        .flatten()
+                    {
+                        error.send(&mut phase)?;
+                    }
+                    return Ok(());
+                }
+            };
+
+        let labels = labels.iter().map(|s| s.as_str()).collect::<Vec<_>>();
+        phase.send(LABELS, &labels, &[])?.unwrap();
 
         match plugin.wrap_file_keys(file_keys, BidirCallbacks(&mut phase))? {
             Ok(files) => {