Fix error in persistent state reading (#1021)

This PR fixes issues with regards to persistent state reading when none
has been written. I.e. it adds additional checks to make sure that
the code does not attempt to read memory that it should not.

Author: Frederik Rothenberger

src/internet_identity/src/storage.rs

@@ -366,29 +366,59 @@ impl<T: candid::CandidType + serde::de::DeserializeOwned, M: Memory> Storage<T,
     /// Reads the persistent state from stable memory just outside of the space allocated to the highest user number.
     /// This is only used to restore state in `post_upgrade`.
     pub fn read_persistent_state(&self) -> Result<PersistentState, PersistentStateError> {
+        const WASM_PAGE_SIZE: u64 = 65536;
         let address = self.unused_memory_start();
-        let mut reader = Reader::new(&self.memory, address);
 
+        if address > self.memory.size() * WASM_PAGE_SIZE {
+            // the address where the persistent state would be is not allocated yet
+            return Err(PersistentStateError::NotFound);
+        }
+
+        let mut reader = Reader::new(&self.memory, address);
         let mut magic_buf: [u8; 4] = [0; 4];
-        reader
+        let bytes_read = reader
             .read(&mut magic_buf)
-            .map_err(|err| PersistentStateError::ReadError(err))?;
+            // if we hit out of bounds here, this means that the persistent state has not been
+            // written at the expected location and thus cannot be found
+            .map_err(|_| PersistentStateError::NotFound)?;
 
-        if magic_buf != PERSISTENT_STATE_MAGIC {
+        if bytes_read != 4 || magic_buf != PERSISTENT_STATE_MAGIC {
+            // less than the expected number of bytes were read or the magic does not match
+            // --> this is not the persistent state
             return Err(PersistentStateError::NotFound);
         }
 
         let mut size_buf: [u8; 8] = [0; 8];
-        reader
+        let bytes_read = reader
             .read(&mut size_buf)
-            .map_err(|err| PersistentStateError::ReadError(err))?;
+            .map_err(|err| PersistentStateError::ReadError(err))? as u64;
+
+        // check if we actually read the required amount of data
+        // note: this will only happen if we hit the memory bounds during read
+        if bytes_read != 8 {
+            let max_address = address + 4 + bytes_read;
+            return Err(PersistentStateError::ReadError(OutOfBounds {
+                max_address,
+                attempted_read_address: max_address + 1,
+            }));
+        }
 
         let size = u64::from_le_bytes(size_buf);
         let mut data_buf = Vec::new();
         data_buf.resize(size as usize, 0);
-        reader
+        let bytes_read = reader
             .read(data_buf.as_mut_slice())
-            .map_err(|err| PersistentStateError::ReadError(err))?;
+            .map_err(|err| PersistentStateError::ReadError(err))? as u64;
+
+        // check if we actually read the required amount of data
+        // note: this will only happen if we hit the memory bounds during read
+        if bytes_read != size {
+            let max_address = address + 4 + 8 + bytes_read;
+            return Err(PersistentStateError::ReadError(OutOfBounds {
+                max_address,
+                attempted_read_address: max_address + 1,
+            }));
+        }
 
         candid::decode_one(&data_buf).map_err(|err| PersistentStateError::CandidError(err))
     }

src/internet_identity/src/storage/tests.rs

@@ -195,6 +195,31 @@ fn should_save_persistent_state_at_expected_memory_address() {
     assert_eq!(buf, PERSISTENT_STATE_MAGIC);
 }
 
+#[test]
+fn should_not_find_persistent_state() {
+    let memory = VectorMemory::default();
+    let mut storage =
+        Storage::<Vec<DeviceDataInternal>, VectorMemory>::new((10_000, 3_784_873), memory.clone());
+    storage.flush();
+
+    let result = storage.read_persistent_state();
+    assert!(matches!(result, Err(PersistentStateError::NotFound)))
+}
+
+#[test]
+fn should_not_find_persistent_state_on_magic_bytes_mismatch() {
+    let memory = VectorMemory::default();
+
+    let mut storage =
+        Storage::<Vec<DeviceDataInternal>, VectorMemory>::new((10_000, 3_784_873), memory.clone());
+    storage.flush();
+
+    memory.write(RESERVED_HEADER_BYTES, b"IIPX"); // correct magic bytes are IIPS
+
+    let result = storage.read_persistent_state();
+    assert!(matches!(result, Err(PersistentStateError::NotFound)))
+}
+
 #[test]
 fn should_save_persistent_state_at_expected_memory_address_with_anchors() {
     const EXPECTED_ADDRESS: u64 = RESERVED_HEADER_BYTES + 100 * 2048; // number of anchors is 100
@@ -233,6 +258,25 @@ fn should_save_persistent_state_at_expected_memory_address_with_many_anchors() {
     assert_eq!(buf, PERSISTENT_STATE_MAGIC);
 }
 
+/// This test verifies that storage correctly reports `NotFound` if the persistent state address
+/// lies outside of the allocated stable memory range. This can happen on upgrade from a version
+/// that did not serialize a persistent state into stable memory.
+#[test]
+fn should_not_panic_on_unallocated_persistent_state_mem_address() {
+    let memory = VectorMemory::default();
+    let mut storage =
+        Storage::<Vec<DeviceDataInternal>, VectorMemory>::new((10_000, 3_784_873), memory.clone());
+    storage.flush();
+    for _ in 0..32 {
+        storage.allocate_user_number();
+    }
+
+    assert!(matches!(
+        storage.read_persistent_state(),
+        Err(PersistentStateError::NotFound)
+    ));
+}
+
 #[test]
 fn should_overwrite_persistent_state_with_next_anchor() {
     const EXPECTED_ADDRESS: u64 = RESERVED_HEADER_BYTES + 2048; // only one anchor exists