diff --git a/src/addr.rs b/src/addr.rs index c330d60e..9c7efb8f 100644 --- a/src/addr.rs +++ b/src/addr.rs @@ -168,6 +168,50 @@ impl VirtAddr { self.0 == 0 } + #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))] + /// Add an offset to a virtual address. + /// + /// This function performs normal arithmetic addition and doesn't jump the + /// address gap. If you're looking for a successor operation that jumps the + /// address gap, use [`Step::forward`]. + /// + /// This function returns `None` on overflow or if the result is not a + /// canonical address. + #[inline] + pub const fn checked_add(self, rhs: u64) -> Option { + let Some(addr) = self.0.checked_add(rhs) else { + return None; + }; + + let Ok(addr) = Self::try_new(addr) else { + return None; + }; + + Some(addr) + } + + #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))] + /// Subtract an offset from a virtual address. + /// + /// This function performs normal arithmetic subtraction and doesn't jump + /// the address gap. If you're looking for a predecessor operation that + /// jumps the address gap, use [`Step::backward`]. + /// + /// This function returns `None` on overflow or if the result is not a + /// canonical address. + #[inline] + pub const fn checked_sub(self, rhs: u64) -> Option { + let Some(addr) = self.0.checked_sub(rhs) else { + return None; + }; + + let Ok(addr) = Self::try_new(addr) else { + return None; + }; + + Some(addr) + } + /// Aligns the virtual address upwards to the given alignment. /// /// See the `align_up` function for more information. @@ -379,41 +423,21 @@ impl fmt::Pointer for VirtAddr { impl Add for VirtAddr { type Output = Self; - #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))] /// Add an offset to a virtual address. /// - /// This function performs normal arithmetic addition and doesn't jump the - /// address gap. If you're looking for a successor operation that jumps the - /// address gap, use [`Step::forward`]. - /// - /// # Panics - /// - /// This function will panic on overflow or if the result is not a - /// canonical address. + /// This function panics when [`VirtAddr::checked_add`] returns `None`. #[inline] #[track_caller] fn add(self, rhs: u64) -> Self::Output { - VirtAddr::try_new( - self.0 - .checked_add(rhs) - .expect("attempt to add with overflow"), - ) - .expect("attempt to add resulted in non-canonical virtual address") + self.checked_add(rhs) + .expect("attempt to add with overflow or resulted in non-canonical virtual address") } } impl AddAssign for VirtAddr { - #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))] /// Add an offset to a virtual address. /// - /// This function performs normal arithmetic addition and doesn't jump the - /// address gap. If you're looking for a successor operation that jumps the - /// address gap, use [`Step::forward`]. - /// - /// # Panics - /// - /// This function will panic on overflow or if the result is not a - /// canonical address. + /// This function panics when [`VirtAddr::checked_add`] returns `None`. #[inline] #[track_caller] fn add_assign(&mut self, rhs: u64) { @@ -424,41 +448,22 @@ impl AddAssign for VirtAddr { impl Sub for VirtAddr { type Output = Self; - #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))] /// Subtract an offset from a virtual address. /// - /// This function performs normal arithmetic subtraction and doesn't jump - /// the address gap. If you're looking for a predecessor operation that - /// jumps the address gap, use [`Step::backward`]. - /// - /// # Panics - /// - /// This function will panic on overflow or if the result is not a - /// canonical address. + /// This function panics when [`VirtAddr::checked_sub`] returns `None`. #[inline] #[track_caller] fn sub(self, rhs: u64) -> Self::Output { - VirtAddr::try_new( - self.0 - .checked_sub(rhs) - .expect("attempt to subtract with overflow"), + self.checked_sub(rhs).expect( + "attempt to subtract with overflow or resulted in non-canonical virtual address", ) - .expect("attempt to subtract resulted in non-canonical virtual address") } } impl SubAssign for VirtAddr { - #[cfg_attr(not(feature = "step_trait"), allow(rustdoc::broken_intra_doc_links))] /// Subtract an offset from a virtual address. /// - /// This function performs normal arithmetic subtraction and doesn't jump - /// the address gap. If you're looking for a predecessor operation that - /// jumps the address gap, use [`Step::backward`]. - /// - /// # Panics - /// - /// This function will panic on overflow or if the result is not a - /// canonical address. + /// This function panics when [`VirtAddr::checked_sub`] returns `None`. #[inline] #[track_caller] fn sub_assign(&mut self, rhs: u64) { @@ -618,6 +623,42 @@ impl PhysAddr { self.0 == 0 } + /// Add an offset to a physical address. + /// + /// This function returns `None` on overflow or if the result is not a + /// valid physical address. + #[inline] + #[const_fn(cfg(not(feature = "memory_encryption")))] + pub const fn checked_add(self, rhs: u64) -> Option { + let Some(addr) = self.0.checked_add(rhs) else { + return None; + }; + + let Ok(addr) = Self::try_new(addr) else { + return None; + }; + + Some(addr) + } + + /// Subtract an offset from a physical address. + /// + /// This function returns `None` on overflow or if the result is not a + /// valid physical address. + #[inline] + #[const_fn(cfg(not(feature = "memory_encryption")))] + pub const fn checked_sub(self, rhs: u64) -> Option { + let Some(addr) = self.0.checked_sub(rhs) else { + return None; + }; + + let Ok(addr) = Self::try_new(addr) else { + return None; + }; + + Some(addr) + } + /// Aligns the physical address upwards to the given alignment. /// /// See the `align_up` function for more information. @@ -717,7 +758,8 @@ impl Add for PhysAddr { #[inline] #[track_caller] fn add(self, rhs: u64) -> Self::Output { - PhysAddr::new(self.0.checked_add(rhs).unwrap()) + self.checked_add(rhs) + .expect("attempt to add with overflow or resulted in invalid physical address") } } @@ -734,7 +776,8 @@ impl Sub for PhysAddr { #[inline] #[track_caller] fn sub(self, rhs: u64) -> Self::Output { - PhysAddr::new(self.0.checked_sub(rhs).unwrap()) + self.checked_sub(rhs) + .expect("attempt to subtract with overflow or resulted in invalid physical address") } } diff --git a/src/structures/paging/frame.rs b/src/structures/paging/frame.rs index 01d3c03d..090a8363 100644 --- a/src/structures/paging/frame.rs +++ b/src/structures/paging/frame.rs @@ -155,6 +155,42 @@ impl PhysFrame { self.start_address.as_u64() / S::SIZE } + /// Add an offset to a frame. + /// + /// This function returns `None` on overflow or if the result does not have a + /// valid physical address. + #[inline] + #[const_fn(cfg(not(feature = "memory_encryption")))] + pub const fn checked_add(self, rhs: u64) -> Option { + let Some(offset) = rhs.checked_mul(S::SIZE) else { + return None; + }; + + let Some(addr) = self.start_address().checked_add(offset) else { + return None; + }; + + Some(Self::containing_address(addr)) + } + + /// Subtract an offset from a frame. + /// + /// This function returns `None` on overflow or if the result does not have a + /// valid physical address. + #[inline] + #[const_fn(cfg(not(feature = "memory_encryption")))] + pub const fn checked_sub(self, rhs: u64) -> Option { + let Some(offset) = rhs.checked_mul(S::SIZE) else { + return None; + }; + + let Some(addr) = self.start_address().checked_sub(offset) else { + return None; + }; + + Some(Self::containing_address(addr)) + } + /// Returns a range of frames, exclusive `end`. #[inline] pub const fn range(start: PhysFrame, end: PhysFrame) -> PhysFrameRange { @@ -186,7 +222,8 @@ impl Add for PhysFrame { #[inline] #[track_caller] fn add(self, rhs: u64) -> Self::Output { - PhysFrame::containing_address(self.start_address() + rhs * S::SIZE) + self.checked_add(rhs) + .expect("attempt to add with overflow or resulted in invalid physical address") } } @@ -203,7 +240,8 @@ impl Sub for PhysFrame { #[inline] #[track_caller] fn sub(self, rhs: u64) -> Self::Output { - PhysFrame::containing_address(self.start_address() - rhs * S::SIZE) + self.checked_sub(rhs) + .expect("attempt to subtract with overflow or resulted in invalid physical address") } } diff --git a/src/structures/paging/mapper/mapped_page_table/iter.rs b/src/structures/paging/mapper/mapped_page_table/iter.rs index 31a701d6..552811b4 100644 --- a/src/structures/paging/mapper/mapped_page_table/iter.rs +++ b/src/structures/paging/mapper/mapped_page_table/iter.rs @@ -4,6 +4,8 @@ use core::ops::Add; +use const_fn::const_fn; + use super::{MappedPageTable, PageTableFrameMapping, PageTableWalkError, PageTableWalker}; use crate::structures::paging::{ Page, PageSize, PageTable, PageTableFlags, PageTableIndex, PhysFrame, Size1GiB, Size2MiB, @@ -23,15 +25,32 @@ pub struct MappedPage { pub flags: PageTableFlags, } +impl MappedPage { + #[const_fn(cfg(not(feature = "memory_encryption")))] + pub const fn checked_add(self, rhs: u64) -> Option { + let Some(page) = self.page.checked_add(rhs) else { + return None; + }; + + let Some(frame) = self.frame.checked_add(rhs) else { + return None; + }; + + Some(Self { + page, + frame, + flags: self.flags, + }) + } +} + impl Add for MappedPage { type Output = Self; + #[track_caller] fn add(self, rhs: u64) -> Self::Output { - Self { - page: self.page + rhs, - frame: self.frame + rhs, - flags: self.flags, - } + self.checked_add(rhs) + .expect("attempt to add with overflow or resulted in invalid mapped page") } } @@ -48,15 +67,42 @@ pub enum MappedPageItem { Size1GiB(MappedPage), } +impl MappedPageItem { + #[const_fn(cfg(not(feature = "memory_encryption")))] + pub const fn checked_add(self, rhs: u64) -> Option { + match self { + Self::Size4KiB(mapped_page) => { + let Some(page) = mapped_page.checked_add(rhs) else { + return None; + }; + + Some(Self::Size4KiB(page)) + } + Self::Size2MiB(mapped_page) => { + let Some(page) = mapped_page.checked_add(rhs) else { + return None; + }; + + Some(Self::Size2MiB(page)) + } + Self::Size1GiB(mapped_page) => { + let Some(page) = mapped_page.checked_add(rhs) else { + return None; + }; + + Some(Self::Size1GiB(page)) + } + } + } +} + impl Add for MappedPageItem { type Output = Self; + #[track_caller] fn add(self, rhs: u64) -> Self::Output { - match self { - Self::Size4KiB(mapped_page) => Self::Size4KiB(mapped_page + rhs), - Self::Size2MiB(mapped_page) => Self::Size2MiB(mapped_page + rhs), - Self::Size1GiB(mapped_page) => Self::Size1GiB(mapped_page + rhs), - } + self.checked_add(rhs) + .expect("attempt to add with overflow or resulted in invalid mapped page item") } } diff --git a/src/structures/paging/page.rs b/src/structures/paging/page.rs index e027caaf..67be7942 100644 --- a/src/structures/paging/page.rs +++ b/src/structures/paging/page.rs @@ -123,6 +123,40 @@ impl Page { S::SIZE } + /// Add an offset to a page. + /// + /// This function returns `None` on overflow or if the result does not have a + /// valid virtual address. + #[inline] + pub const fn checked_add(self, rhs: u64) -> Option { + let Some(offset) = rhs.checked_mul(S::SIZE) else { + return None; + }; + + let Some(addr) = self.start_address().checked_add(offset) else { + return None; + }; + + Some(Self::containing_address(addr)) + } + + /// Subtract an offset from a page. + /// + /// This function returns `None` on overflow or if the result does not have a + /// valid virtual address. + #[inline] + pub const fn checked_sub(self, rhs: u64) -> Option { + let Some(offset) = rhs.checked_mul(S::SIZE) else { + return None; + }; + + let Some(addr) = self.start_address().checked_sub(offset) else { + return None; + }; + + Some(Self::containing_address(addr)) + } + /// Returns the level 4 page table index of this page. #[inline] pub const fn p4_index(self) -> PageTableIndex { @@ -265,7 +299,8 @@ impl Add for Page { #[inline] #[track_caller] fn add(self, rhs: u64) -> Self::Output { - Page::containing_address(self.start_address() + rhs * S::SIZE) + self.checked_add(rhs) + .expect("attempt to add with overflow or resulted in non-canonical virtual address") } } @@ -282,7 +317,9 @@ impl Sub for Page { #[inline] #[track_caller] fn sub(self, rhs: u64) -> Self::Output { - Page::containing_address(self.start_address() - rhs * S::SIZE) + self.checked_sub(rhs).expect( + "attempt to subtract with overflow or resulted in non-canonical virtual address", + ) } } @@ -762,7 +799,7 @@ mod tests { } #[test] - #[should_panic = "attempt to add resulted in non-canonical virtual address: VirtAddrNotValid(0x800000000000)"] + #[should_panic = "attempt to add with overflow or resulted in non-canonical virtual address"] fn test_page_range_next_jumping_gap_panics() { let start = 0x7fff_ffff_f000; let end = 0xffff_8000_0000_0000; @@ -775,7 +812,7 @@ mod tests { // TODO: This probably shouldn't panic, but we can't fix this without a breaking change. #[test] - #[should_panic = "attempt to subtract resulted in non-canonical virtual address: VirtAddrNotValid(0xffff7ffffffff000)"] + #[should_panic = "attempt to subtract with overflow or resulted in non-canonical virtual address"] fn test_page_range_next_back_jumping_gap_panics() { let start = 0x7fff_ffff_f000; let end = 0xffff_8000_0000_0000; @@ -787,7 +824,7 @@ mod tests { } #[test] - #[should_panic = "attempt to add resulted in non-canonical virtual address: VirtAddrNotValid(0x800000000000)"] + #[should_panic = "attempt to add with overflow or resulted in non-canonical virtual address"] fn test_page_range_inclusive_next_not_jumping_gap_panics() { let start = 0x7fff_ffff_f000; let end = 0x7fff_ffff_f000; @@ -799,7 +836,7 @@ mod tests { } #[test] - #[should_panic = "attempt to subtract resulted in non-canonical virtual address: VirtAddrNotValid(0xffff7ffffffff000)"] + #[should_panic = "attempt to subtract with overflow or resulted in non-canonical virtual address"] fn test_page_range_inclusive_next_back_not_jumping_gap_panics() { let start = 0x7fff_ffff_f000; let end = 0xffff_8000_0000_0000; @@ -812,7 +849,7 @@ mod tests { // TODO: This probably shouldn't panic, but we can't fix this without a breaking change. #[test] - #[should_panic = "attempt to add resulted in non-canonical virtual address: VirtAddrNotValid(0x800000000000)"] + #[should_panic = "attempt to add with overflow or resulted in non-canonical virtual address"] fn test_page_range_inclusive_next_jumping_gap_panics() { let start = 0x7fff_ffff_f000; let end = 0x7fff_ffff_f000; @@ -825,7 +862,7 @@ mod tests { } #[test] - #[should_panic = "attempt to subtract resulted in non-canonical virtual address: VirtAddrNotValid(0xffff7ffffffff000)"] + #[should_panic = "attempt to subtract with overflow or resulted in non-canonical virtual address"] fn test_page_range_inclusive_next_back_jumping_gap_panics() { let start = 0xffff_8000_0000_0000; let end = 0xffff_8000_0000_0000;