sys_util: add SeekHole trait

Allow seeking to the next hole or data region in File and QcowFile.

BUG=None
TEST=None

Change-Id: I16e77e4791aa85b4cc96f38327026cd93f02b7e1
Signed-off-by: Daniel Verkamp <dverkamp@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/1274147
Commit-Ready: ChromeOS CL Exonerator Bot <chromiumos-cl-exonerator@appspot.gserviceaccount.com>
Reviewed-by: Zach Reizner <zachr@chromium.org>

3 files changed, 418 insertions, 1 deletions

diff --git a/qcow/src/qcow.rs b/qcow/src/qcow.rs
index 0e08c14..f789b30 100644
--- a/qcow/src/qcow.rs
+++ b/qcow/src/qcow.rs
@@ -23,7 +23,7 @@ use std::io::{self, Read, Seek, SeekFrom, Write};
 use std::mem::size_of;
 use std::os::unix::io::{AsRawFd, RawFd};
 
-use sys_util::{fallocate, FallocateMode, WriteZeroes};
+use sys_util::{fallocate, FallocateMode, SeekHole, WriteZeroes};
 
 #[derive(Debug)]
 pub enum Error {
@@ -669,6 +669,77 @@ impl QcowFile {
         Ok(new_addr)
     }
 
+    // Returns true if the cluster containing `address` is already allocated.
+    fn cluster_allocated(&mut self, address: u64) -> std::io::Result<bool> {
+        if address >= self.virtual_size() as u64 {
+            return Err(std::io::Error::from_raw_os_error(EINVAL));
+        }
+
+        let l1_index = self.l1_table_index(address) as usize;
+        let l2_addr_disk = *self
+            .l1_table
+            .get(l1_index)
+            .ok_or(std::io::Error::from_raw_os_error(EINVAL))?;
+        let l2_index = self.l2_table_index(address) as usize;
+
+        if l2_addr_disk == 0 {
+            // The whole L2 table for this address is not allocated yet,
+            // so the cluster must also be unallocated.
+            return Ok(false);
+        }
+
+        if !self.l2_cache.contains_key(&l1_index) {
+            // Not in the cache.
+            let table =
+                VecCache::from_vec(Self::read_l2_cluster(&mut self.raw_file, l2_addr_disk)?);
+            let l1_table = &self.l1_table;
+            let raw_file = &mut self.raw_file;
+            self.l2_cache.insert(l1_index, table, |index, evicted| {
+                raw_file.write_pointer_table(
+                    l1_table[index],
+                    evicted.get_values(),
+                    CLUSTER_USED_FLAG,
+                )
+            })?;
+        }
+
+        let cluster_addr = self.l2_cache.get(&l1_index).unwrap()[l2_index];
+        // If cluster_addr != 0, the cluster is allocated.
+        Ok(cluster_addr != 0)
+    }
+
+    // Find the first guest address greater than or equal to `address` whose allocation state
+    // matches `allocated`.
+    fn find_allocated_cluster(
+        &mut self,
+        address: u64,
+        allocated: bool,
+    ) -> std::io::Result<Option<u64>> {
+        let size = self.virtual_size();
+        if address >= size {
+            return Ok(None);
+        }
+
+        // If offset is already within a hole, return it.
+        if self.cluster_allocated(address)? == allocated {
+            return Ok(Some(address));
+        }
+
+        // Skip to the next cluster boundary.
+        let cluster_size = self.raw_file.cluster_size();
+        let mut cluster_addr = (address / cluster_size + 1) * cluster_size;
+
+        // Search for clusters with the desired allocation state.
+        while cluster_addr < size {
+            if self.cluster_allocated(cluster_addr)? == allocated {
+                return Ok(Some(cluster_addr));
+            }
+            cluster_addr += cluster_size;
+        }
+
+        Ok(None)
+    }
+
     // Deallocate the storage for the cluster starting at `address`.
     // Any future reads of this cluster will return all zeroes.
     fn deallocate_cluster(&mut self, address: u64) -> std::io::Result<()> {
@@ -995,6 +1066,36 @@ impl WriteZeroes for QcowFile {
     }
 }
 
+impl SeekHole for QcowFile {
+    fn seek_hole(&mut self, offset: u64) -> io::Result<Option<u64>> {
+        match self.find_allocated_cluster(offset, false) {
+            Err(e) => Err(e),
+            Ok(None) => {
+                if offset < self.virtual_size() {
+                    Ok(Some(self.seek(SeekFrom::End(0))?))
+                } else {
+                    Ok(None)
+                }
+            }
+            Ok(Some(o)) => {
+                self.seek(SeekFrom::Start(o))?;
+                Ok(Some(o))
+            }
+        }
+    }
+
+    fn seek_data(&mut self, offset: u64) -> io::Result<Option<u64>> {
+        match self.find_allocated_cluster(offset, true) {
+            Err(e) => Err(e),
+            Ok(None) => Ok(None),
+            Ok(Some(o)) => {
+                self.seek(SeekFrom::Start(o))?;
+                Ok(Some(o))
+            }
+        }
+    }
+}
+
 // Returns an Error if the given offset doesn't align to a cluster boundary.
 fn offset_is_cluster_boundary(offset: u64, cluster_bits: u32) -> Result<()> {
     if offset & ((0x01 << cluster_bits) - 1) != 0 {
@@ -1620,4 +1721,115 @@ mod tests {
             assert_eq!(qcow_file.first_zero_refcount().unwrap(), None);
         });
     }
+
+    fn seek_cur(file: &mut QcowFile) -> u64 {
+        file.seek(SeekFrom::Current(0)).unwrap()
+    }
+
+    #[test]
+    fn seek_data() {
+        with_default_file(0x30000, |mut file| {
+            // seek_data at or after the end of the file should return None
+            assert_eq!(file.seek_data(0x10000).unwrap(), None);
+            assert_eq!(seek_cur(&mut file), 0);
+            assert_eq!(file.seek_data(0x10001).unwrap(), None);
+            assert_eq!(seek_cur(&mut file), 0);
+
+            // Write some data to [0x10000, 0x20000)
+            let b = [0x55u8; 0x10000];
+            file.seek(SeekFrom::Start(0x10000)).unwrap();
+            file.write_all(&b).unwrap();
+            assert_eq!(file.seek_data(0).unwrap(), Some(0x10000));
+            assert_eq!(seek_cur(&mut file), 0x10000);
+
+            // seek_data within data should return the same offset
+            assert_eq!(file.seek_data(0x10000).unwrap(), Some(0x10000));
+            assert_eq!(seek_cur(&mut file), 0x10000);
+            assert_eq!(file.seek_data(0x10001).unwrap(), Some(0x10001));
+            assert_eq!(seek_cur(&mut file), 0x10001);
+            assert_eq!(file.seek_data(0x1FFFF).unwrap(), Some(0x1FFFF));
+            assert_eq!(seek_cur(&mut file), 0x1FFFF);
+
+            assert_eq!(file.seek_data(0).unwrap(), Some(0x10000));
+            assert_eq!(seek_cur(&mut file), 0x10000);
+            assert_eq!(file.seek_data(0x1FFFF).unwrap(), Some(0x1FFFF));
+            assert_eq!(seek_cur(&mut file), 0x1FFFF);
+            assert_eq!(file.seek_data(0x20000).unwrap(), None);
+            assert_eq!(seek_cur(&mut file), 0x1FFFF);
+        });
+    }
+
+    #[test]
+    fn seek_hole() {
+        with_default_file(0x30000, |mut file| {
+            // File consisting entirely of a hole
+            assert_eq!(file.seek_hole(0).unwrap(), Some(0));
+            assert_eq!(seek_cur(&mut file), 0);
+            assert_eq!(file.seek_hole(0xFFFF).unwrap(), Some(0xFFFF));
+            assert_eq!(seek_cur(&mut file), 0xFFFF);
+
+            // seek_hole at or after the end of the file should return None
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x30000).unwrap(), None);
+            assert_eq!(seek_cur(&mut file), 0);
+            assert_eq!(file.seek_hole(0x30001).unwrap(), None);
+            assert_eq!(seek_cur(&mut file), 0);
+
+            // Write some data to [0x10000, 0x20000)
+            let b = [0x55u8; 0x10000];
+            file.seek(SeekFrom::Start(0x10000)).unwrap();
+            file.write_all(&b).unwrap();
+
+            // seek_hole within a hole should return the same offset
+            assert_eq!(file.seek_hole(0).unwrap(), Some(0));
+            assert_eq!(seek_cur(&mut file), 0);
+            assert_eq!(file.seek_hole(0xFFFF).unwrap(), Some(0xFFFF));
+            assert_eq!(seek_cur(&mut file), 0xFFFF);
+
+            // seek_hole within data should return the next hole
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x10000).unwrap(), Some(0x20000));
+            assert_eq!(seek_cur(&mut file), 0x20000);
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x10001).unwrap(), Some(0x20000));
+            assert_eq!(seek_cur(&mut file), 0x20000);
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x1FFFF).unwrap(), Some(0x20000));
+            assert_eq!(seek_cur(&mut file), 0x20000);
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0xFFFF).unwrap(), Some(0xFFFF));
+            assert_eq!(seek_cur(&mut file), 0xFFFF);
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x10000).unwrap(), Some(0x20000));
+            assert_eq!(seek_cur(&mut file), 0x20000);
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x1FFFF).unwrap(), Some(0x20000));
+            assert_eq!(seek_cur(&mut file), 0x20000);
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x20000).unwrap(), Some(0x20000));
+            assert_eq!(seek_cur(&mut file), 0x20000);
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x20001).unwrap(), Some(0x20001));
+            assert_eq!(seek_cur(&mut file), 0x20001);
+
+            // seek_hole at EOF should return None
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x30000).unwrap(), None);
+            assert_eq!(seek_cur(&mut file), 0);
+
+            // Write some data to [0x20000, 0x30000)
+            file.seek(SeekFrom::Start(0x20000)).unwrap();
+            file.write_all(&b).unwrap();
+
+            // seek_hole within [0x20000, 0x30000) should now find the hole at EOF
+            assert_eq!(file.seek_hole(0x20000).unwrap(), Some(0x30000));
+            assert_eq!(seek_cur(&mut file), 0x30000);
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x20001).unwrap(), Some(0x30000));
+            assert_eq!(seek_cur(&mut file), 0x30000);
+            file.seek(SeekFrom::Start(0)).unwrap();
+            assert_eq!(file.seek_hole(0x30000).unwrap(), None);
+            assert_eq!(seek_cur(&mut file), 0);
+        });
+    }
 }
diff --git a/sys_util/src/lib.rs b/sys_util/src/lib.rs
index 94ccebc..e2cf0e6 100644
--- a/sys_util/src/lib.rs
+++ b/sys_util/src/lib.rs
@@ -26,6 +26,7 @@ mod guest_memory;
 mod mmap;
 mod passwd;
 mod poll;
+mod seek_hole;
 mod shm;
 pub mod signal;
 mod signalfd;
@@ -59,6 +60,7 @@ pub use timerfd::*;
 
 pub use guest_memory::Error as GuestMemoryError;
 pub use mmap::Error as MmapError;
+pub use seek_hole::SeekHole;
 pub use signalfd::Error as SignalFdError;
 pub use write_zeroes::WriteZeroes;
 
diff --git a/sys_util/src/seek_hole.rs b/sys_util/src/seek_hole.rs
new file mode 100644
index 0000000..15fee58
--- /dev/null
+++ b/sys_util/src/seek_hole.rs
@@ -0,0 +1,203 @@
+// Copyright 2018 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use std::fs::File;
+use std::io::{Error, Result};
+use std::os::unix::io::AsRawFd;
+
+use libc::{lseek64, ENXIO, SEEK_DATA, SEEK_HOLE};
+
+/// A trait for seeking to the next hole or non-hole position in a file.
+pub trait SeekHole {
+    /// Seek to the first hole in a file at a position greater than or equal to `offset`.
+    /// If no holes exist after `offset`, the seek position will be set to the end of the file.
+    /// If `offset` is at or after the end of the file, the seek position is unchanged, and None is returned.
+    /// Returns the current seek position after the seek or an error.
+    fn seek_hole(&mut self, offset: u64) -> Result<Option<u64>>;
+
+    /// Seek to the first data in a file at a position greater than or equal to `offset`.
+    /// If no data exists after `offset`, the seek position is unchanged, and None is returned.
+    /// Returns the current offset after the seek or an error.
+    fn seek_data(&mut self, offset: u64) -> Result<Option<u64>>;
+}
+
+/// Safe wrapper for `libc::lseek64()`
+fn lseek(file: &mut File, offset: i64, whence: i32) -> Result<Option<u64>> {
+    // This is safe because we pass a known-good file descriptor.
+    let res = unsafe { lseek64(file.as_raw_fd(), offset, whence) };
+
+    if res < 0 {
+        // Convert ENXIO into None; pass any other error as-is.
+        let err = Error::last_os_error();
+        if let Some(errno) = Error::raw_os_error(&err) {
+            if errno == ENXIO {
+                return Ok(None);
+            }
+        }
+        Err(err)
+    } else {
+        Ok(Some(res as u64))
+    }
+}
+
+impl SeekHole for File {
+    fn seek_hole(&mut self, offset: u64) -> Result<Option<u64>> {
+        lseek(self, offset as i64, SEEK_HOLE)
+    }
+
+    fn seek_data(&mut self, offset: u64) -> Result<Option<u64>> {
+        lseek(self, offset as i64, SEEK_DATA)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs::File;
+    use std::io::{Seek, SeekFrom, Write};
+    use std::path::PathBuf;
+    use TempDir;
+
+    fn seek_cur(file: &mut File) -> u64 {
+        file.seek(SeekFrom::Current(0)).unwrap()
+    }
+
+    #[test]
+    fn seek_data() {
+        let tempdir = TempDir::new("/tmp/seek_data_test").unwrap();
+        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        path.push("test_file");
+        let mut file = File::create(&path).unwrap();
+
+        // Empty file
+        assert_eq!(file.seek_data(0).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+
+        // File with non-zero length consisting entirely of a hole
+        file.set_len(0x10000).unwrap();
+        assert_eq!(file.seek_data(0).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+
+        // seek_data at or after the end of the file should return None
+        assert_eq!(file.seek_data(0x10000).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+        assert_eq!(file.seek_data(0x10001).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+
+        // Write some data to [0x10000, 0x20000)
+        let b = [0x55u8; 0x10000];
+        file.seek(SeekFrom::Start(0x10000)).unwrap();
+        file.write_all(&b).unwrap();
+        assert_eq!(file.seek_data(0).unwrap(), Some(0x10000));
+        assert_eq!(seek_cur(&mut file), 0x10000);
+
+        // seek_data within data should return the same offset
+        assert_eq!(file.seek_data(0x10000).unwrap(), Some(0x10000));
+        assert_eq!(seek_cur(&mut file), 0x10000);
+        assert_eq!(file.seek_data(0x10001).unwrap(), Some(0x10001));
+        assert_eq!(seek_cur(&mut file), 0x10001);
+        assert_eq!(file.seek_data(0x1FFFF).unwrap(), Some(0x1FFFF));
+        assert_eq!(seek_cur(&mut file), 0x1FFFF);
+
+        // Extend the file to add another hole after the data
+        file.set_len(0x30000).unwrap();
+        assert_eq!(file.seek_data(0).unwrap(), Some(0x10000));
+        assert_eq!(seek_cur(&mut file), 0x10000);
+        assert_eq!(file.seek_data(0x1FFFF).unwrap(), Some(0x1FFFF));
+        assert_eq!(seek_cur(&mut file), 0x1FFFF);
+        assert_eq!(file.seek_data(0x20000).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0x1FFFF);
+    }
+
+    #[test]
+    fn seek_hole() {
+        let tempdir = TempDir::new("/tmp/seek_hole_test").unwrap();
+        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        path.push("test_file");
+        let mut file = File::create(&path).unwrap();
+
+        // Empty file
+        assert_eq!(file.seek_hole(0).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+
+        // File with non-zero length consisting entirely of a hole
+        file.set_len(0x10000).unwrap();
+        assert_eq!(file.seek_hole(0).unwrap(), Some(0));
+        assert_eq!(seek_cur(&mut file), 0);
+        assert_eq!(file.seek_hole(0xFFFF).unwrap(), Some(0xFFFF));
+        assert_eq!(seek_cur(&mut file), 0xFFFF);
+
+        // seek_hole at or after the end of the file should return None
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x10000).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+        assert_eq!(file.seek_hole(0x10001).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+
+        // Write some data to [0x10000, 0x20000)
+        let b = [0x55u8; 0x10000];
+        file.seek(SeekFrom::Start(0x10000)).unwrap();
+        file.write_all(&b).unwrap();
+
+        // seek_hole within a hole should return the same offset
+        assert_eq!(file.seek_hole(0).unwrap(), Some(0));
+        assert_eq!(seek_cur(&mut file), 0);
+        assert_eq!(file.seek_hole(0xFFFF).unwrap(), Some(0xFFFF));
+        assert_eq!(seek_cur(&mut file), 0xFFFF);
+
+        // seek_hole within data should return the next hole (EOF)
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x10000).unwrap(), Some(0x20000));
+        assert_eq!(seek_cur(&mut file), 0x20000);
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x10001).unwrap(), Some(0x20000));
+        assert_eq!(seek_cur(&mut file), 0x20000);
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x1FFFF).unwrap(), Some(0x20000));
+        assert_eq!(seek_cur(&mut file), 0x20000);
+
+        // seek_hole at EOF after data should return None
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x20000).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+
+        // Extend the file to add another hole after the data
+        file.set_len(0x30000).unwrap();
+        assert_eq!(file.seek_hole(0).unwrap(), Some(0));
+        assert_eq!(seek_cur(&mut file), 0);
+        assert_eq!(file.seek_hole(0xFFFF).unwrap(), Some(0xFFFF));
+        assert_eq!(seek_cur(&mut file), 0xFFFF);
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x10000).unwrap(), Some(0x20000));
+        assert_eq!(seek_cur(&mut file), 0x20000);
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x1FFFF).unwrap(), Some(0x20000));
+        assert_eq!(seek_cur(&mut file), 0x20000);
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x20000).unwrap(), Some(0x20000));
+        assert_eq!(seek_cur(&mut file), 0x20000);
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x20001).unwrap(), Some(0x20001));
+        assert_eq!(seek_cur(&mut file), 0x20001);
+
+        // seek_hole at EOF after a hole should return None
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x30000).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+
+        // Write some data to [0x20000, 0x30000)
+        file.seek(SeekFrom::Start(0x20000)).unwrap();
+        file.write_all(&b).unwrap();
+
+        // seek_hole within [0x20000, 0x30000) should now find the hole at EOF
+        assert_eq!(file.seek_hole(0x20000).unwrap(), Some(0x30000));
+        assert_eq!(seek_cur(&mut file), 0x30000);
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x20001).unwrap(), Some(0x30000));
+        assert_eq!(seek_cur(&mut file), 0x30000);
+        file.seek(SeekFrom::Start(0)).unwrap();
+        assert_eq!(file.seek_hole(0x30000).unwrap(), None);
+        assert_eq!(seek_cur(&mut file), 0);
+    }
+}