From f96298a791f696012c572a1df8396eea46d43266 Mon Sep 17 00:00:00 2001
From: Gabriela Moldovan <gabi@torproject.org>
Date: Mon, 24 Jul 2023 11:56:18 +0100
Subject: [PATCH] keymgr: Add KeyMgr::generate() for generating new keys.

---
 Cargo.lock                   |  1 +
 crates/tor-keymgr/Cargo.toml |  1 +
 crates/tor-keymgr/semver.md  |  1 +
 crates/tor-keymgr/src/mgr.rs | 70 +++++++++++++++++++++++++++++++++++-
 4 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index ae25533b6..248101492 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4567,6 +4567,7 @@ dependencies = [
  "ssh-key",
  "tempfile",
  "thiserror",
+ "tor-basic-utils",
  "tor-config",
  "tor-error",
  "tor-hscrypto",
diff --git a/crates/tor-keymgr/Cargo.toml b/crates/tor-keymgr/Cargo.toml
index ce1ea6e62..c7411f007 100644
--- a/crates/tor-keymgr/Cargo.toml
+++ b/crates/tor-keymgr/Cargo.toml
@@ -42,6 +42,7 @@ zeroize = "1"
 
 [dev-dependencies]
 tempfile = "3"
+tor-basic-utils = { path = "../tor-basic-utils", version = "0.7.1" }
 
 [package.metadata.docs.rs]
 all-features = true
diff --git a/crates/tor-keymgr/semver.md b/crates/tor-keymgr/semver.md
index ce59eaa7e..275eba5c7 100644
--- a/crates/tor-keymgr/semver.md
+++ b/crates/tor-keymgr/semver.md
@@ -15,3 +15,4 @@ ADDED: a `to_bytes` function to `EncodableKey` trait
 ADDED: `Keystore::contains()`
 ADDED: `KeygenRng` trait
 ADDED: `EncodableKey::generate()`
+ADDED: `KeyMgr::generate()`
diff --git a/crates/tor-keymgr/src/mgr.rs b/crates/tor-keymgr/src/mgr.rs
index a1675d581..d38058d7a 100644
--- a/crates/tor-keymgr/src/mgr.rs
+++ b/crates/tor-keymgr/src/mgr.rs
@@ -4,7 +4,8 @@
 //! [`Keystore`].
 
 use crate::{
-    EncodableKey, KeySpecifier, Keystore, KeystoreId, KeystoreSelector, Result, ToEncodableKey,
+    EncodableKey, KeySpecifier, KeygenRng, Keystore, KeystoreId, KeystoreSelector, Result,
+    ToEncodableKey,
 };
 
 use std::iter;
@@ -45,6 +46,32 @@ impl KeyMgr {
         self.get_from_store(key_spec, self.all_stores())
     }
 
+    /// Generate a new key of type `K`, and insert it into the key store specified by `selector`.
+    ///
+    /// If the key already exists in the specified key store, the `overwrite` flag is used to
+    /// decide whether to overwrite it with a newly generated key.
+    pub fn generate<K: ToEncodableKey>(
+        &self,
+        key_spec: &dyn KeySpecifier,
+        selector: KeystoreSelector,
+        rng: &mut dyn KeygenRng,
+        overwrite: bool,
+    ) -> Result<()> {
+        let store = match selector {
+            KeystoreSelector::Id(keystore_id) => self.find_keystore(keystore_id)?,
+            KeystoreSelector::Default => &self.default_store,
+        };
+
+        let key_type = K::Key::key_type();
+
+        if overwrite || !store.contains(key_spec, key_type)? {
+            let key = K::Key::generate(rng);
+            store.insert(&key, key_spec, key_type)
+        } else {
+            Ok(())
+        }
+    }
+
     /// Insert `key` into the [`Keystore`] specified by `selector`.
     ///
     /// If the key already exists, it is overwritten.
@@ -156,6 +183,7 @@ mod tests {
     use std::collections::HashMap;
     use std::str::FromStr;
     use std::sync::RwLock;
+    use tor_basic_utils::test_rng::testing_rng;
 
     /// The type of "key" stored in the test key stores.
     type TestKey = String;
@@ -433,4 +461,44 @@ mod tests {
             .contains(&TestKeySpecifier1, TestKey::key_type())
             .unwrap());
     }
+
+    #[test]
+    fn keygen() {
+        let mgr = KeyMgr::new(Keystore1::default(), vec![]);
+
+        mgr.insert(
+            "coot".to_string(),
+            &TestKeySpecifier1,
+            KeystoreSelector::Default,
+        )
+        .unwrap();
+
+        // Try to generate a new key (overwrite = false)
+        mgr.generate::<TestKey>(
+            &TestKeySpecifier1,
+            KeystoreSelector::Default,
+            &mut testing_rng(),
+            false,
+        )
+        .unwrap();
+
+        assert_eq!(
+            mgr.get::<TestKey>(&TestKeySpecifier1).unwrap(),
+            Some("keystore1_coot".to_string())
+        );
+
+        // Try to generate a new key (overwrite = true)
+        mgr.generate::<TestKey>(
+            &TestKeySpecifier1,
+            KeystoreSelector::Default,
+            &mut testing_rng(),
+            true,
+        )
+        .unwrap();
+
+        assert_eq!(
+            mgr.get::<TestKey>(&TestKeySpecifier1).unwrap(),
+            Some("keystore1_generated_test_key".to_string())
+        );
+    }
 }