From a75e99d59f246784127066d940dcb4e2d06010b4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bogdan-=C8=98tefan=20Neac=C5=9Fu?= <bogdan@nymtech.net>
Date: Tue, 21 Jan 2025 14:18:55 +0100
Subject: [PATCH] Remove empty ephemeral keys

---
 .../src/client/base_client/storage/mod.rs     | 15 ++++++++++----
 .../src/client/key_manager/persistence.rs     | 20 +++++++++++++++----
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/common/client-core/src/client/base_client/storage/mod.rs b/common/client-core/src/client/base_client/storage/mod.rs
index 5129a5fc41..5f24c8ef7f 100644
--- a/common/client-core/src/client/base_client/storage/mod.rs
+++ b/common/client-core/src/client/base_client/storage/mod.rs
@@ -4,6 +4,8 @@
 // TODO: combine those more closely. Perhaps into a single underlying store.
 // Like for persistent, on-disk, storage, what's the point of having 3 different databases?
 
+use rand::rngs::OsRng;
+
 use crate::client::key_manager::persistence::{InMemEphemeralKeys, KeyStore};
 use crate::client::replies::reply_storage;
 use crate::client::replies::reply_storage::ReplyStorageBackend;
@@ -63,7 +65,7 @@ pub trait MixnetClientStorage {
     fn gateway_details_store(&self) -> &Self::GatewaysDetailsStore;
 }
 
-#[derive(Clone, Default)]
+#[derive(Clone)]
 pub struct Ephemeral {
     key_store: InMemEphemeralKeys,
     reply_store: reply_storage::Empty,
@@ -71,9 +73,14 @@ pub struct Ephemeral {
     gateway_details_store: InMemGatewaysDetails,
 }
 
-impl Ephemeral {
-    pub fn new() -> Self {
-        Default::default()
+impl Default for Ephemeral {
+    fn default() -> Self {
+        Ephemeral {
+            key_store: InMemEphemeralKeys::new(&mut OsRng),
+            reply_store: Default::default(),
+            credential_store: Default::default(),
+            gateway_details_store: Default::default(),
+        }
     }
 }
 
diff --git a/common/client-core/src/client/key_manager/persistence.rs b/common/client-core/src/client/key_manager/persistence.rs
index f781cd3da6..42a0b6a19a 100644
--- a/common/client-core/src/client/key_manager/persistence.rs
+++ b/common/client-core/src/client/key_manager/persistence.rs
@@ -3,6 +3,7 @@
 
 use crate::client::key_manager::ClientKeys;
 use async_trait::async_trait;
+use rand::{CryptoRng, RngCore};
 use std::error::Error;
 use std::sync::Arc;
 use tokio::sync::Mutex;
@@ -195,9 +196,20 @@ impl KeyStore for OnDiskKeys {
     }
 }
 
-#[derive(Clone, Default)]
+#[derive(Clone)]
 pub struct InMemEphemeralKeys {
-    keys: Arc<Mutex<Option<ClientKeys>>>,
+    keys: Arc<Mutex<ClientKeys>>,
+}
+
+impl InMemEphemeralKeys {
+    pub fn new<R>(rng: &mut R) -> Self
+    where
+        R: RngCore + CryptoRng,
+    {
+        InMemEphemeralKeys {
+            keys: Arc::new(Mutex::new(ClientKeys::generate_new(rng))),
+        }
+    }
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -210,11 +222,11 @@ impl KeyStore for InMemEphemeralKeys {
     type StorageError = EphemeralKeysError;
 
     async fn load_keys(&self) -> Result<ClientKeys, Self::StorageError> {
-        self.keys.lock().await.clone().ok_or(EphemeralKeysError)
+        Ok(self.keys.lock().await.clone())
     }
 
     async fn store_keys(&self, keys: &ClientKeys) -> Result<(), Self::StorageError> {
-        *self.keys.lock().await = Some(keys.clone());
+        *self.keys.lock().await = keys.clone();
         Ok(())
     }
 }