Use Store instead of Storage in Env

src/env/mod.rs

@@ -2,7 +2,7 @@ use crate::api::upgrade_storage::UpgradeStorage;
 use crate::ctap::hid::ChannelID;
 use crate::ctap::status_code::Ctap2StatusCode;
 use crypto::rng256::Rng256;
-use persistent_store::{Storage, StorageResult};
+use persistent_store::{Storage, Store};
 
 #[cfg(feature = "std")]
 pub mod test;
@@ -24,14 +24,11 @@ pub trait Env {
 
     fn rng(&mut self) -> &mut Self::Rng;
     fn user_presence(&mut self) -> &mut Self::UserPresence;
+    fn store(&mut self) -> &mut Store<Self::Storage>;
 
-    /// Returns the unique storage instance.
+    /// Returns the upgrade storage instance.
     ///
-    /// This function is called at most once. Implementation may panic if called more than once.
-    fn storage(&mut self) -> StorageResult<Self::Storage>;
-
-    /// Returns the unique upgrade storage instance.
-    ///
-    /// This function is called at most once. Implementation may panic if called more than once.
-    fn upgrade_storage(&mut self) -> StorageResult<Self::UpgradeStorage>;
+    /// Upgrade storage is optional, so implementations may return `None`. However, implementations
+    /// should either always return `None` or always return `Some`.
+    fn upgrade_storage(&mut self) -> Option<&mut Self::UpgradeStorage>;
 }

src/env/test/mod.rs

@@ -3,26 +3,55 @@ use crate::ctap::hid::ChannelID;
 use crate::ctap::status_code::Ctap2StatusCode;
 use crate::env::{Env, UserPresence};
 use crypto::rng256::ThreadRng256;
-use persistent_store::{BufferOptions, BufferStorage, StorageResult};
+use persistent_store::{BufferOptions, BufferStorage, Store};
 
 mod upgrade_storage;
 
 pub struct TestEnv {
     rng: ThreadRng256,
     user_presence: TestUserPresence,
+    store: Store<BufferStorage>,
+    upgrade_storage: Option<BufferUpgradeStorage>,
 }
 
 pub struct TestUserPresence {
     check: Box<dyn Fn(ChannelID) -> Result<(), Ctap2StatusCode>>,
 }
 
+fn new_storage() -> BufferStorage {
+    // Use the Nordic configuration.
+    const PAGE_SIZE: usize = 0x1000;
+    const NUM_PAGES: usize = 20;
+    let store = vec![0xff; NUM_PAGES * PAGE_SIZE].into_boxed_slice();
+    let options = BufferOptions {
+        word_size: 4,
+        page_size: PAGE_SIZE,
+        max_word_writes: 2,
+        max_page_erases: 10000,
+        strict_mode: true,
+    };
+    BufferStorage::new(store, options)
+}
+
 impl TestEnv {
     pub fn new() -> Self {
         let rng = ThreadRng256 {};
         let user_presence = TestUserPresence {
             check: Box::new(|_| Ok(())),
         };
-        TestEnv { rng, user_presence }
+        let storage = new_storage();
+        let store = Store::new(storage).ok().unwrap();
+        let upgrade_storage = Some(BufferUpgradeStorage::new().unwrap());
+        TestEnv {
+            rng,
+            user_presence,
+            store,
+            upgrade_storage,
+        }
+    }
+
+    pub fn disable_upgrade_storage(&mut self) {
+        self.upgrade_storage = None;
     }
 }
 
@@ -52,22 +81,11 @@ impl Env for TestEnv {
         &mut self.user_presence
     }
 
-    fn storage(&mut self) -> StorageResult<Self::Storage> {
-        // Use the Nordic configuration.
-        const PAGE_SIZE: usize = 0x1000;
-        const NUM_PAGES: usize = 20;
-        let store = vec![0xff; NUM_PAGES * PAGE_SIZE].into_boxed_slice();
-        let options = BufferOptions {
-            word_size: 4,
-            page_size: PAGE_SIZE,
-            max_word_writes: 2,
-            max_page_erases: 10000,
-            strict_mode: true,
-        };
-        Ok(BufferStorage::new(store, options))
+    fn store(&mut self) -> &mut Store<Self::Storage> {
+        &mut self.store
     }
 
-    fn upgrade_storage(&mut self) -> StorageResult<Self::UpgradeStorage> {
-        BufferUpgradeStorage::new()
+    fn upgrade_storage(&mut self) -> Option<&mut Self::UpgradeStorage> {
+        self.upgrade_storage.as_mut()
     }
 }

src/env/tock/mod.rs

@@ -14,14 +14,14 @@ use libtock_drivers::console::Console;
 use libtock_drivers::result::{FlexUnwrap, TockError};
 use libtock_drivers::timer::Duration;
 use libtock_drivers::{led, timer, usb_ctap_hid};
-use persistent_store::StorageResult;
+use persistent_store::{StorageResult, Store};
 
 mod storage;
 
 pub struct TockEnv {
     rng: TockRng256,
-    storage: bool,
-    upgrade_storage: bool,
+    store: Store<SyscallStorage>,
+    upgrade_storage: SyscallUpgradeStorage,
 }
 
 impl TockEnv {
@@ -34,10 +34,13 @@ impl TockEnv {
         // Make sure the environment was not already taken.
         static TAKEN: AtomicBool = AtomicBool::new(false);
         assert!(!TAKEN.fetch_or(true, Ordering::SeqCst));
+        let storage = unsafe { steal_storage() }.unwrap();
+        let store = Store::new(storage).ok().unwrap();
+        let upgrade_storage = SyscallUpgradeStorage::new().unwrap();
         TockEnv {
             rng: TockRng256 {},
-            storage: false,
-            upgrade_storage: false,
+            store,
+            upgrade_storage,
         }
     }
 }
@@ -75,23 +78,15 @@ impl Env for TockEnv {
         self
     }
 
-    fn storage(&mut self) -> StorageResult<Self::Storage> {
-        assert_once(&mut self.storage);
-        unsafe { steal_storage() }
+    fn store(&mut self) -> &mut Store<Self::Storage> {
+        &mut self.store
     }
 
-    fn upgrade_storage(&mut self) -> StorageResult<Self::UpgradeStorage> {
-        assert_once(&mut self.upgrade_storage);
-        SyscallUpgradeStorage::new()
+    fn upgrade_storage(&mut self) -> Option<&mut Self::UpgradeStorage> {
+        Some(&mut self.upgrade_storage)
     }
 }
 
-/// Asserts a boolean is false and sets it to true.
-fn assert_once(b: &mut bool) {
-    assert!(!*b);
-    *b = true;
-}
-
 // Returns whether the keepalive was sent, or false if cancelled.
 fn send_keepalive_up_needed(
     cid: ChannelID,