Add checked arithmetic operations traits

src/addr.rs

@@ -1,5 +1,6 @@
 //! Physical and virtual addresses manipulation
 
+use crate::ops::{CheckedAdd, CheckedSub};
 use core::fmt;
 use core::ops::{Add, AddAssign, Sub, SubAssign};
 
@@ -265,6 +266,14 @@ impl AddAssign<u64> for VirtAddr {
     }
 }
 
+impl CheckedAdd<u64> for VirtAddr {
+    type Output = Self;
+    #[inline]
+    fn checked_add(self, rhs: u64) -> Option<Self::Output> {
+        self.0.checked_add(rhs).and_then(from_canonical)
+    }
+}
+
 #[cfg(target_pointer_width = "64")]
 impl Add<usize> for VirtAddr {
     type Output = Self;
@@ -297,6 +306,14 @@ impl SubAssign<u64> for VirtAddr {
     }
 }
 
+impl CheckedSub<u64> for VirtAddr {
+    type Output = Self;
+    #[inline]
+    fn checked_sub(self, rhs: u64) -> Option<Self::Output> {
+        self.0.checked_sub(rhs).and_then(from_canonical)
+    }
+}
+
 #[cfg(target_pointer_width = "64")]
 impl Sub<usize> for VirtAddr {
     type Output = Self;
@@ -322,6 +339,14 @@ impl Sub<VirtAddr> for VirtAddr {
     }
 }
 
+impl CheckedSub<VirtAddr> for VirtAddr {
+    type Output = u64;
+    #[inline]
+    fn checked_sub(self, rhs: VirtAddr) -> Option<Self::Output> {
+        self.0.checked_sub(rhs.0)
+    }
+}
+
 /// A passed `u64` was not a valid physical address.
 ///
 /// This means that bits 52 to 64 were not all null.
@@ -479,6 +504,16 @@ impl AddAssign<u64> for PhysAddr {
     }
 }
 
+impl CheckedAdd<u64> for PhysAddr {
+    type Output = Self;
+    #[inline]
+    fn checked_add(self, rhs: u64) -> Option<Self::Output> {
+        self.0
+            .checked_add(rhs)
+            .and_then(|addr| Self::try_new(addr).ok())
+    }
+}
+
 #[cfg(target_pointer_width = "64")]
 impl Add<usize> for PhysAddr {
     type Output = Self;
@@ -511,6 +546,16 @@ impl SubAssign<u64> for PhysAddr {
     }
 }
 
+impl CheckedSub<u64> for PhysAddr {
+    type Output = Self;
+    #[inline]
+    fn checked_sub(self, rhs: u64) -> Option<Self::Output> {
+        self.0
+            .checked_sub(rhs)
+            .and_then(|addr| Self::try_new(addr).ok())
+    }
+}
+
 #[cfg(target_pointer_width = "64")]
 impl Sub<usize> for PhysAddr {
     type Output = Self;
@@ -536,6 +581,28 @@ impl Sub<PhysAddr> for PhysAddr {
     }
 }
 
+impl CheckedSub<PhysAddr> for PhysAddr {
+    type Output = u64;
+    #[inline]
+    fn checked_sub(self, rhs: Self) -> Option<Self::Output> {
+        self.as_u64().checked_sub(rhs.as_u64())
+    }
+}
+
+/// Checks whether the address is canonical.
+
+/// Returns an `Option<VirtAddr>` if the passed u64 falls within canonical
+/// address space.
+#[inline]
+const fn from_canonical(addr: u64) -> Option<VirtAddr> {
+    let trunc = VirtAddr::new_truncate(addr);
+    if trunc.as_u64() == addr {
+        Some(trunc)
+    } else {
+        None
+    }
+}
+
 /// Align address downwards.
 ///
 /// Returns the greatest `x` with alignment `align` so that `x <= addr`.
@@ -593,4 +660,38 @@ mod tests {
         assert_eq!(align_up(0, 2), 0);
         assert_eq!(align_up(0, 0x8000_0000_0000_0000), 0);
     }
+
+    #[test]
+    pub fn checked_operations_only_succeed_on_canonical_addresses() {
+        let upper_space_upper_bound = 0xffff_ffff_ffff_ffffu64;
+        let upper_space_lower_bound = 0xffff_8000_0000_0000u64;
+        let lower_space_upper_bound = 0x7fff_ffff_ffffu64;
+        let lower_space_lower_bound = 0x0u64;
+
+        assert_eq!(
+            None,
+            VirtAddr::new(upper_space_upper_bound).checked_add(1u64)
+        );
+        assert_eq!(
+            Some(VirtAddr::new(upper_space_lower_bound + 1u64)),
+            VirtAddr::new(upper_space_lower_bound).checked_add(1u64)
+        );
+        assert_eq!(
+            None,
+            VirtAddr::new(upper_space_lower_bound).checked_sub(1u64)
+        );
+
+        assert_eq!(
+            Some(VirtAddr::new(lower_space_upper_bound - 1)),
+            VirtAddr::new(lower_space_upper_bound).checked_sub(1u64)
+        );
+        assert_eq!(
+            Some(VirtAddr::new(1)),
+            VirtAddr::new(lower_space_lower_bound).checked_add(0x1u64)
+        );
+        assert_eq!(
+            None,
+            VirtAddr::new(lower_space_lower_bound).checked_sub(1u64)
+        );
+    }
 }

src/lib.rs

@@ -63,6 +63,7 @@ pub(crate) mod asm;
 
 pub mod addr;
 pub mod instructions;
+pub mod ops;
 pub mod registers;
 pub mod structures;
 

src/ops.rs

@@ -0,0 +1,81 @@
+//! Additional arithmetic operators for working with addresses.
+
+/// A checked equivalent to the `+` operator that returns `None` if the
+/// operation would result in a panic or a wrapping overflow. Similar to
+/// [`std::ops::Add`] an `Output` type _must_ be specified and a `Rhs` type
+/// _may_ be specified (defaults to `Self`).
+///
+/// We define our own trait rather than use
+/// [`num::CheckedAdd`](https://docs.rs/num/latest/num/trait.CheckedAdd.html)
+/// due to num's implementation requiring that `Rhs` and `Output` both be `Self`.
+///
+/// # Example
+///
+/// ```
+/// use x86_64::ops::CheckedAdd;
+///
+/// #[derive(Debug, PartialEq)]
+/// struct Addr(u64);
+///
+/// impl CheckedAdd for Addr {
+///     type Output = Self;
+///
+///     fn checked_add(self, other: Self) -> Option<Self> {
+///         self.0.checked_add(other.0).map(Addr)
+///     }
+/// }
+///
+/// assert_eq!(Addr(1).checked_add(Addr(2)),
+///            Some(Addr(3)));
+///
+/// // Overflowing add
+/// assert_eq!(Addr(u64::MAX).checked_add(Addr(1)),
+///            None);
+/// ```
+pub trait CheckedAdd<Rhs = Self> {
+    /// The resulting type returned by the `checked_add` operation.
+    type Output;
+
+    /// Adds two numbers, checking for overflow. If overflow happens, None is returned.
+    fn checked_add(self, rhs: Rhs) -> Option<Self::Output>;
+}
+
+/// A checked equivalent to the `-` operator that returns `None` if the
+/// operation would result in a panic or a wrapping underflow. Similar to
+/// [`std::ops::Sub`] an `Output` type _must_ be specified and a `Rhs` type
+/// _may_ be specified (defaults to `Self`).
+///
+/// We define our own trait rather than use
+/// [`num::CheckedSub`](https://docs.rs/num/latest/num/trait.CheckedSub.html)
+/// due to num's implementation requiring that `Rhs` and `Output` both be `Self`.
+///
+/// # Example
+///
+/// ```
+/// use x86_64::ops::CheckedSub;
+///
+/// #[derive(Debug, PartialEq)]
+/// struct Addr(u64);
+///
+/// impl CheckedSub for Addr {
+///     type Output = Self;
+///
+///     fn checked_sub(self, other: Self) -> Option<Self> {
+///         self.0.checked_sub(other.0).map(Addr)
+///     }
+/// }
+///
+/// assert_eq!(Addr(3).checked_sub(Addr(2)),
+///            Some(Addr(1)));
+///
+/// // Undeflowing sub
+/// assert_eq!(Addr(u64::MIN).checked_sub(Addr(1)),
+///            None);
+/// ```
+pub trait CheckedSub<Rhs = Self> {
+    /// The resulting type returned by the `checked_sub` operation.
+    type Output;
+
+    /// Subtracts two numbers, checking for underflow. If underflow happens, None is returned.
+    fn checked_sub(self, rhs: Rhs) -> Option<Self::Output>;
+}