From 652988f0dae81c8220c65f9e7559d0938e782fd5 Mon Sep 17 00:00:00 2001
From: Andrew Brown <andrew.brown@intel.com>
Date: Tue, 3 Jan 2023 10:32:53 -0800
Subject: [PATCH] review: handle underfill of temporary buffer

---
 crates/wasi-common/src/snapshots/preview_1.rs | 17 +++++++++++++----
 crates/wiggle/src/lib.rs                      |  5 +++++
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/crates/wasi-common/src/snapshots/preview_1.rs b/crates/wasi-common/src/snapshots/preview_1.rs
index 28051ccb10f1..c3cb0035277d 100644
--- a/crates/wasi-common/src/snapshots/preview_1.rs
+++ b/crates/wasi-common/src/snapshots/preview_1.rs
@@ -1111,15 +1111,24 @@ impl wasi_snapshot_preview1::WasiSnapshotPreview1 for WasiCtx {
         buf: &GuestPtr<'a, u8>,
         buf_len: types::Size,
     ) -> Result<(), Error> {
-        let buf = buf.as_array(buf_len).as_unsafe_slice_mut()?;
+        let buf = buf.as_array(buf_len);
         if buf.is_shared_memory() {
             // If the Wasm memory is shared, copy to an intermediate buffer to
             // avoid Rust unsafety (i.e., the called function could rely on
             // `&mut [u8]`'s exclusive ownership which is not guaranteed due to
             // potential access from other threads).
-            let mut tmp = vec![0; buf.len().min(MAX_SHARED_BUFFER_SIZE)];
-            self.random.try_fill_bytes(&mut tmp)?;
-            buf.copy_from_slice(&tmp)?;
+            let mut copied: u32 = 0;
+            while copied < buf.len() {
+                let len = (buf.len() - copied).min(MAX_SHARED_BUFFER_SIZE as u32);
+                let mut tmp = vec![0; len as usize];
+                self.random.try_fill_bytes(&mut tmp)?;
+                let dest = buf
+                    .get_range(copied..copied + len)
+                    .unwrap()
+                    .as_unsafe_slice_mut()?;
+                dest.copy_from_slice(&tmp)?;
+                copied += len;
+            }
         } else {
             // If the Wasm memory is non-shared, copy directly into the linear
             // memory.
diff --git a/crates/wiggle/src/lib.rs b/crates/wiggle/src/lib.rs
index 92a6c643830d..9b7451bd7139 100644
--- a/crates/wiggle/src/lib.rs
+++ b/crates/wiggle/src/lib.rs
@@ -482,6 +482,11 @@ impl<'a, T: ?Sized + Pointee> GuestPtr<'a, T> {
     {
         GuestPtr::new(self.mem, (self.pointer, elems))
     }
+
+    /// Check if this pointer references WebAssembly shared memory.
+    pub fn is_shared_memory(&self) -> bool {
+        self.mem.is_shared_memory()
+    }
 }
 
 impl<'a, T> GuestPtr<'a, [T]> {