From 9d818f164b51a2cd9e8fb39b9ce65104512d9223 Mon Sep 17 00:00:00 2001
From: Gabriela Moldovan <gabi@torproject.org>
Date: Wed, 19 Jul 2023 15:47:50 +0100
Subject: [PATCH] keymgr: Require callers to be explicit about where to remove
 keys from.

As with `KeyMgr::insert`, only `KeystoreSelector::Id` and
`KeystoreSelector::Default` are supported.
---
 crates/tor-keymgr/semver.md  |  2 +-
 crates/tor-keymgr/src/mgr.rs | 81 +++++++++++++++++++++++++++++-------
 2 files changed, 67 insertions(+), 16 deletions(-)

diff --git a/crates/tor-keymgr/semver.md b/crates/tor-keymgr/semver.md
index 96fea116d..56b71d9bd 100644
--- a/crates/tor-keymgr/semver.md
+++ b/crates/tor-keymgr/semver.md
@@ -7,7 +7,7 @@ BREAKING: `KeyMgr::new` takes an extra argument (the default keystore)
 ADDED: `KeystoreSelector`
 BREAKING: `KeystoreError` now has a `boxed` function
 ADDED: an `id` function to `Keystore` trait
-BREAKING: `KeyMgr::insert`, `KeyMgr::get` now take an additional
+BREAKING: `KeyMgr::insert`, `KeyMgr::get`, `KeyMgr::remove` now take an
 `selector: KeystoreSelector` argument
 REMOVED: the `has_key_bundle` function (from the `Keystore` trait)
 ADDED: `PartialEq`, `Eq`, `Hash` derives for `ArtiPath` and `KeyType`
diff --git a/crates/tor-keymgr/src/mgr.rs b/crates/tor-keymgr/src/mgr.rs
index 02948bdfa..85a2a1429 100644
--- a/crates/tor-keymgr/src/mgr.rs
+++ b/crates/tor-keymgr/src/mgr.rs
@@ -120,27 +120,35 @@ impl KeyMgr {
         }
     }
 
-    /// Remove the specified key.
+    /// Remove the key identified by `key_spec` from the [`Keystore`] specified by `selector`.
     ///
-    /// If the key exists in multiple key stores, this will only remove it from the first one.
+    /// This function can only be used with [`KeystoreSelector::Default`] and
+    /// [`KeystoreSelector::Id`]. It returns an error if the specified `selector`
+    /// is not supported.
     ///
-    /// A return vaue of `Ok(None)` indicates the key doesn't exist in any of the key stores,
-    /// whereas `Ok(Some(())` means the key was successfully removed.
+    /// Returns `Ok(None)` if the key does not exist in the requested keystore.
+    /// Returns `Ok(Some(())` if the key was successfully removed.
     ///
     /// Returns `Err` if an error occurred while trying to remove the key.
-    pub fn remove<K: ToEncodableKey>(&self, key_spec: &dyn KeySpecifier) -> Result<Option<()>> {
-        for store in &self.key_stores {
-            match store.remove(key_spec, K::Key::key_type()) {
-                Ok(None) => {
-                    // This key store doesn't have the key we're trying to remove, so we search the
-                    // next key store...
-                    continue;
-                }
-                res => return res,
+    pub fn remove<K: ToEncodableKey>(
+        &self,
+        key_spec: &dyn KeySpecifier,
+        selector: KeystoreSelector,
+    ) -> Result<Option<()>> {
+        match selector {
+            KeystoreSelector::Id(keystore_id) => {
+                let Some(keystore) = self.find_keystore(keystore_id) else {
+                    return Err(KeyMgrError::KeystoreNotFound(keystore_id).boxed());
+                };
+                keystore.remove(key_spec, K::Key::key_type())
             }
+            KeystoreSelector::Default => self.default_store.remove(key_spec, K::Key::key_type()),
+            KeystoreSelector::All => Err(KeyMgrError::UnsupportedKeystoreSelector {
+                op: "remove",
+                selector,
+            }
+            .boxed()),
         }
-
-        Ok(None)
     }
 
     /// Attempt to retrieve a key from one of the specified `stores`.
@@ -400,4 +408,47 @@ mod tests {
             Some("keystore1_cormorant".to_string())
         );
     }
+
+    #[test]
+    fn remove() {
+        let mgr = KeyMgr::new(
+            Keystore1::default(),
+            vec![Keystore2::new_boxed(), Keystore3::new_boxed()],
+        );
+
+        // Insert a key into Keystore2
+        mgr.insert(
+            "coot".to_string(),
+            &TestKeySpecifier1,
+            KeystoreSelector::Id("keystore2"),
+        )
+        .unwrap();
+        assert_eq!(
+            mgr.get::<TestKey>(&TestKeySpecifier1, KeystoreSelector::All)
+                .unwrap(),
+            Some("keystore2_coot".to_string())
+        );
+
+        // Try to remove the key from a non-existent key store
+        assert!(mgr
+            .remove::<TestKey>(
+                &TestKeySpecifier1,
+                KeystoreSelector::Id("not_an_id_we_know_of")
+            )
+            .is_err());
+
+        // Try to remove the key from the default key store
+        assert_eq!(
+            mgr.remove::<TestKey>(&TestKeySpecifier1, KeystoreSelector::Default)
+                .unwrap(),
+            None
+        );
+
+        // Removing from Keystore2 should succeed.
+        assert_eq!(
+            mgr.remove::<TestKey>(&TestKeySpecifier1, KeystoreSelector::Id("keystore2"))
+                .unwrap(),
+            Some(())
+        );
+    }
 }