Fix seeking / reading / writing past the end (#17)

crates/icfs/internal.rs

@@ -0,0 +1,59 @@
+// Copied from: https://github.com/rust-lang/rust/blob/a2af9cf1cf6ccb195eae40cdd793939bc77e7e73/library/std/src/io/mod.rs#L356
+// This may be avoidable in future: https://github.com/rust-lang/rfcs/pull/1210
+pub (crate) fn default_read_to_end<R: std::io::Read + ?Sized>(r: &mut R, buf: &mut Vec<u8>) -> std::io::Result<usize> {
+    let start_len = buf.len();
+    let start_cap = buf.capacity();
+
+    let mut initialized = 0; // Extra initialized bytes from previous loop iteration
+    loop {
+        if buf.len() == buf.capacity() {
+            buf.reserve(32); // buf is full, need more space
+        }
+
+        let mut read_buf = std::io::ReadBuf::uninit(buf.spare_capacity_mut());
+
+        // SAFETY: These bytes were initialized but not filled in the previous loop
+        unsafe {
+            read_buf.assume_init(initialized);
+        }
+
+        match r.read_buf(&mut read_buf) {
+            Ok(()) => {}
+            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
+            Err(e) => return Err(e),
+        }
+
+        if read_buf.filled_len() == 0 {
+            return Ok(buf.len() - start_len);
+        }
+
+        // store how much was initialized but not filled
+        initialized = read_buf.initialized_len() - read_buf.filled_len();
+        let new_len = read_buf.filled_len() + buf.len();
+
+        // SAFETY: ReadBuf's invariants mean this much memory is init
+        unsafe {
+            buf.set_len(new_len);
+        }
+
+        if buf.len() == buf.capacity() && buf.capacity() == start_cap {
+            // The buffer might be an exact fit. Let's read into a probe buffer
+            // and see if it returns `Ok(0)`. If so, we've avoided an
+            // unnecessary doubling of the capacity. But if not, append the
+            // probe buffer to the primary buffer and let its capacity grow.
+            let mut probe = [0u8; 32];
+
+            loop {
+                match r.read(&mut probe) {
+                    Ok(0) => return Ok(buf.len() - start_len),
+                    Ok(n) => {
+                        buf.extend_from_slice(&probe[..n]);
+                        break;
+                    }
+                    Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
+                    Err(e) => return Err(e),
+                }
+            }
+        }
+    }
+}

crates/icfs/lib.rs

@@ -1,2 +1,6 @@
+#![feature(read_buf)]
+#![feature(result_flattening)]
+
+mod internal;
 mod stable_memory;
 pub use stable_memory::StableMemory;

crates/icfs/stable_memory.rs

@@ -1,7 +1,7 @@
 // Based on https://github.com/dfinity/cdk-rs/blob/a253119adb08929b6304d007ee0a6a37960656ed/src/ic-cdk/src/api/stable.rs
 // * Supports 64-bit addressed memory
 use ic_cdk::api::stable::{
-    stable64_grow, stable64_read, stable64_size, stable64_write, StableMemoryError,
+    stable64_read, stable64_write, StableMemoryError,
 };
 use std::io;
 
@@ -40,20 +40,20 @@ pub fn capacity() -> usize {
 
 /// Attempts to grow the memory by adding new pages.
 pub fn grow(added_pages: u64) -> Result<u64, StableMemoryError> {
-    stable64_grow(added_pages)
+    ic_cdk::api::stable::stable64_grow(added_pages)
 }
 
 /// Gets current size of the stable memory in WebAssembly pages.
 pub fn size() -> u64 {
-    stable64_size()
+    ic_cdk::api::stable::stable64_size()
 }
 
 /// Reads data from the stable memory location specified by an offset.
 pub fn read(stable_memory: &mut StableMemory, buf: &mut [u8]) -> Result<usize, StableMemoryError> {
     let offset = get_offset(stable_memory);
     let capacity = capacity();
     let read_buf = if buf.len() + offset > capacity {
-        if offset < capacity {
+        if offset <= capacity {
             &mut buf[..capacity - offset]
         } else {
             return Err(StableMemoryError::OutOfBounds);
@@ -99,18 +99,27 @@ fn seek(stable_memory: &mut StableMemory, pos: io::SeekFrom) -> Result<u64, Stab
 /// error out is if it cannot grow the memory.
 pub fn write(stable_memory: &mut StableMemory, buf: &[u8]) -> Result<usize, StableMemoryError> {
     let offset = get_offset(stable_memory);
-    let new_offset = offset + buf.len();
-    let memory_end_bytes = (offset + buf.len()) as u64;
+    let memory_end_bytes = offset + buf.len();
     let memory_end_pages =
-        (memory_end_bytes + WASM_PAGE_SIZE_IN_BYTES - 1) / WASM_PAGE_SIZE_IN_BYTES;
-    let current_pages = capacity() as u64;
-    let additional_pages_required = memory_end_pages.saturating_sub(current_pages);
+        (memory_end_bytes as u64 + WASM_PAGE_SIZE_IN_BYTES - 1) / WASM_PAGE_SIZE_IN_BYTES;
+    let additional_pages_required = memory_end_pages.saturating_sub(capacity() as u64);
     if additional_pages_required > 0 {
-        stable64_grow(additional_pages_required)?;
+        grow(additional_pages_required)?;
     }
-    stable64_write(offset as u64, buf);
+    let capacity = capacity();
+    let write_buf = if memory_end_bytes > capacity {
+        if offset <= capacity {
+            &buf[..capacity - offset]
+        } else {
+            return Err(StableMemoryError::OutOfBounds);
+        }
+    } else {
+        buf
+    };
+    stable64_write(offset as u64, write_buf);
+    let new_offset = offset + write_buf.len();
     set_offset(stable_memory, new_offset);
-    Ok(buf.len())
+    Ok(write_buf.len())
 }
 
 impl StableMemory {
@@ -147,8 +156,11 @@ impl Default for StableMemory {
 
 impl std::io::Read for StableMemory {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
-        // read(self, buf).map_err(|e| io::Error::new(io::ErrorKind::Other, e))
-        read(self, buf).or(Ok(0)) // Read defines EOF to be success
+        read(self, buf).map_err(|e| io::Error::new(io::ErrorKind::Other, e))
+    }
+
+    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
+        crate::internal::default_read_to_end(self, buf).or(Ok(0)) // Read defines EOF to be success
     }
 }
 

examples/icfs/icfs.did

@@ -2,10 +2,12 @@ service : {
   test_writer : () -> ();
   test_writer_vectored : () -> ();
   test_writer_seek : () -> ();
+  test_writer_error : () -> ();
   test_reader : () -> ();
   test_reader_vectored : () -> ();
   test_read_to_end : () -> ();
   test_read_exact : () -> ();
+  test_reader_error : () -> ();
   test_seek_past_end : () -> ();
   test_seek_before_0 : () -> ();
 }

examples/icfs/lib.rs

@@ -106,6 +106,20 @@ fn test_writer_seek() {
     })
 }
 
+#[update]
+fn test_writer_error() {
+    setup();
+    STABLE_MEMORY.with(|stable_memory| {
+        let mut stable_memory = *stable_memory.borrow();
+        let capacity = icfs::StableMemory::capacity();
+        let offset = capacity as u64 - 2;
+        assert_eq!(stable_memory.seek(SeekFrom::End(-2)).unwrap(), offset);
+        assert_eq!(stable_memory.write(&[0]).unwrap(), 1);
+        assert_eq!(stable_memory.write(&[0, 0]).unwrap(), 1);
+        assert_eq!(stable_memory.write(&[0, 0]).unwrap(), 0);
+    })
+}
+
 #[update]
 fn test_reader() {
     setup();
@@ -234,15 +248,39 @@ fn test_read_exact() {
     })
 }
 
+#[update]
+fn test_reader_error() {
+    setup();
+    STABLE_MEMORY.with(|stable_memory| {
+        let mut stable_memory = *stable_memory.borrow();
+        let capacity = icfs::StableMemory::capacity();
+        let offset = capacity as u64 - 2;
+        assert_eq!(stable_memory.seek(SeekFrom::End(-2)).unwrap(), offset);
+
+        let mut buf = [0];
+        assert_eq!(stable_memory.read(&mut buf).unwrap(), 1);
+
+        let mut buf = [0, 0];
+        assert_eq!(stable_memory.read(&mut buf).unwrap(), 1);
+
+        let mut buf = [0, 0];
+        assert_eq!(stable_memory.read(&mut buf).unwrap(), 0);
+    })
+}
+
 #[update]
 fn test_seek_past_end() {
     setup();
     STABLE_MEMORY.with(|stable_memory| {
         let mut stable_memory = *stable_memory.borrow();
         let capacity = icfs::StableMemory::capacity();
         let offset = capacity as u64 + 1;
+
+        assert_eq!(stable_memory.seek(SeekFrom::Start(offset)).unwrap(), offset);
+        assert!(stable_memory.read(&mut [0]).is_err());
+
         assert_eq!(stable_memory.seek(SeekFrom::Start(offset)).unwrap(), offset);
-        assert_eq!(stable_memory.read(&mut [0]).unwrap(), 0);
+        assert!(stable_memory.write(&[3]).is_err());
     })
 }
 

examples/icfs/test.ic-repl

@@ -13,6 +13,9 @@ assert result == null;
 let result = call icfs.test_writer_seek();
 assert result == null;
 
+let result = call icfs.test_writer_error();
+assert result == null;
+
 let result = call icfs.test_reader();
 assert result == null;
 
@@ -25,6 +28,9 @@ assert result == null;
 let result = call icfs.test_read_exact();
 assert result == null;
 
+let result = call icfs.test_reader_error();
+assert result == null;
+
 let result = call icfs.test_seek_past_end();
 assert result == null;
 

flake.nix

@@ -41,9 +41,9 @@
         rust = (mozilla.rustChannelOf {
           # channel = "1.56.0";
           # sha256 = "L1e0o7azRjOHd0zBa+xkFnxdFulPofTedSTEYZSjj2s=";
-          channel = "nightly"; # for write_all_vectored
-          date = "2021-10-21"; # date of 1.56.0 release
-          sha256 = "KNIEZgJQBD31InRkeRX9SGh49qOkiwS5s+RXl28uA48=";
+          channel = "nightly";
+          date = "2022-04-07"; # day of 1.60.0 release
+          sha256 = "z+elrzVPDgtdqSMg8NTSGqkmfsK6vOn9XUFXcsSXhXo=";
           # sha256 = pkgs.lib.fakeSha256;
         }).rust.override {
           extensions = [