Remove steal_storage

examples/erase_storage.rs

@@ -17,7 +17,7 @@
 extern crate lang_items;
 
 use core::fmt::Write;
-use ctap2::env::tock::steal_storage;
+use ctap2::env::tock::take_storage;
 use libtock_drivers::console::Console;
 use libtock_drivers::led;
 use libtock_drivers::result::FlexUnwrap;
@@ -37,7 +37,7 @@ fn is_page_erased(storage: &dyn Storage, page: usize) -> bool {
 
 fn main() {
     led::get(1).flex_unwrap().on().flex_unwrap(); // red on dongle
-    let mut storage = unsafe { steal_storage() }.unwrap();
+    let mut storage = take_storage().unwrap();
     let num_pages = storage.num_pages();
     writeln!(Console::new(), "Erase {} pages of storage:", num_pages).unwrap();
     for page in 0..num_pages {

examples/store_latency.rs

@@ -21,8 +21,7 @@ use alloc::string::{String, ToString};
 use alloc::vec::Vec;
 use alloc::{format, vec};
 use core::fmt::Write;
-use ctap2::env::tock::{steal_storage, TockEnv};
-use ctap2::env::Env;
+use ctap2::env::tock::{take_storage, TockStorage};
 use libtock_drivers::console::Console;
 use libtock_drivers::timer::{self, Duration, Timer, Timestamp};
 use persistent_store::Store;
@@ -40,10 +39,8 @@ fn measure<T>(timer: &Timer, operation: impl FnOnce() -> T) -> (T, Duration<f64>
     (result, after - before)
 }
 
-// Only use one store at a time.
-unsafe fn boot_store(erase: bool) -> Store<<TockEnv as Env>::Storage> {
+fn boot_store(mut storage: TockStorage, erase: bool) -> Store<TockStorage> {
     use persistent_store::Storage;
-    let mut storage = steal_storage().unwrap();
     let num_pages = storage.num_pages();
     if erase {
         for page in 0..num_pages {
@@ -58,9 +55,8 @@ struct StorageConfig {
     num_pages: usize,
 }
 
-fn storage_config() -> StorageConfig {
+fn storage_config(storage: &TockStorage) -> StorageConfig {
     use persistent_store::Storage;
-    let storage = unsafe { steal_storage() }.unwrap();
     StorageConfig {
         num_pages: storage.num_pages(),
     }
@@ -77,11 +73,12 @@ struct Stat {
 }
 
 fn compute_latency(
+    storage: TockStorage,
     timer: &Timer,
     num_pages: usize,
     key_increment: usize,
     word_length: usize,
-) -> Stat {
+) -> (TockStorage, Stat) {
     let mut stat = Stat {
         key_increment,
         entry_length: word_length,
@@ -96,7 +93,7 @@ fn compute_latency(
     )
     .unwrap();
 
-    let mut store = unsafe { boot_store(true) };
+    let mut store = boot_store(storage, true);
     let total_capacity = store.capacity().unwrap().total();
     assert_eq!(store.capacity().unwrap().used(), 0);
     assert_eq!(store.lifetime().unwrap().used(), 0);
@@ -130,7 +127,8 @@ fn compute_latency(
     );
 
     // Measure latency of boot.
-    let (mut store, time) = measure(timer, || unsafe { boot_store(false) });
+    let storage = store.extract_storage();
+    let (mut store, time) = measure(timer, || boot_store(storage, false));
     writeln!(console, "Boot: {:.1}ms.", time.ms()).unwrap();
     stat.boot_ms = time.ms();
 
@@ -153,20 +151,23 @@ fn compute_latency(
     stat.compaction_ms = time.ms();
     assert!(store.lifetime().unwrap().used() > total_capacity + num_pages);
 
-    stat
+    (store.extract_storage(), stat)
 }
 
 fn main() {
     let mut with_callback = timer::with_callback(|_, _| {});
     let timer = with_callback.init().ok().unwrap();
-    let config = storage_config();
+    let storage = take_storage().unwrap();
+    let config = storage_config(&storage);
     let mut stats = Vec::new();
 
     writeln!(Console::new(), "\nRunning 2 tests...").unwrap();
     // Simulate a store full of credentials (of 50 words).
-    stats.push(compute_latency(&timer, config.num_pages, 1, 50));
+    let (storage, stat) = compute_latency(storage, &timer, config.num_pages, 1, 50);
+    stats.push(stat);
     // Simulate a store full of increments of a single counter.
-    stats.push(compute_latency(&timer, config.num_pages, 0, 1));
+    let (_storage, stat) = compute_latency(storage, &timer, config.num_pages, 0, 1);
+    stats.push(stat);
     writeln!(Console::new(), "\nDone.\n").unwrap();
 
     const HEADERS: &[&str] = &[

libraries/persistent_store/src/store.rs

@@ -238,6 +238,11 @@ impl<S: Storage> Store<S> {
         Ok(store)
     }
 
+    /// Extracts the storage.
+    pub fn extract_storage(self) -> S {
+        self.storage
+    }
+
     /// Iterates over the entries.
     pub fn iter<'a>(&'a self) -> StoreResult<StoreIter<'a>> {
         let head = or_invalid(self.head)?;
@@ -1162,11 +1167,6 @@ impl Store<BufferStorage> {
         &mut self.storage
     }
 
-    /// Extracts the storage.
-    pub fn extract_storage(self) -> BufferStorage {
-        self.storage
-    }
-
     /// Returns the value of a possibly deleted entry.
     ///
     /// If the value has been partially compacted, only return the non-compacted part. Returns an

src/env/tock/mod.rs

@@ -1,4 +1,4 @@
-use self::storage::{SyscallStorage, SyscallUpgradeStorage};
+pub use self::storage::{TockStorage, TockUpgradeStorage};
 use crate::ctap::hid::{ChannelID, CtapHid, CtapHidCommand, KeepaliveStatus, ProcessedPacket};
 use crate::ctap::status_code::Ctap2StatusCode;
 use crate::env::{Env, UserPresence};
@@ -20,8 +20,8 @@ mod storage;
 
 pub struct TockEnv {
     rng: TockRng256,
-    store: Store<SyscallStorage>,
-    upgrade_storage: Option<SyscallUpgradeStorage>,
+    store: Store<TockStorage>,
+    upgrade_storage: Option<TockUpgradeStorage>,
 }
 
 impl TockEnv {
@@ -31,12 +31,10 @@ impl TockEnv {
     ///
     /// - If called a second time.
     pub fn new() -> Self {
-        // 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();
+        // We rely on `take_storage` to ensure that this function is called only once.
+        let storage = take_storage().unwrap();
         let store = Store::new(storage).ok().unwrap();
-        let upgrade_storage = SyscallUpgradeStorage::new().ok();
+        let upgrade_storage = TockUpgradeStorage::new().ok();
         TockEnv {
             rng: TockRng256 {},
             store,
@@ -45,17 +43,16 @@ impl TockEnv {
     }
 }
 
-/// Creates a new storage instance.
+/// Returns the unique storage instance.
 ///
-/// # Safety
+/// # Panics
 ///
-/// It is probably technically memory-safe to have multiple storage instances at the same time, but
-/// for extra precaution we mark the function as unsafe. To ensure correct usage, this function
-/// should only be called if the previous storage instance was dropped.
-// This function is exposed to example binaries testing the hardware. This could probably be cleaned
-// up by having the persistent store return its storage.
-pub unsafe fn steal_storage() -> StorageResult<SyscallStorage> {
-    SyscallStorage::new()
+/// - If called a second time.
+pub fn take_storage() -> StorageResult<TockStorage> {
+    // Make sure the storage was not already taken.
+    static TAKEN: AtomicBool = AtomicBool::new(false);
+    assert!(!TAKEN.fetch_or(true, Ordering::SeqCst));
+    TockStorage::new()
 }
 
 impl UserPresence for TockEnv {
@@ -67,8 +64,8 @@ impl UserPresence for TockEnv {
 impl Env for TockEnv {
     type Rng = TockRng256;
     type UserPresence = Self;
-    type Storage = SyscallStorage;
-    type UpgradeStorage = SyscallUpgradeStorage;
+    type Storage = TockStorage;
+    type UpgradeStorage = TockUpgradeStorage;
 
     fn rng(&mut self) -> &mut Self::Rng {
         &mut self.rng

src/env/tock/storage.rs

@@ -115,7 +115,7 @@ fn erase_page(ptr: usize, page_length: usize) -> StorageResult<()> {
     block_command(DRIVER_NUMBER, command_nr::ERASE_PAGE, ptr, page_length)
 }
 
-pub struct SyscallStorage {
+pub struct TockStorage {
     word_size: usize,
     page_size: usize,
     num_pages: usize,
@@ -124,7 +124,7 @@ pub struct SyscallStorage {
     storage_locations: Vec<&'static [u8]>,
 }
 
-impl SyscallStorage {
+impl TockStorage {
     /// Provides access to the embedded flash if available.
     ///
     /// # Errors
@@ -134,8 +134,8 @@ impl SyscallStorage {
     /// - The page size is a power of two.
     /// - The page size is a multiple of the word size.
     /// - The storage is page-aligned.
-    pub fn new() -> StorageResult<SyscallStorage> {
-        let mut syscall = SyscallStorage {
+    pub fn new() -> StorageResult<TockStorage> {
+        let mut syscall = TockStorage {
             word_size: get_info(command_nr::get_info_nr::WORD_SIZE, 0)?,
             page_size: get_info(command_nr::get_info_nr::PAGE_SIZE, 0)?,
             num_pages: 0,
@@ -175,7 +175,7 @@ impl SyscallStorage {
     }
 }
 
-impl Storage for SyscallStorage {
+impl Storage for TockStorage {
     fn word_size(&self) -> usize {
         self.word_size
     }
@@ -217,13 +217,13 @@ impl Storage for SyscallStorage {
     }
 }
 
-pub struct SyscallUpgradeStorage {
+pub struct TockUpgradeStorage {
     page_size: usize,
     partition: ModRange,
     metadata: ModRange,
 }
 
-impl SyscallUpgradeStorage {
+impl TockUpgradeStorage {
     /// Provides access to the other upgrade partition and metadata if available.
     ///
     /// The implementation assumes that storage locations returned by the kernel through
@@ -238,8 +238,8 @@ impl SyscallUpgradeStorage {
     /// Returns a `NotAligned` error if partitions or metadata ranges are
     /// - not exclusive or,
     /// - not consecutive.
-    pub fn new() -> StorageResult<SyscallUpgradeStorage> {
-        let mut locations = SyscallUpgradeStorage {
+    pub fn new() -> StorageResult<TockUpgradeStorage> {
+        let mut locations = TockUpgradeStorage {
             page_size: get_info(command_nr::get_info_nr::PAGE_SIZE, 0)?,
             partition: ModRange::new_empty(),
             metadata: ModRange::new_empty(),
@@ -287,7 +287,7 @@ impl SyscallUpgradeStorage {
     }
 }
 
-impl UpgradeStorage for SyscallUpgradeStorage {
+impl UpgradeStorage for TockUpgradeStorage {
     fn read_partition(&self, offset: usize, length: usize) -> StorageResult<&[u8]> {
         if length == 0 {
             return Err(StorageError::OutOfBounds);