From f47e1e2a8663e60a3c9b61db246cebfeb193d8f5 Mon Sep 17 00:00:00 2001
From: Jean-Michel Picod <jmichel@google.com>
Date: Wed, 25 Nov 2020 17:44:19 +0100
Subject: [PATCH] Ensure store behaves as expected in prod

---
 src/ctap/storage.rs | 33 +++++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/ctap/storage.rs b/src/ctap/storage.rs
index 5793e6c..da8828c 100644
--- a/src/ctap/storage.rs
+++ b/src/ctap/storage.rs
@@ -229,6 +229,18 @@ impl PersistentStore {
                 })
                 .unwrap();
         }
+        // TODO(jmichel): remove this when vendor command is in place
+        #[cfg(not(any(test, feature = "ram_storage")))]
+        self.load_attestation_from_firmware();
+
+        if self.store.find_one(&Key::Aaguid).is_none() {
+            self.set_aaguid(key_material::AAGUID).unwrap();
+        }
+    }
+
+    // TODO(jmichel): remove this function when vendor command is in place.
+    #[cfg(not(any(test, feature = "ram_storage")))]
+    fn load_attestation_from_firmware(&mut self) {
         // The following 2 entries are meant to be written by vendor-specific commands.
         if self.store.find_one(&Key::AttestationPrivateKey).is_none() {
             self.set_attestation_private_key(key_material::ATTESTATION_PRIVATE_KEY)
@@ -238,10 +250,6 @@ impl PersistentStore {
             self.set_attestation_certificate(key_material::ATTESTATION_CERTIFICATE)
                 .unwrap();
         }
-
-        if self.store.find_one(&Key::Aaguid).is_none() {
-            self.set_aaguid(key_material::AAGUID).unwrap();
-        }
     }
 
     pub fn find_credential(
@@ -1000,6 +1008,23 @@ mod test {
         let mut rng = ThreadRng256 {};
         let mut persistent_store = PersistentStore::new(&mut rng);
 
+        // Make sure the attestation are absent. There is no batch attestation in tests.
+        assert!(persistent_store
+            .attestation_private_key()
+            .unwrap()
+            .is_none());
+        assert!(persistent_store
+            .attestation_certificate()
+            .unwrap()
+            .is_none());
+
+        // Make sure the persistent keys are initialized.
+        persistent_store
+            .set_attestation_private_key(key_material::ATTESTATION_PRIVATE_KEY)
+            .unwrap();
+        persistent_store
+            .set_attestation_certificate(key_material::ATTESTATION_CERTIFICATE)
+            .unwrap();
         assert_eq!(&persistent_store.aaguid().unwrap(), key_material::AAGUID);
 
         // The persistent keys stay initialized and preserve their value after a reset.