[PATCH v4] rust: Return Option from page_align and ensure no usize overflow

[PATCH v4] rust: Return Option from page_align and ensure no usize overflow

[Brendan Shephard]

Change `page_align()` to return `Option<usize>` to allow validation
of the provided `addr` value. This ensures that any value that is
within one `PAGE_SIZE` of `usize::MAX` will not panic, and instead
returns `None` to indicate overflow.

Signed-off-by: Brendan Shephard <bshephar@bne-home.net>
---
Changes in v2:
- Reworded commit message to follow the imperative form.
- Expanded the documentation to explain the `Some` and `None` return cases.
- Added a period at the end of the documentation comment.
- Link to v1 (and v2): https://lore.kernel.org/rust-for-linux/aSheTh-T1oroAUHR@fedora/T/#t

Changes in v3:
- Fix documentation layout for better rustdoc rendering
- Add doc examples and doctest
- Ensure function is always inlined for performance optimisation
- Restructure function so that early return is the None case and the
  default is the happy path.

Changes in v4:
- Fix rustdoc missing comment (//) prefix
- Rebase on master
- Link to v3: https://lore.kernel.org/rust-for-linux/aSoY31U3uDI2y7V1@fedora/T/#u

 rust/kernel/page.rs | 34 ++++++++++++++++++++++++++++------
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/rust/kernel/page.rs b/rust/kernel/page.rs
index 432fc0297d4a..2049ff859ac9 100644
--- a/rust/kernel/page.rs
+++ b/rust/kernel/page.rs
@@ -27,12 +27,34 @@

 /// Round up the given number to the next multiple of [`PAGE_SIZE`].
 ///
-/// It is incorrect to pass an address where the next multiple of [`PAGE_SIZE`] doesn't fit in a
-/// [`usize`].
-pub const fn page_align(addr: usize) -> usize {
-    // Parentheses around `PAGE_SIZE - 1` to avoid triggering overflow sanitizers in the wrong
-    // cases.
-    (addr + (PAGE_SIZE - 1)) & PAGE_MASK
+/// Returns a page aligned [`usize`] in cases where the value can be aligned. Otherwise, returns `None`
+/// if the aligned size will overflow a [`usize`].
+/// # Examples
+///
+/// Assuming a `PAGE_SIZE` of 4096 (0x1000):
+///
+/// ```rust
+/// use kernel::page::{page_align, PAGE_SIZE};
+/// // Case 1: Already aligned
+/// assert_eq!(page_align(0x0), Some(0x0));
+/// assert_eq!(page_align(0x1000), Some(0x1000));
+///
+/// // Case 2: Needs alignment up
+/// assert_eq!(page_align(0x1), Some(0x1000));
+/// assert_eq!(page_align(0x1001), Some(0x2000));
+///
+/// // Case 3: Requested address causes overflow (returns None)
+/// // The check asserts that None is returned when a value is requested within one PAGE_SIZE of
+/// // usize::MAX.
+/// let overflow_addr = usize::MAX - (PAGE_SIZE / 2);
+/// assert_eq!(page_align(overflow_addr), None);
+/// ```
+#[inline(always)]
+pub const fn page_align(addr: usize) -> Option<usize> {
+    let Some(sum) = addr.checked_add(PAGE_SIZE - 1) else {
+        return None;
+    };
+    Some(sum & PAGE_MASK)
 }

 /// Representation of a non-owning reference to a [`Page`].

base-commit: e6640487845061255af9614ec0a192e4fafa486e
--
2.51.1

Re: [PATCH v4] rust: Return Option from page_align and ensure no usize overflow

[Miguel Ojeda]

Hi Brendan,

This looks much better now with the examples!

On Sat, Nov 29, 2025 at 1:54 AM Brendan Shephard <bshephar@bne-home.net> wrote:
>
> +/// Returns a page aligned [`usize`] in cases where the value can be aligned. Otherwise, returns `None`
> +/// if the aligned size will overflow a [`usize`].

[`None`]

> +/// # Examples

Newline before `# Examples` header.

> +/// Assuming a `PAGE_SIZE` of 4096 (0x1000):

Can we assume that? i.e. these tests run as KUnit tests and we support
architectures that allow for other sizes, so we may need to guard
these with a `cfg` or `if` (which may be hidden), or perhaps better,
you could compute the values based on `PAGE_SIZE`, e.g. something
like:

    assert_eq!(page_align(PAGE_SIZE + 1), Some(2 * PAGE_SIZE));

It would be nice if you can confirm the tests fail (and then work
again) running them in a kernel config with a different page size:

    https://docs.kernel.org/rust/testing.html#the-kunit-tests

> +/// // The check asserts that None is returned when a value is requested within one PAGE_SIZE of
> +/// // usize::MAX.

`None`
`PAGE_SIZE`
`usize::MAX`

> +/// let overflow_addr = usize::MAX - (PAGE_SIZE / 2);
> +/// assert_eq!(page_align(overflow_addr), None);

Is there a reason for that particular address? If not, perhaps just
using `usize::MAX` (the maximum value) would be simpler. Or perhaps
the edge case (the first value that returns `None`) would be better.
Or even better, both.

Thanks!

Cheers,
Miguel

Re: [PATCH v4] rust: Return Option from page_align and ensure no usize overflow

[Brendan Shephard]

On Sun, Nov 30, 2025 at 01:01:58PM +0100, Miguel Ojeda wrote:
> Hi Brendan,
> 
> This looks much better now with the examples!
> 
> On Sat, Nov 29, 2025 at 1:54 AM Brendan Shephard <bshephar@bne-home.net> wrote:
> >
> > +/// Returns a page aligned [`usize`] in cases where the value can be aligned. Otherwise, returns `None`
> > +/// if the aligned size will overflow a [`usize`].
> 
> [`None`]
> 
> > +/// # Examples
> 
> Newline before `# Examples` header.
> 
> > +/// Assuming a `PAGE_SIZE` of 4096 (0x1000):
> 
> Can we assume that? i.e. these tests run as KUnit tests and we support
> architectures that allow for other sizes, so we may need to guard
> these with a `cfg` or `if` (which may be hidden), or perhaps better,
> you could compute the values based on `PAGE_SIZE`, e.g. something
> like:
> 
>     assert_eq!(page_align(PAGE_SIZE + 1), Some(2 * PAGE_SIZE));
> 
> It would be nice if you can confirm the tests fail (and then work
> again) running them in a kernel config with a different page size:
> 
>     https://docs.kernel.org/rust/testing.html#the-kunit-tests
> 
Yeah, good plan. I rewrote all of these tests to use the PAGE_SIZE
rather than static values, and I've tested this change on Asahi with 16k
PAGE_SIZE.

> > +/// // The check asserts that None is returned when a value is requested within one PAGE_SIZE of
> > +/// // usize::MAX.
> 
> `None`
> `PAGE_SIZE`
> `usize::MAX`
> 
> > +/// let overflow_addr = usize::MAX - (PAGE_SIZE / 2);
> > +/// assert_eq!(page_align(overflow_addr), None);
> 
> Is there a reason for that particular address? If not, perhaps just
> using `usize::MAX` (the maximum value) would be simpler. Or perhaps
> the edge case (the first value that returns `None`) would be better.
> Or even better, both.
> 
Yeah, I originally just used usize::MAX, but I wanted to demonstrate
that a value within 1 PAGE_SIZE of usize::MAX would overflow and result
in the None condition. The function would accept usize::MAX as a valid
addr value, so I thought doing it this way with a clearly named variable
showing the logic might be more representative of the logic.
> Thanks!
> 
> Cheers,
> Miguel

Thanks again. I'll fix up those docs issues and send a new patch with
the modified tests as well.