Implement aligned WDK allocator #353

Zero-Tang · 2025-05-13T16:06:52Z

Now the WDK allocator can allocate aligned memory according to the layout argument.

Note that this allocator won't waste additional space for alignment if the alignment is normal (e.g.: less than or equal to 16-byte boundary on 64-bit platform). In average scenario, this is better than _aligned_malloc and _aligned_free in msvcrt, in that I guess we rarely need big alignments in drivers for some stuff in the pool. Still, some performance-critical codes might require cache-aligned (e.g.: 64-byte in x86) objects, which is usually bigger than ExAllocatePool2's alignment.

Close #303.

Now the WDK allocator can allocate aligned memory according to the layout argument.

wmmc88

Thanks for the PR @Zero-Tang!

How did you go about validating these changes? I think some concrete examples of the non-aligned case should be added the example drivers (ie. sample wdm & sample kmdf)

crates/wdk-alloc/src/lib.rs

wmmc88 · 2025-05-13T17:59:13Z

crates/wdk-alloc/src/lib.rs

+                            let q = ((p as usize & mask) + layout.align()) as *mut *mut u8;
+                            // Store the original pointer right before the pointer to return,
+                            // so that ExFreePoolWithTag can receive the correct pointer.
+                            // This is also how `_aligned_malloc` is implemented in msvcrt.dll


Where is this documented?

It is not documented, but _aligned_malloc is actually somewhat open-source. Check the %ProgramFiles%\Windows Kits\10\Source\<Your SDK version>\ucrt\heap\align.cpp file, and you can find its implementation.

crates/wdk-alloc/src/lib.rs

wmmc88 · 2025-05-13T18:07:00Z

related to #303? @gurry could you also take a look at this and provide feedback :)

crates/wdk-alloc/src/lib.rs

Zero-Tang · 2025-05-14T02:07:43Z

Thanks for the feedback!
I did wrote an example, but in a local separate crate, because rust-analyzer marks most parts of wdk-alloc crate as inactive codes for some strange reasons. It won't be marked as inactive if I check it in a separate crate.
I will copy'n'paste my validation example here soon.

gurry · 2025-05-14T07:52:59Z

related to #303? @gurry could you also take a look at this and provide feedback :)

Thanks @wmmc88. I will have a look.

Also address Melvin's reviews.

gurry

@Zero-Tang The overall logic looks good to me. Thanks. Added comments for some small improvements.

crates/wdk-alloc/src/lib.rs

gurry · 2025-05-17T06:51:23Z

examples/sample-wdm-driver/src/lib.rs

+    {
+        // Example of using WDK allocator.
+        // Allocations will be properly aligned on their boundaries!
+        // Allocate a single instance.
+        let ah=Box::new(AlignedHigher(1234));
+        println!("ah is allocated at {:p}, value={ah:?}",&raw const *ah);
+        // Verify its alignment.
+        assert_eq!((&raw const *ah) as usize & (align_of::<AlignedHigher>() - 1), 0);
+        // Allocate a vector that occupies more than a page.
+        let vh=vec![AlignedHigher(1234), AlignedHigher(5678), AlignedHigher(9012), AlignedHigher(3456), AlignedHigher(7890)];
+        println!("vh is allocated at {:p}, value={vh:?}",vh.as_ptr());
+        // Verify their alignments.
+        for x in &vh {
+            assert_eq!((core::ptr::from_ref::<AlignedHigher>(x) as usize) & (align_of::<AlignedHigher>() - 1), 0);
+        }
+    }


@wmmc88 Is there a more preferred way to write tests?

How about we write unit tests in the crates/wdk-alloc/src/lib.rs file itself. We can mock APIs like ExAllocatePool2 by putting them in global statics and then then in the test replacing them with mock implementations.

I think we can mock ExAllocatePool2 by calling VirtualAlloc on Windows or mmap on Linux, then add the returned pointer by 16 or 8. In this case, we can guarantee no coincidental big alignments returned by CRT's malloc.

It may be easier for the mock ExAllocatePool2 to use a global array as its source of allocations instead of VirtualAlloc or malloc.

It would be preferable if there are unit-tests with mocked wdk api calls in the wdk-alloc.

Ideally, integration tests for the allocator would be facilitated in wdk-alloc/examples folder, but i don't think either our cargo-make scripts, nor the new cargo-wdk tool support packaging drivers from examples right now

It would be preferable if there are unit-tests with mocked wdk api calls in the wdk-alloc.

In that case @Zero-Tang can you please see if you can write some key unit tests?

Emmmmmm, I'm not so sure how to do this.

An approach like this can work:

Declare a static pointing either to the real ExAllocatePool2 or a mock depending on whether we're running tests or not:

type ExAllocatePool2Fn = unsafe extern "system" fn(pool_flags: u32, number_of_bytes: usize, tag: u32) -> *mut c_void; #[cfg(not(test))] pub static mut EX_ALLOCATE_POOL2_FN: ExAllocatePool2Fn = ExAllocatePool2; #[cfg(test)] pub static mut EX_ALLOCATE_POOL2_FN: ExAllocatePool2Fn = mock_ex_allocate_pool2;

Declare a static buffer that mock_ex_allocate_pool2 will "allocate" from:

static ALLOC_BUFFER: [u8; 5120] = [0; 5120]; // 5 Pages worth of buffer to allocate from static mut FREE_INDEX: usize = 0; // Index to keep track of available memory in buffer

Implement mock_ex_allocate_pool2 using this buffer as the source of allocations.

In WdkAllcator call EX_ALLOCATE_POOL2_FN instead of calling ExAllocatePoo2 directly.

Now you should be ready to write your tests.

I kinda give up on this in that I couldn't make the test link to wdk_alloc::kernel_mode module due to the #[cfg(any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"))] attribute.
I tried cargo test --config driver_model__driver_type=\"WDM\" but it doesn't work.

@wmmc88 Can you please help with this?

I kinda give up on this in that I couldn't make the test link to wdk_alloc::kernel_mode module due to the #[cfg(any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"))] attribute.
I tried cargo test --config driver_model__driver_type="WDM" but it doesn't work.

I don't think this will work as currently structured. I am working on some tooling to enable better unit testing of code that relies on the metadata section. unfortunately, since this is such a high impact change, I think we'll wait for that tooling to finish before merging this (with tests).

crates/wdk-alloc/src/lib.rs

Address Gurinder's review comments.

gurry

Thanks for making the requested changes @Zero-Tang. Here are a few more suggestions.

crates/wdk-alloc/src/lib.rs

examples/sample-wdm-driver/src/lib.rs

gurry

A few comments on comments.

crates/wdk-alloc/src/lib.rs

Further clarify how the original pointer is stored and why it's safe.

Zero-Tang · 2025-05-20T16:11:06Z

Just FYI that I edited the first comment of this PR thread so that this PR will close #303 on merge.

gurry · 2025-05-21T01:50:06Z

Just FYI that I edited the first comment of this PR thread so that this PR will close #303 on merge.

Thanks. Yes, #303 should be closed.

I am done with my review. Will wait for inputs from @wmmc88.

Implement aligned allocator

083fb4a

Now the WDK allocator can allocate aligned memory according to the layout argument.

wmmc88 requested changes May 13, 2025

View reviewed changes

wmmc88 reviewed May 13, 2025

View reviewed changes

crates/wdk-alloc/src/lib.rs Outdated Show resolved Hide resolved

Add samples to verify the alignments

071895b

Also address Melvin's reviews.

Zero-Tang requested a review from wmmc88 May 15, 2025 04:07

gurry reviewed May 17, 2025

View reviewed changes

crates/wdk-alloc/src/lib.rs Show resolved Hide resolved

Zero-Tang added 2 commits May 18, 2025 12:35

Minor improvements

1ac9d3f

Address Gurinder's review comments.

Add MIN_POOL_ALIGNMENT constant

fddcf82

gurry reviewed May 19, 2025

View reviewed changes

Bunch of nits

d06890d

gurry reviewed May 20, 2025

View reviewed changes

crates/wdk-alloc/src/lib.rs Outdated Show resolved Hide resolved

crates/wdk-alloc/src/lib.rs Outdated Show resolved Hide resolved

crates/wdk-alloc/src/lib.rs Outdated Show resolved Hide resolved

crates/wdk-alloc/src/lib.rs Outdated Show resolved Hide resolved

Some nits for comments

8fa375b

Further clarify how the original pointer is stored and why it's safe.

wmmc88 self-assigned this May 21, 2025

Implement mock and test (incomplete)

e3183e0

Implement aligned WDK allocator #353

Are you sure you want to change the base?

Implement aligned WDK allocator #353

Conversation

Zero-Tang commented May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wmmc88 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

wmmc88 commented May 13, 2025

Uh oh!

Uh oh!

Zero-Tang commented May 14, 2025

Uh oh!

gurry commented May 14, 2025

Uh oh!

gurry left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gurry May 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gurry May 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gurry May 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gurry left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gurry left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Zero-Tang commented May 20, 2025

Uh oh!

gurry commented May 21, 2025

Uh oh!

Zero-Tang commented May 13, 2025 •

edited

Loading

gurry May 17, 2025 •

edited

Loading

gurry May 18, 2025 •

edited

Loading

gurry May 26, 2025 •

edited

Loading