From 0e6eef71c555955798e4c5840472ad226a77102f Mon Sep 17 00:00:00 2001
From: Liu Jiang <gerry@linux.alibaba.com>
Date: Tue, 7 Jul 2020 11:38:45 +0800
Subject: [PATCH] Add check_range() to validate guest address range

Current implementation of GuestMemory::checked_offset() validates
the GuestAddress(base + offset) is valid. By auditting invocations of
GuestMemory::checked_offset(), most callers expect that the whole range
[base, base + offset) is valid. So add a new interface check_range()
to validate the address range.

Signed-off-by: Liu Jiang <gerry@linux.alibaba.com>
---
 coverage_config_x86_64.json |  2 +-
 src/guest_memory.rs         | 15 +++++++++++
 src/mmap.rs                 | 53 +++++++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/coverage_config_x86_64.json b/coverage_config_x86_64.json
index fe5f2670..a3936db6 100644
--- a/coverage_config_x86_64.json
+++ b/coverage_config_x86_64.json
@@ -1,5 +1,5 @@
 {
-  "coverage_score": 85.2,
+  "coverage_score": 84.9,
   "exclude_path": "mmap_windows.rs",
   "crate_features": "backend-mmap,backend-atomic"
 }
diff --git a/src/guest_memory.rs b/src/guest_memory.rs
index 19de1a6b..a70cd25b 100644
--- a/src/guest_memory.rs
+++ b/src/guest_memory.rs
@@ -513,6 +513,21 @@ pub trait GuestMemory {
         self.find_region(addr).map(|_| addr)
     }
 
+    /// Returns the GuestAddress(base + offset) if the range [base, base + offset) is valid.
+    fn check_range(&self, base: GuestAddress, offset: usize) -> Option<GuestAddress> {
+        let end = base.checked_add(offset as u64)?;
+        let region = self.find_region(base)?;
+        if end <= region.last_addr() {
+            return Some(end);
+        }
+
+        let len = offset.checked_add(1)?;
+        match self.try_access(len, base, |_, count, _, _| -> Result<usize> { Ok(count) }) {
+            Ok(count) if count == len => Some(end),
+            _ => None,
+        }
+    }
+
     /// Returns the address plus the offset if it is present within the memory of the guest.
     fn checked_offset(&self, base: GuestAddress, offset: usize) -> Option<GuestAddress> {
         base.checked_add(offset as u64)
diff --git a/src/mmap.rs b/src/mmap.rs
index 406bc867..ad7916d0 100644
--- a/src/mmap.rs
+++ b/src/mmap.rs
@@ -1394,4 +1394,57 @@ mod tests {
         assert!(guest_mem.get_slice(GuestAddress(0x600), 0x100).is_err());
         assert!(guest_mem.get_slice(GuestAddress(0xc00), 0x100).is_err());
     }
+
+    #[test]
+    fn test_checked_offset() {
+        let start_addr1 = GuestAddress(0);
+        let start_addr2 = GuestAddress(0x800);
+        let start_addr3 = GuestAddress(0xc00);
+        let guest_mem = GuestMemoryMmap::from_ranges(&[
+            (start_addr1, 0x400),
+            (start_addr2, 0x400),
+            (start_addr3, 0x400),
+        ])
+        .unwrap();
+
+        assert_eq!(
+            guest_mem.checked_offset(start_addr1, 0x200),
+            Some(GuestAddress(0x200))
+        );
+        assert_eq!(
+            guest_mem.checked_offset(start_addr1, 0xa00),
+            Some(GuestAddress(0xa00))
+        );
+        assert_eq!(
+            guest_mem.checked_offset(start_addr2, 0x7ff),
+            Some(GuestAddress(0xfff))
+        );
+        assert_eq!(guest_mem.checked_offset(start_addr2, 0xc00), None);
+        assert_eq!(guest_mem.checked_offset(start_addr1, std::usize::MAX), None);
+    }
+
+    #[test]
+    fn test_check_range() {
+        let start_addr1 = GuestAddress(0);
+        let start_addr2 = GuestAddress(0x800);
+        let start_addr3 = GuestAddress(0xc00);
+        let guest_mem = GuestMemoryMmap::from_ranges(&[
+            (start_addr1, 0x400),
+            (start_addr2, 0x400),
+            (start_addr3, 0x400),
+        ])
+        .unwrap();
+
+        assert_eq!(
+            guest_mem.check_range(start_addr1, 0x200),
+            Some(GuestAddress(0x200))
+        );
+        assert_eq!(guest_mem.check_range(start_addr1, 0xa00), None);
+        assert_eq!(
+            guest_mem.check_range(start_addr2, 0x7ff),
+            Some(GuestAddress(0xfff))
+        );
+        assert_eq!(guest_mem.check_range(start_addr2, 0xc00), None);
+        assert_eq!(guest_mem.check_range(start_addr1, std::usize::MAX), None);
+    }
 }