sys_util: add write_zeroes_all() function

In the same spirit as write_all() for the standard io::Write::write()
function, add a write_zeroes_all() function with a default
implementation that calls write_zeroes() in a loop until the requested
length is met.  This will allow write_zeroes implementations that don't
necessarily fulfill the entire requested length.

BUG=None
TEST=cargo test -p sys_util write_zeroes

Change-Id: I0fc3a4b3fe8904946e253ab8a2687555b12657be
Signed-off-by: Daniel Verkamp <dverkamp@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/1811466
Reviewed-by: Zach Reizner <zachr@chromium.org>
Reviewed-by: Cody Schuffelen <schuffelen@google.com>
Tested-by: kokoro <noreply+kokoro@google.com>

5 files changed, 36 insertions, 14 deletions

diff --git a/devices/src/virtio/block.rs b/devices/src/virtio/block.rs
index 5939652..73fd416 100644
--- a/devices/src/virtio/block.rs
+++ b/devices/src/virtio/block.rs
@@ -673,7 +673,7 @@ impl Block {
                     } else {
                         disk.seek(SeekFrom::Start(offset))
                             .map_err(|e| ExecuteError::Seek { ioerr: e, sector })?;
-                        disk.write_zeroes(length as usize).map_err(|e| {
+                        disk.write_zeroes_all(length as usize).map_err(|e| {
                             ExecuteError::DiscardWriteZeroes {
                                 ioerr: Some(e),
                                 sector,
diff --git a/qcow/src/qcow.rs b/qcow/src/qcow.rs
index aaff5f1..633b21c 100644
--- a/qcow/src/qcow.rs
+++ b/qcow/src/qcow.rs
@@ -1239,7 +1239,7 @@ impl QcowFile {
                 if let Some(offset) = self.file_offset_read(curr_addr)? {
                     // Partial cluster - zero it out.
                     self.raw_file.file_mut().seek(SeekFrom::Start(offset))?;
-                    self.raw_file.file_mut().write_zeroes(count)?;
+                    self.raw_file.file_mut().write_zeroes_all(count)?;
                 }
             }
 
@@ -1822,8 +1822,7 @@ mod tests {
             q.write(&b).expect("Failed to write test string.");
             // Overwrite the test data with zeroes.
             q.seek(SeekFrom::Start(0xfff2000)).expect("Failed to seek.");
-            let nwritten = q.write_zeroes(0x200).expect("Failed to write zeroes.");
-            assert_eq!(nwritten, 0x200);
+            q.write_zeroes_all(0x200).expect("Failed to write zeroes.");
             // Verify that the correct part of the data was zeroed out.
             let mut buf = [0u8; 0x1000];
             q.seek(SeekFrom::Start(0xfff2000)).expect("Failed to seek.");
@@ -1848,8 +1847,8 @@ mod tests {
             q.write(&b).expect("Failed to write test string.");
             // Overwrite the full cluster with zeroes.
             q.seek(SeekFrom::Start(0)).expect("Failed to seek.");
-            let nwritten = q.write_zeroes(CHUNK_SIZE).expect("Failed to write zeroes.");
-            assert_eq!(nwritten, CHUNK_SIZE);
+            q.write_zeroes_all(CHUNK_SIZE)
+                .expect("Failed to write zeroes.");
             // Verify that the data was zeroed out.
             let mut buf = [0u8; CHUNK_SIZE];
             q.seek(SeekFrom::Start(0)).expect("Failed to seek.");
diff --git a/qcow/src/qcow_raw_file.rs b/qcow/src/qcow_raw_file.rs
index 99a8d71..ede28d8 100644
--- a/qcow/src/qcow_raw_file.rs
+++ b/qcow/src/qcow_raw_file.rs
@@ -142,7 +142,7 @@ impl QcowRawFile {
     pub fn zero_cluster(&mut self, address: u64) -> io::Result<()> {
         let cluster_size = self.cluster_size as usize;
         self.file.seek(SeekFrom::Start(address))?;
-        self.file.write_zeroes(cluster_size)?;
+        self.file.write_zeroes_all(cluster_size)?;
         Ok(())
     }
 }
diff --git a/qcow_utils/src/qcow_img.rs b/qcow_utils/src/qcow_img.rs
index 08040a1..b4ad422 100644
--- a/qcow_utils/src/qcow_img.rs
+++ b/qcow_utils/src/qcow_img.rs
@@ -264,7 +264,7 @@ fn dd(file_path: &str, source_path: &str, count: Option<usize>) -> std::result::
         let nread = src_file.read(&mut buf[..this_count]).map_err(|_| ())?;
         // If this block is all zeros, then use write_zeros so the output file is sparse.
         if buf.iter().all(|b| *b == 0) {
-            qcow_file.write_zeroes(CHUNK_SIZE).map_err(|_| ())?;
+            qcow_file.write_zeroes_all(CHUNK_SIZE).map_err(|_| ())?;
         } else {
             qcow_file.write(&buf).map_err(|_| ())?;
         }
diff --git a/sys_util/src/write_zeroes.rs b/sys_util/src/write_zeroes.rs
index e3b531e..0e733c7 100644
--- a/sys_util/src/write_zeroes.rs
+++ b/sys_util/src/write_zeroes.rs
@@ -4,7 +4,7 @@
 
 use std::cmp::min;
 use std::fs::File;
-use std::io::{self, Seek, SeekFrom, Write};
+use std::io::{self, Error, ErrorKind, Seek, SeekFrom, Write};
 
 use crate::fallocate;
 use crate::FallocateMode;
@@ -24,8 +24,31 @@ impl PunchHole for File {
 
 /// A trait for writing zeroes to a stream.
 pub trait WriteZeroes {
-    /// Write `length` bytes of zeroes to the stream, returning how many bytes were written.
+    /// Write up to `length` bytes of zeroes to the stream, returning how many bytes were written.
     fn write_zeroes(&mut self, length: usize) -> io::Result<usize>;
+
+    /// Write zeroes to the stream until `length` bytes have been written.
+    ///
+    /// This method will continuously call `write_zeroes` until the requested
+    /// `length` is satisfied or an error is encountered.
+    fn write_zeroes_all(&mut self, mut length: usize) -> io::Result<()> {
+        while length > 0 {
+            match self.write_zeroes(length) {
+                Ok(0) => return Err(Error::from(ErrorKind::WriteZero)),
+                Ok(bytes_written) => {
+                    length = length
+                        .checked_sub(bytes_written)
+                        .ok_or(Error::from(ErrorKind::Other))?
+                }
+                Err(e) => {
+                    if e.kind() != ErrorKind::Interrupted {
+                        return Err(e);
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
 }
 
 impl<T: PunchHole + Seek + Write> WriteZeroes for T {
@@ -98,8 +121,8 @@ mod tests {
 
         // Overwrite some of the data with zeroes
         f.seek(SeekFrom::Start(2345)).unwrap();
-        f.write_zeroes(4321).expect("write_zeroes failed");
-        // Verify seek position after write_zeroes()
+        f.write_zeroes_all(4321).expect("write_zeroes failed");
+        // Verify seek position after write_zeroes_all()
         assert_eq!(f.seek(SeekFrom::Current(0)).unwrap(), 2345 + 4321);
 
         // Read back the data and verify that it is now zero
@@ -147,8 +170,8 @@ mod tests {
 
         // Overwrite some of the data with zeroes
         f.seek(SeekFrom::Start(0)).unwrap();
-        f.write_zeroes(0x10001).expect("write_zeroes failed");
-        // Verify seek position after write_zeroes()
+        f.write_zeroes_all(0x10001).expect("write_zeroes failed");
+        // Verify seek position after write_zeroes_all()
         assert_eq!(f.seek(SeekFrom::Current(0)).unwrap(), 0x10001);
 
         // Read back the data and verify that it is now zero