From 2d2966871275b4b9eca8bee3300f0762d63f1039 Mon Sep 17 00:00:00 2001
From: Jack Grigg <thestr4d@gmail.com>
Date: Sun, 11 Aug 2024 22:12:26 +0000
Subject: [PATCH] age: Add labels extension to client side of `recipient-v1`

---
 age/CHANGELOG.md  |  1 +
 age/src/plugin.rs | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/age/CHANGELOG.md b/age/CHANGELOG.md
index e699e76..ab7b138 100644
--- a/age/CHANGELOG.md
+++ b/age/CHANGELOG.md
@@ -21,6 +21,7 @@ to 1.0.0 are beta releases.
 - `age::Recipient::wrap_file_key` now returns `(Vec<Stanza>, HashSet<String>)`:
   a tuple of the stanzas to be placed in an age file header, and labels that
   constrain how the stanzas may be combined with those from other recipients.
+- `age::plugin::RecipientPluginV1` now supports the labels extension.
 
 ### Removed
 - `age::decryptor::PassphraseDecryptor` (use `age::Decryptor` with
diff --git a/age/src/plugin.rs b/age/src/plugin.rs
index 4551574..933a67a 100644
--- a/age/src/plugin.rs
+++ b/age/src/plugin.rs
@@ -33,6 +33,7 @@ const PLUGIN_IDENTITY_PREFIX: &str = "age-plugin-";
 
 const CMD_ERROR: &str = "error";
 const CMD_RECIPIENT_STANZA: &str = "recipient-stanza";
+const CMD_LABELS: &str = "labels";
 const CMD_MSG: &str = "msg";
 const CMD_CONFIRM: &str = "confirm";
 const CMD_REQUEST_PUBLIC: &str = "request-public";
@@ -395,12 +396,13 @@ impl<C: Callbacks> crate::Recipient for RecipientPluginV1<C> {
             for identity in &self.identities {
                 phase.send("add-identity", &[&identity.identity], &[])?;
             }
+            phase.send("extension-labels", &[], &[])?;
             phase.send("wrap-file-key", &[], file_key.expose_secret())
         })?;
 
         // Phase 2: collect either stanzas or errors
         let mut stanzas = vec![];
-        let labels = HashSet::new();
+        let mut labels = None;
         let mut errors = vec![];
         if let Err(e) = conn.bidir_receive(
             &[
@@ -409,6 +411,7 @@ impl<C: Callbacks> crate::Recipient for RecipientPluginV1<C> {
                 CMD_REQUEST_PUBLIC,
                 CMD_REQUEST_SECRET,
                 CMD_RECIPIENT_STANZA,
+                CMD_LABELS,
                 CMD_ERROR,
             ],
             |mut command, reply| match command.tag.as_str() {
@@ -464,6 +467,34 @@ impl<C: Callbacks> crate::Recipient for RecipientPluginV1<C> {
                     }
                     reply.ok(None)
                 }
+                CMD_LABELS => {
+                    if labels.is_none() {
+                        let labels_count = command.args.len();
+                        let label_set = command.args.into_iter().collect::<HashSet<_>>();
+                        if label_set.len() == labels_count {
+                            labels = Some(label_set);
+                        } else {
+                            errors.push(PluginError::Other {
+                                kind: "internal".to_owned(),
+                                metadata: vec![],
+                                message: format!(
+                                    "{} command must not contain duplicate labels",
+                                    CMD_LABELS
+                                ),
+                            });
+                        }
+                    } else {
+                        errors.push(PluginError::Other {
+                            kind: "internal".to_owned(),
+                            metadata: vec![],
+                            message: format!(
+                                "{} command must not be sent more than once",
+                                CMD_LABELS
+                            ),
+                        });
+                    }
+                    reply.ok(None)
+                }
                 CMD_ERROR => {
                     if command.args.len() == 2 && command.args[0] == "recipient" {
                         let index: usize = command.args[1].parse().unwrap();
@@ -489,7 +520,7 @@ impl<C: Callbacks> crate::Recipient for RecipientPluginV1<C> {
             return Err(e.into());
         };
         match (stanzas.is_empty(), errors.is_empty()) {
-            (false, true) => Ok((stanzas, labels)),
+            (false, true) => Ok((stanzas, labels.unwrap_or_default())),
             (a, b) => {
                 if a & b {
                     errors.push(PluginError::Other {