mount_or_else: ensure that allocation has not been overwritten in the else clause

Author: sosthene-nitrokey

src/fs.rs

@@ -1104,6 +1104,13 @@ impl<'a, Storage: driver::Storage> Filesystem<'a, Storage> {
         Ok(fs)
     }
 
+    fn set_alloc_config(alloc: &mut Allocation<Storage>, storage: &mut Storage) {
+        alloc.config.context = storage as *mut _ as *mut c_void;
+        alloc.config.read_buffer = alloc.cache.read.get() as *mut c_void;
+        alloc.config.prog_buffer = alloc.cache.write.get() as *mut c_void;
+        alloc.config.lookahead_buffer = alloc.cache.lookahead.get() as *mut c_void;
+    }
+
     /// Mount the filesystem or, if that fails, call `f` with the mount error and the storage and then try again.
     pub fn mount_or_else<F>(
         alloc: &'a mut Allocation<Storage>,
@@ -1113,9 +1120,11 @@ impl<'a, Storage: driver::Storage> Filesystem<'a, Storage> {
     where
         F: FnOnce(Error, &mut Storage, &mut Allocation<Storage>) -> Result<()>,
     {
-        let fs = Self::new(alloc, storage);
+        let mut fs = Self::new(alloc, storage);
         if let Err(err) = fs.raw_mount() {
-            f(err, fs.storage, &mut fs.alloc.borrow_mut())?;
+            let alloc = fs.alloc.get_mut();
+            f(err, fs.storage, alloc)?;
+            Self::set_alloc_config(alloc, fs.storage);
             fs.raw_mount()?;
         }
         Ok(fs)
@@ -1130,12 +1139,7 @@ impl<'a, Storage: driver::Storage> Filesystem<'a, Storage> {
 
     // Not public, user should use `mount`, possibly after `format`
     fn new(alloc: &'a mut Allocation<Storage>, storage: &'a mut Storage) -> Self {
-        alloc.config.context = storage as *mut _ as *mut c_void;
-
-        alloc.config.read_buffer = alloc.cache.read.get() as *mut c_void;
-        alloc.config.prog_buffer = alloc.cache.write.get() as *mut c_void;
-        alloc.config.lookahead_buffer = alloc.cache.lookahead.get() as *mut c_void;
-
+        Self::set_alloc_config(alloc, storage);
         Filesystem {
             alloc: RefCell::new(alloc),
             storage,

src/tests.rs

@@ -2,7 +2,7 @@ use core::convert::TryInto;
 use generic_array::typenum::consts;
 
 use crate::{
-    fs::{Attribute, File, Filesystem},
+    fs::{Allocation, Attribute, File, Filesystem},
     io::{Error, OpenSeekFrom, Read, Result, SeekFrom},
     path, BACKEND_VERSION, DISK_VERSION,
 };
@@ -521,6 +521,19 @@ fn test_iter_dirs() {
     .unwrap();
 }
 
+#[test]
+fn test_mount_or_else_clobber_alloc() {
+    let mut backend = Ram::default();
+    let mut storage = RamStorage::new(&mut backend);
+    let alloc = &mut Allocation::new();
+    Filesystem::mount_or_else(alloc, &mut storage, |_, storage, alloc| {
+        *alloc = Allocation::new();
+        Filesystem::format(storage).unwrap();
+        Ok(())
+    })
+    .unwrap();
+}
+
 // // These are some tests that ensure our type constructions
 // // actually do what we intend them to do.
 // // Since dev-features cannot be optional, trybuild is not `no_std`,