From 6319451d5896604437dbc588e447e5f04da38423 Mon Sep 17 00:00:00 2001
From: Andreea Florescu <fandree@amazon.com>
Date: Tue, 28 Nov 2023 17:28:56 +0100
Subject: [PATCH] try_access: use checked_add instead of overflowing_add

We were allowing overflows if the overflow value would've been 0.
This can be problematic if there is a memory region that starts
at GuestAddress(0x0), but the access was not intended to spill
into that region. We don't consider memory to be a circular ring
buffer which would be the only situation in which we should allow
overflowing accesses.

While at it, also refactor the checked_add for the toal number of
bytes to use `ok_or` instead of a match pattern, as this is
easier to read.

Writing a regression test for the overflowing add is not possible
with the current infrastructure because GuestMemoryMmap does not
allow creating memory regions with an access that would overflow
(i.e. having a memory region starting at u64::MAX of size 1).
Also, we cannot call try_access with a size bigger than the
memory region size because we only allow accesses that fit in
the memory region.

Signed-off-by: Andreea Florescu <fandree@amazon.com>
---
 src/guest_memory.rs | 35 ++++++++++++++++++++++++++---------
 1 file changed, 26 insertions(+), 9 deletions(-)

diff --git a/src/guest_memory.rs b/src/guest_memory.rs
index 251735ab..904cdc9a 100644
--- a/src/guest_memory.rs
+++ b/src/guest_memory.rs
@@ -654,15 +654,13 @@ pub trait GuestMemory {
                 Ok(0) => return Ok(total),
                 // made some progress
                 Ok(len) => {
-                    total = match total.checked_add(len) {
-                        Some(x) if x < count => x,
-                        Some(x) if x == count => return Ok(x),
-                        _ => return Err(Error::CallbackOutOfRange),
-                    };
-                    cur = match cur.overflowing_add(len as GuestUsize) {
-                        (x @ GuestAddress(0), _) | (x, false) => x,
-                        (_, true) => return Err(Error::GuestAddressOverflow),
-                    };
+                    total = total.checked_add(len).ok_or(Error::CallbackOutOfRange)?;
+                    if total == count {
+                        return Ok(total);
+                    }
+                    cur = cur
+                        .checked_add(len as GuestUsize)
+                        .ok_or(Error::GuestAddressOverflow)?;
                 }
                 // error happened
                 e => return e,
@@ -1304,6 +1302,25 @@ mod tests {
         non_atomic_access_helper::<u16>()
     }
 
+    #[cfg(feature = "backend-mmap")]
+    #[test]
+    // This test makes sure that computation for regions that end at u64::MAX do not cause
+    // overflows.
+    fn test_region_at_usize_max() {
+        let region_len: usize = 0x1000;
+        let addr = GuestAddress(u64::MAX - region_len as u64);
+        let mem = GuestMemoryMmap::from_ranges(&[(addr, region_len)]).unwrap();
+        let mut image = vec![0; region_len + 1];
+
+        // This access could cause an overflow if we wouldn't have a check for memory region length.
+        let invalid_access_size = region_len + 1;
+        let count = mem
+            .write_volatile_to(addr, &mut image.as_mut_slice(), invalid_access_size)
+            .unwrap();
+        // We can only write maximum region_len bytes.
+        assert_eq!(count, region_len);
+    }
+
     #[cfg(feature = "backend-mmap")]
     #[test]
     fn test_zero_length_accesses() {