From ff6c700cd9dad77ec5186005056c3419a40c5356 Mon Sep 17 00:00:00 2001
From: Julien Cretin <cretin@google.com>
Date: Thu, 30 Jun 2022 15:03:38 +0200
Subject: [PATCH] Use indirection to implement the default KeyStore

---
 src/api/key_store.rs | 52 +++++++++++++++++++++++++-------------------
 src/env/test/mod.rs  |  3 +++
 src/env/tock/mod.rs  |  3 +++
 3 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/src/api/key_store.rs b/src/api/key_store.rs
index 9723228..f3714a2 100644
--- a/src/api/key_store.rs
+++ b/src/api/key_store.rs
@@ -50,7 +50,10 @@ pub struct Error;
 /// Key of the environment store reserved for the key store.
 pub const STORE_KEY: usize = 2046;
 
-impl<T: Env> KeyStore for T {
+/// Implements a default key store using the environment rng and store.
+pub trait Helper: Env {}
+
+impl<T: Helper> KeyStore for T {
     fn key_handle_encryption(&mut self) -> Result<[u8; 32], Error> {
         Ok(get_master_keys(self)?.encryption)
     }
@@ -114,28 +117,33 @@ impl From<StoreError> for Error {
     }
 }
 
-#[test]
-fn test_key_store() {
-    let mut env = crate::env::test::TestEnv::new();
-    let key_store = env.key_store();
+#[cfg(test)]
+mod test {
+    use super::*;
 
-    // Master keys are well-defined and stable.
-    let encryption_key = key_store.key_handle_encryption().unwrap();
-    let authentication_key = key_store.key_handle_authentication().unwrap();
-    assert_eq!(key_store.key_handle_encryption(), Ok(encryption_key));
-    assert_eq!(
-        key_store.key_handle_authentication(),
-        Ok(authentication_key)
-    );
+    #[test]
+    fn test_key_store() {
+        let mut env = crate::env::test::TestEnv::new();
+        let key_store = env.key_store();
 
-    // ECDSA seeds are well-defined and stable.
-    let ecdsa_seed = key_store.generate_ecdsa_seed().unwrap();
-    let ecdsa_key = key_store.derive_ecdsa(&ecdsa_seed).unwrap();
-    assert_eq!(key_store.derive_ecdsa(&ecdsa_seed), Ok(ecdsa_key));
+        // Master keys are well-defined and stable.
+        let encryption_key = key_store.key_handle_encryption().unwrap();
+        let authentication_key = key_store.key_handle_authentication().unwrap();
+        assert_eq!(key_store.key_handle_encryption(), Ok(encryption_key));
+        assert_eq!(
+            key_store.key_handle_authentication(),
+            Ok(authentication_key)
+        );
 
-    // Master keys change after reset. We don't require this for ECDSA seeds because it's not the
-    // case, but it might be better.
-    key_store.reset().unwrap();
-    assert!(key_store.key_handle_encryption().unwrap() != encryption_key);
-    assert!(key_store.key_handle_authentication().unwrap() != authentication_key);
+        // ECDSA seeds are well-defined and stable.
+        let ecdsa_seed = key_store.generate_ecdsa_seed().unwrap();
+        let ecdsa_key = key_store.derive_ecdsa(&ecdsa_seed).unwrap();
+        assert_eq!(key_store.derive_ecdsa(&ecdsa_seed), Ok(ecdsa_key));
+
+        // Master keys change after reset. We don't require this for ECDSA seeds because it's not
+        // the case, but it might be better.
+        key_store.reset().unwrap();
+        assert!(key_store.key_handle_encryption().unwrap() != encryption_key);
+        assert!(key_store.key_handle_authentication().unwrap() != authentication_key);
+    }
 }
diff --git a/src/env/test/mod.rs b/src/env/test/mod.rs
index ea5f2b0..1cd6d24 100644
--- a/src/env/test/mod.rs
+++ b/src/env/test/mod.rs
@@ -16,6 +16,7 @@ use self::upgrade_storage::BufferUpgradeStorage;
 use crate::api::connection::{HidConnection, SendOrRecvResult, SendOrRecvStatus};
 use crate::api::customization::DEFAULT_CUSTOMIZATION;
 use crate::api::firmware_protection::FirmwareProtection;
+use crate::api::key_store;
 use crate::api::user_presence::{UserPresence, UserPresenceResult};
 use crate::clock::ClockInt;
 use crate::env::Env;
@@ -147,6 +148,8 @@ impl FirmwareProtection for TestEnv {
     }
 }
 
+impl key_store::Helper for TestEnv {}
+
 impl Env for TestEnv {
     type Rng = TestRng256;
     type UserPresence = TestUserPresence;
diff --git a/src/env/tock/mod.rs b/src/env/tock/mod.rs
index e76088f..8cb6274 100644
--- a/src/env/tock/mod.rs
+++ b/src/env/tock/mod.rs
@@ -16,6 +16,7 @@ pub use self::storage::{TockStorage, TockUpgradeStorage};
 use crate::api::connection::{HidConnection, SendOrRecvError, SendOrRecvResult, SendOrRecvStatus};
 use crate::api::customization::{CustomizationImpl, DEFAULT_CUSTOMIZATION};
 use crate::api::firmware_protection::FirmwareProtection;
+use crate::api::key_store;
 use crate::api::user_presence::{UserPresence, UserPresenceError, UserPresenceResult};
 use crate::clock::{ClockInt, KEEPALIVE_DELAY_MS};
 use crate::env::Env;
@@ -193,6 +194,8 @@ impl FirmwareProtection for TockEnv {
     }
 }
 
+impl key_store::Helper for TockEnv {}
+
 impl Env for TockEnv {
     type Rng = TockRng256;
     type UserPresence = Self;