ntalloc/lib.rs
1//! Windows kernel-mode pool allocator for `#![no_std]` Rust drivers.
2//!
3//! `ntalloc` wraps the WDK pool APIs (`ExAllocatePool2`, `ExFreePool`,
4//! `ExFreePoolWithTag`) behind the standard [`GlobalAlloc`] trait and, with
5//! the `allocator-api` feature, the [`Allocator`] trait, so ordinary
6//! `alloc`-crate types (`Vec`, `Box`, ...) can be backed by kernel pool
7//! memory from a `#![no_std]` driver. It builds on stable Rust 1.85 or
8//! newer; `allocator-api` needs Rust 1.100, and `allocator-api-markers` a
9//! nightly compiler.
10//!
11//! [`Allocator`]: core::alloc::Allocator
12//!
13//! # Overview
14//!
15//! - [`NtAlloc`] is a zero-sized allocator type parameterized over the pool
16//! flags, the memory tag, whether deallocation uses the tagged or untagged
17//! free routine, and whether it realigns requests the pool cannot align by
18//! itself, which it does by default. The pool flags are exported as
19//! `POOL_FLAG_*` constants, such as [`POOL_FLAG_NON_PAGED`] and
20//! [`POOL_FLAG_UNINITIALIZED`].
21//! - [`PagedAlloc`] and [`NonPagedAlloc`] are convenience aliases for the
22//! paged and non-paged pool, respectively.
23//! - Allocations are tagged with a 4-byte ASCII tag (see
24//! [`DEFAULT_MEMORY_TAG`]) so they can be attributed back to this allocator
25//! with tools such as `!poolused`.
26//!
27//! # Design
28//!
29//! `ExAllocatePool2` only guarantees an allocation aligned to
30//! [`MIN_ALIGNMENT`] (the WDK's `MEMORY_ALLOCATION_ALIGNMENT`: two pointers,
31//! i.e. 16 bytes on 64-bit targets), or to the page size for allocations of
32//! at least one page (see [`NtAlloc::kernel_alignment`]). Rust [`Layout`]s can
33//! request stricter alignment than that, so when a request exceeds what the
34//! pool allocator natively guarantees, [`NtAlloc::allocate_aligned`]
35//! over-allocates by `layout.align()` extra bytes, rounds the returned
36//! pointer up to the requested alignment, and stashes the original
37//! (unaligned) pointer in the pointer-sized slot immediately before the
38//! aligned pointer. There is always room for that slot because
39//! `ExAllocatePool2` never returns a base below [`MIN_ALIGNMENT`], so
40//! rounding up moves the pointer by a multiple of [`MIN_ALIGNMENT`] — two
41//! pointers — and never by less. [`NtAlloc::deallocate_aligned`] reads that
42//! slot back out to recover the pointer that must be passed to
43//! `ExFreePool`/`ExFreePoolWithTag`, which is why an over-aligned allocation
44//! has to be released through it rather than through [`NtAlloc::deallocate`].
45//!
46//! Realignment is the default. A driver that never needs more than the pool's
47//! own alignment can turn it off with the `REALIGN` parameter, as described
48//! under [Turning realignment off](NtAlloc#turning-realignment-off).
49//!
50//! # Example
51//!
52//! ```no_run
53//! use ntalloc::NonPagedAlloc;
54//!
55//! type Alloc = NonPagedAlloc;
56//!
57//! let ptr = Alloc::allocate(64);
58//! assert!(!ptr.is_null());
59//! // SAFETY: `ptr` came from `Alloc::allocate` and needed no realignment.
60//! unsafe { Alloc::deallocate(ptr) };
61//! ```
62//!
63//! # Features
64//!
65//! - **`allocator-api`** — implements the `Allocator` trait for
66//! [`NtAlloc`], so a collection can use it directly, as in `Vec::new_in`
67//! and `Box::new_in`, not only as the `#[global_allocator]`. It needs Rust
68//! 1.100 or newer, where the trait is stable.
69//! - **`allocator-api-markers`** — implements the `StaticAllocator` and
70//! `AllocatorClone` marker traits for [`NtAlloc`]. `alloc` requires them in
71//! places: `Box::pin_in` needs `StaticAllocator`, and cloning an `Arc` or
72//! `Rc`, or using a `BTreeMap` at all, needs `AllocatorClone`. It implies
73//! `allocator-api`. The marker traits are still unstable (`allocator_ext`),
74//! so it needs a nightly compiler from October 2026 or later.
75//! - **`usermode-backend`** — for tests only. Allocates from the process heap
76//! (`std::alloc::System`) instead of the kernel pool, so code that allocates
77//! through [`NtAlloc`] can run in tests on a normal machine. Where the pool
78//! would bugcheck with `BAD_POOL_CALLER`, the test process prints what went
79//! wrong and aborts instead: on a double free, on a pointer it never handed
80//! out, and, with `TAGGED_DEALLOC`, on the wrong tag. On Windows any abort
81//! is reported as `STATUS_STACK_BUFFER_OVERRUN`, whatever its cause, and
82//! running the tests with `-- --test-threads=1` shows which test it was.
83//! It brings in `std`, along with the `hashbrown` and `allocator-api2`
84//! crates, so enable it only under `[dev-dependencies]`: through a normal
85//! dependency it reaches the driver itself, which then fails to build with
86//! a duplicate `panic_impl` lang item, as `std`'s panic handler clashes with
87//! the driver's own. [`NtAlloc::kernel_allocate`] always uses the pool,
88//! with or without it.
89//!
90//! ```toml
91//! [dev-dependencies]
92//! ntalloc = { version = "0.1", features = ["usermode-backend"] }
93//! ```
94
95#![no_std]
96#![cfg_attr(feature = "allocator-api-markers", feature(allocator_ext))]
97#![deny(unsafe_op_in_unsafe_fn)]
98#![deny(missing_docs)]
99// Doctests are this crate's documentation, so let them fail on unused imports
100// rather than accumulating stale `use` lines. Scoped to `unused` rather than
101// all warnings: this only ever runs for our own doctests, but a blanket deny
102// would also break on an unrelated future lint in a dependency's API.
103#![doc(test(attr(deny(unused))))]
104
105#[cfg(any(test, feature = "usermode-backend"))]
106extern crate std;
107
108mod bindings;
109#[cfg(any(test, feature = "usermode-backend"))]
110mod usermode;
111
112use crate::bindings::POOL_FLAG_RAISE_ON_FAILURE;
113pub use crate::bindings::{
114 POOL_FLAG_CACHE_ALIGNED, POOL_FLAG_NON_PAGED, POOL_FLAG_PAGED, POOL_FLAG_SPECIAL_POOL,
115 POOL_FLAG_UNINITIALIZED, POOL_FLAG_USE_QUOTA, POOL_FLAGS,
116};
117#[cfg(feature = "allocator-api")]
118use core::alloc::{AllocError, Allocator};
119#[cfg(feature = "allocator-api-markers")]
120use core::alloc::{AllocatorClone, StaticAllocator};
121#[cfg(feature = "allocator-api")]
122use core::ptr::{NonNull, slice_from_raw_parts_mut, without_provenance_mut};
123use core::{
124 alloc::{GlobalAlloc, Layout},
125 num::NonZeroUsize,
126 ptr::null_mut,
127};
128use windows_sys::Wdk::System::SystemServices::{
129 ExAllocatePool2, ExFreePool, ExFreePoolWithTag, PAGE_SIZE,
130};
131
132// The README's examples are compiled as doctests, all but the `Allocator` one,
133// which needs Rust 1.100 and the allocator-api feature and is marked ignore.
134#[cfg(doctest)]
135#[doc = include_str!("../README.md")]
136struct ReadmeDoctests;
137
138/// Default memory tag for allocations, visible via kernel debugger pool
139/// inspection tools such as `!poolused`/`!poolfind`.
140///
141/// Built with `from_ne_bytes` deliberately. Every Windows target is
142/// little-endian and the kernel reads a pool tag as a little-endian `u32`,
143/// so native byte order is what makes the debugger print `Rust` rather than
144/// `tsuR`. This is not a portability oversight; do not change it to
145/// `from_be_bytes`.
146pub const DEFAULT_MEMORY_TAG: u32 = u32::from_ne_bytes(*b"Rust");
147
148/// `"BIG "`, a tag the kernel reserves. Microsoft lists allocating with it as
149/// bug check `BAD_POOL_CALLER` (`0x9C`): such memory would be untrackable, and
150/// could corrupt the pool's tag tables.
151const BIG_POOL_TAG: u32 = u32::from_ne_bytes(*b"BIG ");
152
153/// The lowest byte `ExAllocatePool2` allows in a pool tag character: `0x20`,
154/// a space.
155const MIN_TAG_CHAR: u8 = 0x20;
156
157/// The highest byte `ExAllocatePool2` allows in a pool tag character: `0x7E`,
158/// a tilde.
159const MAX_TAG_CHAR: u8 = 0x7E;
160
161/// Whether `tag` is one to four printable ASCII characters, each from
162/// [`MIN_TAG_CHAR`] to [`MAX_TAG_CHAR`], in the order the debugger shows them,
163/// followed only by zero bytes when it is shorter, as `ExAllocatePool2`'s
164/// documentation requires.
165const fn tag_is_well_formed(tag: u32) -> bool {
166 let bytes = tag.to_ne_bytes();
167 let mut padding = false;
168 let mut i = 0;
169 while i < bytes.len() {
170 let byte = bytes[i];
171 if byte == 0 {
172 padding = true;
173 } else if padding || !matches!(byte, MIN_TAG_CHAR..=MAX_TAG_CHAR) {
174 return false;
175 }
176 i += 1;
177 }
178 // A first byte of zero means no characters at all.
179 bytes[0] != 0
180}
181
182/// Whether any byte of `tag` is an ASCII letter or digit. Microsoft lists
183/// allocating with a tag that has neither as bug check `BAD_POOL_CALLER`
184/// (`0x9D`), since such tags make pool usage hard to track.
185const fn tag_has_letter_or_digit(tag: u32) -> bool {
186 let bytes = tag.to_ne_bytes();
187 let mut i = 0;
188 while i < bytes.len() {
189 if bytes[i].is_ascii_alphanumeric() {
190 return true;
191 }
192 i += 1;
193 }
194 false
195}
196
197/// A stateless, `#![no_std]`-friendly Windows kernel pool allocator.
198///
199/// `NtAlloc` is a zero-sized type: all of its behavior is selected entirely at
200/// compile time via its const generic parameters, and every operation is
201/// exposed as an associated function, so no instance ever needs to be
202/// constructed to use them directly. An instance is still needed to satisfy
203/// the [`GlobalAlloc`] trait bound (and, with the `allocator-api` feature,
204/// the [`Allocator`] one), and can always be created with
205/// [`Default::default()`] at zero runtime cost.
206///
207/// [`Allocator`]: core::alloc::Allocator
208///
209/// # Type Parameters
210/// - `ALLOC_POOL_FLAGS`: the pool flags passed to `ExAllocatePool2`.
211/// [`POOL_FLAG_PAGED`] or [`POOL_FLAG_NON_PAGED`] selects the pool, and
212/// [`POOL_FLAG_UNINITIALIZED`], [`POOL_FLAG_CACHE_ALIGNED`],
213/// [`POOL_FLAG_USE_QUOTA`] or [`POOL_FLAG_SPECIAL_POOL`] can be added to
214/// it, as in the example below. Flags from `wdm.h` that this crate does not
215/// export can be passed as raw values. `POOL_FLAG_RAISE_ON_FAILURE` is
216/// refused; see [below](#unsupported-configurations).
217/// - `MEMORY_TAG`: the 4-byte pool tag attached to every allocation, used by
218/// tools such as `!poolused`/`!poolfind` to attribute pool memory back to
219/// this allocator. Defaults to [`DEFAULT_MEMORY_TAG`]. It must be one to
220/// four printable ASCII characters (`0x20` to `0x7E`), padded with zero
221/// bytes at the end when shorter, with at least one letter or digit; see
222/// [below](#unsupported-configurations).
223/// - `TAGGED_DEALLOC`: selects the free routine used by [`NtAlloc::deallocate`]:
224/// `true` uses `ExFreePoolWithTag` (the tag must match the one the memory
225/// was allocated with), `false` uses the untagged `ExFreePool`. Defaults to
226/// `true`.
227/// - `REALIGN`: whether [`NtAlloc::allocate_aligned`], and through it the
228/// [`GlobalAlloc`] and [`Allocator`] implementations, serve a [`Layout`]
229/// aligned more strictly than the pool guarantees for its size (see
230/// [`NtAlloc::kernel_alignment`]) by over-allocating and realigning, as the
231/// [crate documentation](crate#design) describes. Defaults to `true`;
232/// `false` turns realignment off, as described
233/// [below](#turning-realignment-off).
234///
235/// Prefer the [`PagedAlloc`]/[`NonPagedAlloc`] aliases over naming `NtAlloc`
236/// directly, unless the configuration needs more flags than the pool:
237///
238/// ```
239/// use ntalloc::{NtAlloc, POOL_FLAG_NON_PAGED, POOL_FLAG_UNINITIALIZED};
240///
241/// // Non-paged pool that skips the zeroing ExAllocatePool2 otherwise does.
242/// type ScratchAlloc = NtAlloc<{ POOL_FLAG_NON_PAGED | POOL_FLAG_UNINITIALIZED }>;
243///
244/// assert!(!ScratchAlloc::POOL_ZEROES);
245/// ```
246///
247/// Call associated functions through an alias that names the whole type, such
248/// as `type Alloc = NonPagedAlloc;` followed by `Alloc::allocate(64)`. Writing
249/// `NonPagedAlloc::allocate(64)` does not compile (E0284): Rust fills in the
250/// aliases' default parameters only where a type is expected.
251///
252/// # Turning realignment off
253///
254/// Realignment is on by default. A driver that never needs more than the
255/// pool's own alignment can turn it off by setting `REALIGN` to `false`:
256/// [`NtAlloc::allocate_aligned`], and with it the [`GlobalAlloc`] and
257/// [`Allocator`] implementations, then pass every request to the pool by size
258/// and ignore the alignment it asks for. Nothing is over-allocated or stored
259/// next to a block, and a block is aligned only as the pool aligns it: to
260/// [`MIN_ALIGNMENT`] below a page, and to a page from one page up.
261///
262/// Only do so if nothing allocated through such a configuration needs more,
263/// including types from other crates, which a `#[global_allocator]` serves
264/// too. Anything `#[repr(align(32))]` or larger needs more, as do 256- and
265/// 512-bit SIMD vectors and cache-padded wrappers such as crossbeam's
266/// `CachePadded`. A misaligned `Box` or reference is undefined behaviour, and
267/// an aligned SIMD load from one faults, which in a driver is a bug check.
268/// Debug builds usually stop at the first misaligned access, with "misaligned
269/// pointer dereference"; release builds do not.
270///
271/// ```no_run
272/// use core::alloc::{GlobalAlloc, Layout};
273/// use ntalloc::{DEFAULT_MEMORY_TAG, MIN_ALIGNMENT, NonPagedAlloc};
274///
275/// type NoRealignAlloc = NonPagedAlloc<DEFAULT_MEMORY_TAG, true, false>;
276///
277/// let alloc = NoRealignAlloc::default();
278///
279/// // Asks for 64-byte alignment, but gets only what the pool gives a small block.
280/// let layout = Layout::from_size_align(64, 64).unwrap();
281/// // SAFETY: the layout has a non-zero size.
282/// let ptr = unsafe { alloc.alloc(layout) };
283/// assert!(!ptr.is_null());
284/// assert_eq!(ptr as usize % MIN_ALIGNMENT, 0);
285///
286/// // SAFETY: from `alloc` with this layout, freed once.
287/// unsafe { alloc.dealloc(ptr, layout) };
288/// ```
289///
290/// # Unsupported configurations
291///
292/// These configurations break `ExAllocatePool2`'s documented rules, fail every
293/// allocation, are listed by Microsoft as bug checks, or fail in a way Rust
294/// cannot handle, so a type that uses one fails to build as soon as anything
295/// allocates with it. The error comes from `cargo build`; `cargo check` does
296/// not evaluate it.
297///
298/// - `POOL_FLAG_RAISE_ON_FAILURE` (`0x20`) makes `ExAllocatePool2` raise an
299/// SEH exception on failure instead of returning null. Rust cannot catch
300/// SEH: left unhandled the exception bugchecks the system, and handled
301/// further up it unwinds through Rust frames, which is undefined behaviour.
302/// - A zero `MEMORY_TAG`, which `ExAllocatePool2` rejects as invalid.
303/// - A `MEMORY_TAG` that is not one to four printable ASCII characters
304/// (`0x20` to `0x7E`) followed only by zero padding, as `ExAllocatePool2`'s
305/// documentation requires: one with a control byte, a byte above `0x7E`, or
306/// a zero byte before its last character.
307/// - A `MEMORY_TAG` with no ASCII letter or digit, such as four spaces, which
308/// Microsoft lists as bug check `BAD_POOL_CALLER` (`0x9D`).
309/// - `"BIG "`, a tag the kernel reserves, which Microsoft lists as bug check
310/// `BAD_POOL_CALLER` (`0x9C`).
311///
312/// Without `POOL_FLAG_RAISE_ON_FAILURE`, failure is reported as null.
313/// Fallible APIs such as `Box::try_new_in` return it as an error, and
314/// infallible ones such as `Box::new` pass it to `handle_alloc_error`, which
315/// in a `#![no_std]` build panics by default: the fail-fast behaviour the flag
316/// would otherwise give.
317///
318/// The flags are not validated beyond that: any other combination the pool
319/// considers invalid fails at run time, like any other allocation failure.
320///
321/// ```compile_fail,E0080
322/// use ntalloc::{NtAlloc, POOL_FLAG_NON_PAGED};
323///
324/// // 0x20 is POOL_FLAG_RAISE_ON_FAILURE, which this crate does not export.
325/// type Raising = NtAlloc<{ POOL_FLAG_NON_PAGED | 0x20 }>;
326///
327/// let _ = Raising::allocate(64);
328/// ```
329///
330/// ```compile_fail,E0080
331/// use ntalloc::{NtAlloc, POOL_FLAG_NON_PAGED};
332///
333/// type Untagged = NtAlloc<POOL_FLAG_NON_PAGED, 0>;
334///
335/// let _ = Untagged::allocate(64);
336/// ```
337///
338/// ```compile_fail,E0080
339/// use ntalloc::{NtAlloc, POOL_FLAG_NON_PAGED};
340///
341/// // A zero byte before the last character.
342/// type Gapped = NtAlloc<POOL_FLAG_NON_PAGED, { u32::from_ne_bytes(*b"ab\0c") }>;
343///
344/// let _ = Gapped::allocate(64);
345/// ```
346///
347/// ```compile_fail,E0080
348/// use ntalloc::{NtAlloc, POOL_FLAG_NON_PAGED};
349///
350/// type Blank = NtAlloc<POOL_FLAG_NON_PAGED, { u32::from_ne_bytes(*b" ") }>;
351///
352/// let _ = Blank::allocate(64);
353/// ```
354///
355/// ```compile_fail,E0080
356/// use ntalloc::{NtAlloc, POOL_FLAG_NON_PAGED};
357///
358/// type Big = NtAlloc<POOL_FLAG_NON_PAGED, { u32::from_ne_bytes(*b"BIG ") }>;
359///
360/// let _ = Big::allocate(64);
361/// ```
362#[derive(Default, Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
363pub struct NtAlloc<
364 const ALLOC_POOL_FLAGS: POOL_FLAGS,
365 const MEMORY_TAG: u32 = DEFAULT_MEMORY_TAG,
366 const TAGGED_DEALLOC: bool = true,
367 const REALIGN: bool = true,
368>;
369
370/// [`NtAlloc`] backed by the paged pool ([`POOL_FLAG_PAGED`]).
371///
372/// Memory from the paged pool may be paged out to disk, so it may only be
373/// touched at IRQL `APC_LEVEL` or below — that is, strictly below
374/// `DISPATCH_LEVEL`. Never touch it from a DPC, from an ISR, or while
375/// holding a spin lock.
376pub type PagedAlloc<
377 const MEMORY_TAG: u32 = DEFAULT_MEMORY_TAG,
378 const TAGGED_DEALLOC: bool = true,
379 const REALIGN: bool = true,
380> = NtAlloc<POOL_FLAG_PAGED, MEMORY_TAG, TAGGED_DEALLOC, REALIGN>;
381/// [`NtAlloc`] backed by the non-paged pool ([`POOL_FLAG_NON_PAGED`]).
382///
383/// Non-paged pool memory is always resident and safe to access at any IRQL,
384/// at the cost of being a more constrained system-wide resource than paged
385/// pool.
386pub type NonPagedAlloc<
387 const MEMORY_TAG: u32 = DEFAULT_MEMORY_TAG,
388 const TAGGED_DEALLOC: bool = true,
389 const REALIGN: bool = true,
390> = NtAlloc<POOL_FLAG_NON_PAGED, MEMORY_TAG, TAGGED_DEALLOC, REALIGN>;
391
392/// Size, in bytes, of a pointer on the target platform.
393pub const POINTER_SIZE: usize = size_of::<*const u8>();
394/// The minimum alignment `ExAllocatePool2` guarantees for allocations smaller
395/// than a page. Mirrors the WDK's `MEMORY_ALLOCATION_ALIGNMENT` (two
396/// pointers: 16 bytes on 64-bit targets, 8 bytes on 32-bit).
397pub const MIN_ALIGNMENT: usize = POINTER_SIZE * 2;
398
399impl<
400 const ALLOC_POOL_FLAGS: POOL_FLAGS,
401 const MEMORY_TAG: u32,
402 const TAGGED_DEALLOC: bool,
403 const REALIGN: bool,
404> NtAlloc<ALLOC_POOL_FLAGS, MEMORY_TAG, TAGGED_DEALLOC, REALIGN>
405{
406 /// Whether allocations from this configuration arrive zero-initialized.
407 ///
408 /// `ExAllocatePool2` zeroes every block it returns unless
409 /// [`POOL_FLAG_UNINITIALIZED`] is among `ALLOC_POOL_FLAGS`, so this is a
410 /// property of the type, resolved at compile time. It is `true` for
411 /// [`PagedAlloc`] and [`NonPagedAlloc`].
412 ///
413 /// [`GlobalAlloc::alloc_zeroed`], and with the `allocator-api` feature
414 /// [`Allocator::allocate_zeroed`], consult it to skip the redundant second
415 /// pass their default implementations would otherwise make over memory the
416 /// pool has already zeroed.
417 ///
418 /// [`Allocator::allocate_zeroed`]: core::alloc::Allocator::allocate_zeroed
419 pub const POOL_ZEROES: bool = (ALLOC_POOL_FLAGS & POOL_FLAG_UNINITIALIZED) == 0;
420
421 /// Fails the build for a configuration [`NtAlloc`] refuses: one whose
422 /// `ALLOC_POOL_FLAGS` include `POOL_FLAG_RAISE_ON_FAILURE`, or whose
423 /// `MEMORY_TAG` is zero, is not one to four printable ASCII characters,
424 /// has no letter or digit, or is `"BIG "`.
425 ///
426 /// With `POOL_FLAG_RAISE_ON_FAILURE`, `ExAllocatePool2` raises an SEH
427 /// exception on failure instead of returning null. Unhandled, the
428 /// exception bugchecks the system; handled further up, it unwinds through
429 /// Rust frames that call `ExAllocatePool2` through a non-unwinding ABI,
430 /// which is undefined behaviour. Either way the null that [`GlobalAlloc`]
431 /// and [`Allocator`] report failure with never arrives. A zero tag fails
432 /// every allocation, since `ExAllocatePool2` rejects it as invalid. The
433 /// other tags break the format `ExAllocatePool2` documents, or are ones
434 /// Microsoft lists as bug check `BAD_POOL_CALLER`.
435 ///
436 /// [`Allocator`]: core::alloc::Allocator
437 ///
438 /// Both allocation backends evaluate this before anything else, so the
439 /// error appears wherever such a type allocates. A constant that depends
440 /// on the type's parameters is only evaluated when code is generated for
441 /// it, which is why `cargo build` reports it and `cargo check` does not.
442 const CONFIGURATION_SUPPORTED: () = {
443 assert!(
444 (ALLOC_POOL_FLAGS & POOL_FLAG_RAISE_ON_FAILURE) == 0,
445 "POOL_FLAG_RAISE_ON_FAILURE is not supported: ExAllocatePool2 would \
446 raise an SEH exception on failure, which Rust cannot catch"
447 );
448 assert!(
449 MEMORY_TAG != 0,
450 "MEMORY_TAG is zero: ExAllocatePool2 fails every allocation with a \
451 zero tag"
452 );
453 assert!(
454 tag_is_well_formed(MEMORY_TAG),
455 "MEMORY_TAG is not one to four printable ASCII characters (0x20 to \
456 0x7E) followed only by zero padding, as ExAllocatePool2 requires"
457 );
458 assert!(
459 tag_has_letter_or_digit(MEMORY_TAG),
460 "MEMORY_TAG has no ASCII letter or digit (bug check \
461 BAD_POOL_CALLER, 0x9D)"
462 );
463 assert!(
464 MEMORY_TAG != BIG_POOL_TAG,
465 "MEMORY_TAG is \"BIG \", a tag the kernel reserves (bug check \
466 BAD_POOL_CALLER, 0x9C)"
467 );
468 };
469
470 /// Returns the address, one pointer-width before `ptr`, where the
471 /// original (pre-alignment) pointer is stored by
472 /// [`NtAlloc::save_original_pointer`].
473 ///
474 /// `ptr` must be a pointer returned by [`NtAlloc::allocate_aligned`] for
475 /// a layout that required realignment; the slot this names is only
476 /// guaranteed to exist, and to belong to the caller, for such a pointer.
477 /// That restriction is why this is not part of the public API: nothing
478 /// in the signature can enforce it.
479 #[inline]
480 #[expect(
481 clippy::cast_ptr_alignment,
482 reason = "a realigned address is aligned to more than MIN_ALIGNMENT"
483 )]
484 fn original_pointer_location(ptr: *mut u8) -> *mut *mut u8 {
485 // SAFETY: `ptr` is an address allocate_aligned realigned, which lies at
486 // least MIN_ALIGNMENT bytes into the same pool block, so one pointer
487 // back is still inside it.
488 unsafe { ptr.cast::<*mut u8>().sub(1) }
489 }
490
491 /// Stores `original_ptr` in the pointer-sized slot immediately before
492 /// `ptr`, so it can later be recovered by [`NtAlloc::original_pointer`]
493 /// and passed to the real free routine.
494 ///
495 /// `ptr` must have at least [`POINTER_SIZE`] bytes of valid, writable
496 /// memory immediately before it. [`NtAlloc::allocate_aligned`] guarantees
497 /// that: the gap [`NtAlloc::align_up_to`] opens is a multiple of
498 /// [`MIN_ALIGNMENT`], which is two pointers wide. `original_ptr` is the
499 /// pointer the underlying allocation actually returned, to be restored
500 /// on deallocation.
501 #[inline]
502 fn save_original_pointer(ptr: *mut u8, original_ptr: *mut u8) {
503 let original_ptr_location = Self::original_pointer_location(ptr);
504 // SAFETY: the slot lies in the gap allocate_aligned opened before `ptr`,
505 // inside the pool block and outside the region handed out, and is
506 // pointer-aligned because `ptr` is aligned to more than MIN_ALIGNMENT.
507 unsafe {
508 *original_ptr_location = original_ptr;
509 }
510 }
511
512 /// Reads back the pointer previously stored by
513 /// [`NtAlloc::save_original_pointer`].
514 ///
515 /// `ptr` must be a pointer returned by [`NtAlloc::allocate_aligned`] for
516 /// which realignment was required; for any other pointer the slot read
517 /// here holds unrelated memory.
518 #[inline]
519 fn original_pointer(ptr: *mut u8) -> *mut u8 {
520 let original_ptr_location = Self::original_pointer_location(ptr);
521 // SAFETY: this is the slot save_original_pointer wrote when
522 // allocate_aligned handed `ptr` out; nothing else writes to it.
523 unsafe { *original_ptr_location }
524 }
525
526 /// Returns the alignment guarantee `ExAllocatePool2` provides for an
527 /// allocation of `size` bytes.
528 ///
529 /// Per the [`ExAllocatePool2` documentation][remarks], allocations of at
530 /// least one page are page-aligned; smaller allocations are only
531 /// guaranteed [`MIN_ALIGNMENT`].
532 ///
533 /// [remarks]: https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wdm/nf-wdm-exallocatepool2#remarks
534 ///
535 /// # Parameters
536 /// - `size`: the requested allocation size, in bytes.
537 ///
538 /// # Returns
539 /// `PAGE_SIZE` if `size` is at least one page, otherwise [`MIN_ALIGNMENT`].
540 ///
541 /// # Examples
542 /// ```
543 /// use ntalloc::{PagedAlloc, MIN_ALIGNMENT};
544 /// use windows_sys::Wdk::System::SystemServices::PAGE_SIZE;
545 ///
546 /// type Alloc = PagedAlloc;
547 ///
548 /// assert_eq!(Alloc::kernel_alignment(16), MIN_ALIGNMENT);
549 /// assert_eq!(Alloc::kernel_alignment(PAGE_SIZE as usize), PAGE_SIZE as usize);
550 /// ```
551 #[must_use]
552 #[inline]
553 pub fn kernel_alignment(size: usize) -> usize {
554 if size >= PAGE_SIZE as usize {
555 PAGE_SIZE as usize
556 } else {
557 MIN_ALIGNMENT
558 }
559 }
560
561 /// Allocates `size` bytes directly from the kernel pool via
562 /// `ExAllocatePool2`, bypassing the process-heap backend that
563 /// [`NtAlloc::allocate`] uses in tests and under the `usermode-backend`
564 /// feature. It always reaches the pool, so outside a kernel it can only
565 /// be called with a `size` of `0`.
566 ///
567 /// # Parameters
568 /// - `size`: the number of bytes to allocate. A zero-size request never
569 /// reaches the pool: `ExAllocatePool2` documents `NumberOfBytes` as
570 /// "a non-zero number of bytes", and lists only memory exhaustion and
571 /// a bad tag or flags as reasons it returns null — so it defines no
572 /// behaviour at all for `0`. This returns null instead.
573 ///
574 /// # Returns
575 /// A pointer to the allocated memory, or a null pointer if `size` is `0`
576 /// or the allocation fails.
577 ///
578 /// # Examples
579 /// ```no_run
580 /// use ntalloc::NonPagedAlloc;
581 ///
582 /// type Alloc = NonPagedAlloc;
583 ///
584 /// let ptr = Alloc::kernel_allocate(64);
585 /// assert!(!ptr.is_null());
586 /// ```
587 #[must_use = "discarding the pointer leaks the allocation"]
588 #[inline]
589 pub fn kernel_allocate(size: usize) -> *mut u8 {
590 // Fails the build, not anything at runtime, for a configuration this
591 // crate refuses; see CONFIGURATION_SUPPORTED.
592 let () = Self::CONFIGURATION_SUPPORTED;
593
594 if size == 0 {
595 return null_mut();
596 }
597
598 // SAFETY: ExAllocatePool2 takes no pointers: only the flags and tag,
599 // which CONFIGURATION_SUPPORTED has checked, and a non-zero size. The
600 // paged pool's IRQL limit is documented on PagedAlloc.
601 unsafe { ExAllocatePool2(ALLOC_POOL_FLAGS, size, MEMORY_TAG) }.cast()
602 }
603
604 /// Allocates `size` bytes, returning a null pointer for a zero-size
605 /// request.
606 ///
607 /// With the `usermode-backend` feature, and in this crate's own tests,
608 /// allocation is redirected to the process heap so that the allocator's
609 /// logic can run without a kernel; otherwise this calls
610 /// [`NtAlloc::kernel_allocate`].
611 ///
612 /// The zero-size check lives in each backend rather than here, so the
613 /// shipped path performs it once instead of twice.
614 ///
615 /// # Parameters
616 /// - `size`: the number of bytes to allocate.
617 ///
618 /// # Returns
619 /// A pointer to the allocated memory, or a null pointer if `size` is `0`
620 /// or the allocation fails.
621 ///
622 /// # Examples
623 /// ```no_run
624 /// use ntalloc::NonPagedAlloc;
625 ///
626 /// type Alloc = NonPagedAlloc;
627 ///
628 /// let ptr = Alloc::allocate(64);
629 /// assert!(!ptr.is_null());
630 /// // SAFETY: `ptr` came from `Alloc::allocate` and needed no realignment.
631 /// unsafe { Alloc::deallocate(ptr) };
632 /// ```
633 #[must_use = "discarding the pointer leaks the allocation"]
634 pub fn allocate(size: usize) -> *mut u8 {
635 #[cfg(any(test, feature = "usermode-backend"))]
636 {
637 Self::usermode_allocate(size)
638 }
639
640 #[cfg(not(any(test, feature = "usermode-backend")))]
641 {
642 Self::kernel_allocate(size)
643 }
644 }
645
646 /// Returns whether `layout`'s alignment is already satisfied by whatever
647 /// alignment the kernel pool allocator natively guarantees for
648 /// `layout.size()` (see [`NtAlloc::kernel_alignment`]), meaning no
649 /// over-allocation/realignment is needed to satisfy it.
650 ///
651 /// This runs twice per allocation: once in [`NtAlloc::allocate_aligned`] to
652 /// decide whether to stash the pool's own pointer, and again in
653 /// [`NtAlloc::deallocate_aligned`] to decide whether to read it back. The
654 /// two answers must agree — if the free path decides "not stashed" for a
655 /// pointer that was, it hands `ExFreePool` an address inside the block and
656 /// corrupts the pool.
657 ///
658 /// Agreement is not automatic, because this depends on `layout.size()` and
659 /// [`Allocator::deallocate`] accepts any size that *fits* rather than the
660 /// exact one. Two separate properties currently guarantee it; both are
661 /// checked by `free_path_decision_cannot_flip`.
662 ///
663 /// [`Allocator::deallocate`]: core::alloc::Allocator::deallocate
664 #[inline]
665 fn is_kernel_aligned(layout: Layout) -> bool {
666 let pool_alignment = Self::kernel_alignment(layout.size());
667
668 layout.align() <= pool_alignment
669 }
670
671 /// Rounds `allocated_buffer` up to the next address that is a multiple
672 /// of `alignment`. The result is always strictly greater than
673 /// `allocated_buffer`, even when it was already aligned, which is what
674 /// leaves [`NtAlloc::allocate_aligned`] room to stash a pointer just
675 /// before the address it hands out.
676 ///
677 /// How much room is not guaranteed by this function alone. For an
678 /// arbitrary address the gap can be a single byte, too small to hold a
679 /// pointer. It is sufficient in [`NtAlloc::allocate_aligned`] only
680 /// because the address there comes from `ExAllocatePool2` and is
681 /// therefore aligned to at least [`MIN_ALIGNMENT`], while `alignment`
682 /// is a larger power of two — so the gap is a multiple of
683 /// [`MIN_ALIGNMENT`] and never smaller than it.
684 ///
685 /// # Parameters
686 /// - `allocated_buffer`: the pointer to align. This function does not
687 /// dereference the pointer, so any address is acceptable.
688 /// - `alignment`: the desired alignment, in bytes. [`NonZeroUsize`]
689 /// rather than `usize` because `alignment - 1` below would otherwise
690 /// underflow on `0`, wrapping to a null result in release builds. It
691 /// must also be a power of two — the masking trick silently returns a
692 /// non-multiple otherwise — which the type cannot express, so a
693 /// `debug_assert!` covers that half. [`Layout`] guarantees both
694 /// properties, and every caller's alignment comes from a [`Layout`].
695 ///
696 /// # Returns
697 /// The next address after `allocated_buffer` that is a multiple of
698 /// `alignment`, as a pointer still derived from `allocated_buffer` and
699 /// therefore still carrying its provenance.
700 ///
701 /// # Note
702 /// Does not check for address overflow; with overflow checks disabled
703 /// (the default in release builds), an `allocated_buffer` near the top
704 /// of the address space combined with a large `alignment` would wrap
705 /// silently — worth noting since kernel-mode pointers often sit near
706 /// that boundary.
707 #[inline]
708 fn align_up_to(allocated_buffer: *mut u8, alignment: NonZeroUsize) -> *mut u8 {
709 debug_assert!(
710 alignment.is_power_of_two(),
711 "alignment must be a power of two; the mask below is meaningless otherwise"
712 );
713
714 let alignment = alignment.get();
715 let alignment_mask_base = !(alignment - 1);
716
717 // `map_addr` rather than a cast through `usize`: it produces the same
718 // address but keeps the result derived from `allocated_buffer`, so the
719 // pointer still carries provenance over the pool block it points into.
720 // Casting out to an integer and back drops that lineage, leaving the
721 // compiler to guess which allocation was meant — defined under
722 // exposed provenance, but the one construct Miri cannot check.
723 allocated_buffer.map_addr(|addr| (addr & alignment_mask_base) + alignment)
724 }
725
726 /// Allocates memory for `layout`, handling alignment requirements stronger
727 /// than what the kernel pool allocator natively provides.
728 ///
729 /// If `layout.align()` is no stronger than [`NtAlloc::kernel_alignment`]
730 /// for the requested size, this simply forwards to [`NtAlloc::allocate`].
731 /// Otherwise it over-allocates by `layout.align()` extra bytes, rounds the
732 /// result up to `layout.align()`, and stashes the original pointer just
733 /// before the returned one (see the crate-level documentation for
734 /// details). With `REALIGN` set to `false` it forwards every request to
735 /// [`NtAlloc::allocate`] instead: see
736 /// [Turning realignment off](NtAlloc#turning-realignment-off).
737 ///
738 /// # Parameters
739 /// - `layout`: the size and alignment to allocate for.
740 ///
741 /// # Returns
742 /// A pointer to `layout.size()` bytes aligned to `layout.align()` (only to
743 /// the pool's own alignment when `REALIGN` is `false`), or a null pointer
744 /// on allocation failure.
745 ///
746 /// A zero-size `layout` always returns a null pointer, and never reaches
747 /// the pool. [`GlobalAlloc`] calls a zero-size request undefined
748 /// behaviour but invites implementations to "guarantee a sentinel address
749 /// or a null pointer in response to a zero-size allocation request" —
750 /// this is that guarantee. Callers needing a usable pointer for an empty
751 /// allocation want the [`Allocator`] impl (the `allocator-api` feature),
752 /// which returns a dangling one as that trait requires.
753 ///
754 /// # Note
755 /// A pointer returned from this function when realignment was required
756 /// cannot be freed with the plain [`NtAlloc::deallocate`] — that function
757 /// has no way to recover the original pointer. Free it with
758 /// [`NtAlloc::deallocate_aligned`], passing the same `layout`, or through
759 /// this type's [`GlobalAlloc`]/[`Allocator`] implementations, which do
760 /// the same thing. (For a pointer that did *not* need realignment,
761 /// [`NtAlloc::deallocate`] is fine either way — but since only `layout`
762 /// distinguishes the two cases, always pairing with
763 /// [`NtAlloc::deallocate_aligned`] is the rule that cannot go wrong.)
764 ///
765 /// [`Allocator`]: core::alloc::Allocator
766 ///
767 /// # Examples
768 /// ```no_run
769 /// use core::alloc::Layout;
770 /// use ntalloc::NonPagedAlloc;
771 ///
772 /// type Alloc = NonPagedAlloc;
773 ///
774 /// let layout = Layout::from_size_align(16, 64).unwrap();
775 /// let ptr = Alloc::allocate_aligned(layout);
776 /// assert!(!ptr.is_null());
777 /// assert_eq!(ptr as usize % layout.align(), 0);
778 ///
779 /// // SAFETY: `ptr` came from `allocate_aligned` with this exact layout.
780 /// unsafe { Alloc::deallocate_aligned(ptr, layout) };
781 /// ```
782 #[must_use = "discarding the pointer leaks the allocation"]
783 pub fn allocate_aligned(layout: Layout) -> *mut u8 {
784 if layout.size() == 0 {
785 return null_mut();
786 }
787
788 if !REALIGN || Self::is_kernel_aligned(layout) {
789 // Without REALIGN the alignment asked for is ignored, as NtAlloc's
790 // documentation warns: the block is aligned as the pool aligns it.
791 Self::allocate(layout.size())
792 } else {
793 // Cannot overflow: Layout guarantees size rounded up to align is
794 // at most isize::MAX, so size + align is at most isize::MAX + 1.
795 let extended_buffer_size = layout.size() + layout.align();
796 let buffer = Self::allocate(extended_buffer_size);
797 if buffer.is_null() {
798 return null_mut();
799 }
800
801 // SAFETY: Layout guarantees a non-zero power-of-two alignment.
802 let alignment = unsafe { NonZeroUsize::new_unchecked(layout.align()) };
803 let aligned = Self::align_up_to(buffer, alignment);
804
805 debug_assert!(
806 (aligned as usize).saturating_add(layout.size())
807 <= (buffer as usize).saturating_add(extended_buffer_size),
808 "the realigned region must fit inside the block actually allocated"
809 );
810 debug_assert!(
811 (aligned as usize) - (buffer as usize) >= POINTER_SIZE,
812 "the gap must be wide enough to hold the stashed pointer"
813 );
814
815 Self::save_original_pointer(aligned, buffer);
816
817 aligned
818 }
819 }
820
821 /// Implements the [`Allocator::allocate`] contract: allocates memory for
822 /// `layout` via [`NtAlloc::allocate_aligned`] and wraps it as a
823 /// `NonNull<[u8]>` slice.
824 ///
825 /// # Returns
826 /// - `Ok(_)` — a `NonNull<[u8]>` slice of `layout.size()` bytes, aligned as
827 /// [`NtAlloc::allocate_aligned`] describes, or, for a zero-size
828 /// `layout`, a dangling pointer aligned to `layout.align()` that never
829 /// reached the pool. The [`Allocator::allocate`] implementation
830 /// documents that case in full.
831 /// - `Err(AllocError)` — if the underlying pool allocation fails.
832 #[cfg(feature = "allocator-api")]
833 #[clippy::msrv = "1.100"]
834 fn allocate_slice(layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
835 if layout.size() == 0 {
836 // Allocator, unlike GlobalAlloc, requires zero-sized allocations
837 // to succeed, and ExAllocatePool2 has no defined behaviour for a
838 // zero-size request — so this must not reach the pool. A dangling
839 // but correctly aligned pointer is exactly what the contract asks
840 // for, and deallocate_aligned discards it without freeing.
841 let dangling = without_provenance_mut::<u8>(layout.align());
842 let empty = slice_from_raw_parts_mut(dangling, 0);
843
844 // SAFETY: layout.align() is a non-zero power of two, so the
845 // address is non-null.
846 return Ok(unsafe { NonNull::new_unchecked(empty) });
847 }
848
849 let allocated_raw = Self::allocate_aligned(layout);
850
851 if allocated_raw.is_null() {
852 Err(AllocError)
853 } else {
854 // Exactly layout.size(), never the larger block really taken from
855 // the pool. Allocator permits reporting more, but whatever is
856 // reported becomes `max` in the range of sizes that
857 // Allocator::deallocate will then accept — and the free path
858 // recomputes the stash decision from the size it is handed. Keeping
859 // max == min leaves callers one legal size and no way to move that
860 // decision. This line is load-bearing; see is_kernel_aligned.
861 let allocated_slice = slice_from_raw_parts_mut(allocated_raw, layout.size());
862
863 // SAFETY: allocated_raw was checked non-null above, and the slice
864 // pointer keeps its address.
865 Ok(unsafe { NonNull::new_unchecked(allocated_slice) })
866 }
867 }
868
869 /// Frees `ptr` with `ExFreePoolWithTag`, using this allocator's
870 /// `MEMORY_TAG`.
871 ///
872 /// `ptr` must be non-null (unlike C's `free`, passing null here is not a
873 /// safe no-op) and must have been allocated from the same pool with the
874 /// same `MEMORY_TAG`; passing a pointer allocated with a different tag
875 /// corrupts the pool.
876 // The heap backend frees through usermode_free instead; these stay
877 // compiled, and checked, in every configuration.
878 #[cfg_attr(any(test, feature = "usermode-backend"), allow(dead_code))]
879 fn deallocate_tagged(ptr: *mut u8) {
880 // SAFETY: the only caller, deallocate, passes on its own contract: a
881 // non-null pool block of this configuration, allocated with this
882 // MEMORY_TAG and not yet freed.
883 unsafe {
884 ExFreePoolWithTag(ptr.cast(), MEMORY_TAG);
885 }
886 }
887
888 /// Frees `ptr` with the untagged `ExFreePool`.
889 ///
890 /// `ptr` must be non-null (unlike C's `free`, passing null here is not a
891 /// safe no-op) and must have been allocated from the same pool as this
892 /// allocator targets.
893 #[cfg_attr(any(test, feature = "usermode-backend"), allow(dead_code))]
894 fn deallocate_non_tagged(ptr: *mut u8) {
895 // SAFETY: the only caller, deallocate, passes on its own contract: a
896 // non-null pool block of this configuration, not yet freed.
897 unsafe {
898 ExFreePool(ptr.cast());
899 }
900 }
901
902 /// Frees a pointer previously returned by [`NtAlloc::allocate`]. Normally
903 /// that is a pool block, which [`NtAlloc::kernel_allocate`] returns too,
904 /// freed with the routine `TAGGED_DEALLOC` selects. Under the
905 /// `usermode-backend` feature both functions use the process heap instead.
906 ///
907 /// This does **not** account for the pointer-stashing done by
908 /// [`NtAlloc::allocate_aligned`] for over-aligned requests — see that
909 /// function's documentation for when it is safe to call this directly.
910 ///
911 /// # Safety
912 /// The caller must guarantee all of the following:
913 ///
914 /// - `ptr` is non-null. Unlike [`NtAlloc::deallocate_aligned`], this
915 /// function does not check, and `ExFreePool`/`ExFreePoolWithTag`
916 /// bugchecks on a null pointer rather than ignoring it the way C's
917 /// `free` would.
918 /// - `ptr` came from the same `NtAlloc` configuration: the same pool
919 /// flags, the same `MEMORY_TAG`, and the same `TAGGED_DEALLOC`.
920 /// Freeing with a mismatched tag corrupts the pool.
921 /// - `ptr` is the address the pool returned, not one that
922 /// [`NtAlloc::allocate_aligned`] adjusted for over-alignment. Those go
923 /// to [`NtAlloc::deallocate_aligned`] instead.
924 /// - `ptr` has not already been freed, and is not used afterwards.
925 ///
926 /// # Examples
927 /// ```no_run
928 /// use ntalloc::NonPagedAlloc;
929 ///
930 /// type Alloc = NonPagedAlloc;
931 ///
932 /// let ptr = Alloc::allocate(64);
933 /// assert!(!ptr.is_null());
934 /// // SAFETY: `ptr` came from `Alloc::allocate` and needed no realignment.
935 /// unsafe { Alloc::deallocate(ptr) };
936 /// ```
937 pub unsafe fn deallocate(ptr: *mut u8) {
938 #[cfg(any(test, feature = "usermode-backend"))]
939 {
940 Self::usermode_free(ptr);
941 }
942
943 #[cfg(not(any(test, feature = "usermode-backend")))]
944 {
945 if TAGGED_DEALLOC {
946 Self::deallocate_tagged(ptr);
947 } else {
948 Self::deallocate_non_tagged(ptr);
949 }
950 }
951 }
952
953 /// Frees a pointer returned by [`NtAlloc::allocate_aligned`] for `layout`,
954 /// recovering and freeing the original (pre-realignment) pointer when
955 /// `layout`'s alignment required over-allocation.
956 ///
957 /// This is the counterpart to [`NtAlloc::allocate_aligned`]: a pointer
958 /// from that function must be freed through this one, because for an
959 /// over-aligned request the address handed out is not the address the
960 /// pool returned, and only this function knows how to recover it.
961 ///
962 /// Passing a null `ptr` is a no-op.
963 ///
964 /// # Safety
965 /// The caller must guarantee all of the following:
966 ///
967 /// - `ptr` was returned by [`NtAlloc::allocate_aligned`] on this same
968 /// `NtAlloc` instantiation — the same pool flags, the same
969 /// `MEMORY_TAG`, the same `TAGGED_DEALLOC` and the same `REALIGN` — or
970 /// is null.
971 /// - `layout` is the *exact* `Layout` that was passed to that call.
972 /// `layout` alone decides whether the original pointer is read back
973 /// from the slot before `ptr` or `ptr` is freed directly, so a
974 /// mismatched layout will either hand `ExFreePool` an interior
975 /// pointer or read [`POINTER_SIZE`] bytes of unrelated memory and
976 /// free whatever they happen to contain.
977 /// - `ptr` has not already been freed, and is not used afterwards.
978 ///
979 /// None of these can be checked at runtime: an interior pointer and a
980 /// valid one are indistinguishable, and the pool allocator has no way to
981 /// report that a tag or a size did not match.
982 pub unsafe fn deallocate_aligned(ptr: *mut u8, layout: Layout) {
983 if ptr.is_null() {
984 return;
985 }
986
987 if layout.size() == 0 {
988 // Nothing was ever taken from the pool for a zero-size layout:
989 // allocate_aligned returns null, and allocate_slice returns a
990 // dangling pointer. Either way there is nothing to give back.
991 return;
992 }
993
994 // Without REALIGN nothing is ever stashed: allocate_aligned hands out the
995 // pool's own pointer for every layout.
996 let original_ptr = if !REALIGN || Self::is_kernel_aligned(layout) {
997 ptr
998 } else {
999 Self::original_pointer(ptr)
1000 };
1001
1002 debug_assert!(
1003 !original_ptr.is_null(),
1004 "Original pointer is null. This may indicate that the pointer was not allocated by this allocator or has already been deallocated."
1005 );
1006
1007 // SAFETY: original_ptr is the address `allocate` returned — either
1008 // `ptr` itself when no realignment happened, or the value stashed by
1009 // allocate_aligned — and this fn's own contract covers the rest.
1010 unsafe { Self::deallocate(original_ptr) };
1011 }
1012}
1013
1014// SAFETY: [`GlobalAlloc`] asks for blocks that fit the layout and stay valid
1015// until freed, and for no method to unwind. For [`NtAlloc`]:
1016//
1017// - `alloc` returns null or `layout.size()` bytes aligned to `layout.align()`:
1018// the pool's own block when its alignment suffices, otherwise an address
1019// realigned inside a block over-allocated by `layout.align()`. With `REALIGN`
1020// set to `false` it returns the pool's own block whatever the layout asks
1021// for, so that promise holds only for layouts within the pool's alignment;
1022// NtAlloc documents that such a configuration must not be asked for more;
1023// - a block belongs to the pool, not to the allocator value, so it stays valid
1024// until `dealloc` hands it back, recovering the pool's own address from the
1025// same layout;
1026// - `alloc_zeroed` returns zeroed memory: the pool zeroes it unless
1027// [`POOL_FLAG_UNINITIALIZED`] is set, and then it writes the zeroes itself;
1028// - in a release build nothing on the pool path panics; its `debug_assert!`s
1029// fire only on a bug here or on a caller that broke the contract, and the
1030// heap backend aborts rather than unwinds on misuse.
1031unsafe impl<
1032 const ALLOC_POOL_FLAGS: POOL_FLAGS,
1033 const MEMORY_TAG: u32,
1034 const TAGGED_DEALLOC: bool,
1035 const REALIGN: bool,
1036> GlobalAlloc for NtAlloc<ALLOC_POOL_FLAGS, MEMORY_TAG, TAGGED_DEALLOC, REALIGN>
1037{
1038 /// Allocates memory for `layout` via [`NtAlloc::allocate_aligned`].
1039 ///
1040 /// # Safety
1041 /// As [`GlobalAlloc::alloc`] requires: `layout` must have non-zero size,
1042 /// and the caller must pass the exact same `layout` to
1043 /// [`GlobalAlloc::dealloc`] (or `realloc`) when freeing the returned
1044 /// pointer.
1045 ///
1046 /// Should a zero-size `layout` arrive anyway, this returns null rather
1047 /// than reaching `ExAllocatePool2`, which documents no behaviour for a
1048 /// zero-size request. See [`NtAlloc::allocate_aligned`].
1049 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
1050 Self::allocate_aligned(layout)
1051 }
1052
1053 /// Frees a pointer previously returned by [`GlobalAlloc::alloc`] (or
1054 /// `realloc`) on this same allocator, correctly recovering any pointer
1055 /// adjustment [`NtAlloc::allocate_aligned`] made for over-aligned
1056 /// requests.
1057 ///
1058 /// # Safety
1059 /// `ptr` must have been returned by a prior call to [`GlobalAlloc::alloc`]
1060 /// (or `realloc`) on this same `NtAlloc` instantiation, and `layout` must
1061 /// be the exact layout used for that allocation.
1062 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
1063 // SAFETY: forwarded straight from GlobalAlloc::dealloc, whose own
1064 // contract already requires exactly what deallocate_aligned does.
1065 unsafe { Self::deallocate_aligned(ptr, layout) };
1066 }
1067
1068 /// Allocates zero-initialized memory for `layout`.
1069 ///
1070 /// The default implementation of this method allocates and then writes
1071 /// zeroes over the whole block. `ExAllocatePool2` has already done that,
1072 /// so when [`NtAlloc::POOL_ZEROES`] holds this simply returns the
1073 /// allocation and the second pass is skipped — the condition is a const,
1074 /// so only one branch survives compilation.
1075 ///
1076 /// For a realigned request the returned region still arrives zeroed: it
1077 /// lies inside the larger block the pool zeroed, and the pointer
1078 /// [`NtAlloc::allocate_aligned`] stashes goes in the bytes *before* it.
1079 ///
1080 /// # Safety
1081 /// As [`GlobalAlloc::alloc`]: `layout` must have non-zero size, and the
1082 /// same `layout` must come back on the free path.
1083 unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
1084 let ptr = Self::allocate_aligned(layout);
1085
1086 // The left operand is a const, so when the pool zeroes this whole
1087 // branch — the null test included — is compiled away. When it does
1088 // not, null still has to be excluded: it means the allocation failed
1089 // or `layout` was zero-sized (see [`NtAlloc::allocate_aligned`]), and
1090 // null is not a valid write target even for a zero-length write.
1091 if !Self::POOL_ZEROES && !ptr.is_null() {
1092 // SAFETY: the allocation succeeded, so layout.size() bytes from
1093 // `ptr` are valid for writes.
1094 unsafe { ptr.write_bytes(0, layout.size()) };
1095 }
1096
1097 ptr
1098 }
1099}
1100
1101// SAFETY: [`Allocator`] asks that a block stay valid until it is deallocated,
1102// however the allocator value is copied, moved or dropped, and that a copy
1103// behave like the original. For [`NtAlloc`]:
1104//
1105// - blocks come from the pool (or the heap backend), never from the value,
1106// which is a zero-sized `Copy` type that owns nothing, so no copy, move or
1107// drop of it can invalidate one;
1108// - every value of one type makes the same pool calls, so any of them can free
1109// what another allocated;
1110// - `allocate` reports exactly `layout.size()` bytes, so the layout the free
1111// path receives is the one the stash decision was made with;
1112// - with `REALIGN` set to `false`, `allocate` returns the pool's own block
1113// whatever `layout.align()` asks for, so the block fits the layout only when
1114// the pool's alignment covers it, which NtAlloc documents as the condition
1115// for using such a configuration;
1116// - a zero-size request gets a dangling, aligned pointer, which `deallocate`
1117// recognizes and never frees.
1118#[cfg(feature = "allocator-api")]
1119#[clippy::msrv = "1.100"]
1120unsafe impl<
1121 const ALLOC_POOL_FLAGS: POOL_FLAGS,
1122 const MEMORY_TAG: u32,
1123 const TAGGED_DEALLOC: bool,
1124 const REALIGN: bool,
1125> Allocator for NtAlloc<ALLOC_POOL_FLAGS, MEMORY_TAG, TAGGED_DEALLOC, REALIGN>
1126{
1127 /// Allocates memory for `layout` via [`NtAlloc::allocate_aligned`],
1128 /// returning it as a `NonNull<[u8]>` slice.
1129 ///
1130 /// # Returns
1131 /// - `Ok(_)` — a `NonNull<[u8]>` slice of exactly `layout.size()` bytes
1132 /// (never more, even though the real underlying allocation may be
1133 /// larger than that once realignment padding is accounted for).
1134 /// - `Ok(_)` — for a zero-size `layout`, a dangling pointer aligned to
1135 /// `layout.align()` with a length of `0`. [`Allocator`] requires
1136 /// zero-sized allocations to succeed, unlike [`GlobalAlloc`], and the
1137 /// pool is never involved. Passing it back to
1138 /// [`Allocator::deallocate`] is correct and does nothing.
1139 /// - `Err(AllocError)` — only if the underlying pool allocation fails.
1140 fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
1141 Self::allocate_slice(layout)
1142 }
1143
1144 /// Allocates zero-initialized memory for `layout`.
1145 ///
1146 /// Skips the second zeroing pass the default implementation would make
1147 /// whenever [`NtAlloc::POOL_ZEROES`] holds, for the reasons given on
1148 /// [`GlobalAlloc::alloc_zeroed`]. A zero-sized `layout` yields the same
1149 /// dangling, length-`0` slice [`Allocator::allocate`] returns, which is
1150 /// trivially zeroed.
1151 fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
1152 let allocated = Self::allocate_slice(layout)?;
1153
1154 if !Self::POOL_ZEROES {
1155 // SAFETY: allocate_slice returned a block of exactly this length.
1156 unsafe {
1157 allocated
1158 .cast::<u8>()
1159 .as_ptr()
1160 .write_bytes(0, allocated.len());
1161 }
1162 }
1163
1164 Ok(allocated)
1165 }
1166
1167 /// Frees a pointer previously returned by [`Allocator::allocate`] (or
1168 /// `grow`/`shrink`) on this same allocator.
1169 ///
1170 /// # Safety
1171 /// `ptr` must denote memory currently allocated via this allocator, and
1172 /// `layout` must be the same layout that was used to allocate it, per
1173 /// the [`Allocator::deallocate`] contract.
1174 unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
1175 // SAFETY: forwarded straight from Allocator::deallocate, whose own
1176 // contract already requires exactly what deallocate_aligned does.
1177 unsafe { Self::deallocate_aligned(ptr.as_ptr(), layout) };
1178 }
1179}
1180
1181// SAFETY: [`StaticAllocator`] asks that nothing about the allocator's own
1182// lifetime can invalidate memory it handed out. For [`NtAlloc`]:
1183//
1184// - it has no `Drop` impl, so dropping one releases nothing;
1185// - it is a zero-sized type holding no state;
1186// - its lifetime expiring frees nothing, because it owns nothing;
1187// - nothing reachable from safe code can invalidate an allocation. The safe
1188// public surface is [`NtAlloc::POOL_ZEROES`], [`NtAlloc::kernel_alignment`],
1189// [`NtAlloc::kernel_allocate`], [`NtAlloc::allocate`] and
1190// [`NtAlloc::allocate_aligned`], none of which free anything. Of the five
1191// functions that do return memory, two are `unsafe`
1192// ([`NtAlloc::deallocate`] and [`NtAlloc::deallocate_aligned`]) and the other
1193// three are private and called only from inside those two;
1194// - all of the above hold for every equivalent allocator, since instances
1195// sharing the same const parameters are indistinguishable.
1196//
1197// A pool block therefore stays valid until the matching `ExFreePool` or
1198// `ExFreePoolWithTag`, regardless of what happens to the allocator value.
1199#[cfg(feature = "allocator-api-markers")]
1200unsafe impl<
1201 const ALLOC_POOL_FLAGS: POOL_FLAGS,
1202 const MEMORY_TAG: u32,
1203 const TAGGED_DEALLOC: bool,
1204 const REALIGN: bool,
1205> StaticAllocator for NtAlloc<ALLOC_POOL_FLAGS, MEMORY_TAG, TAGGED_DEALLOC, REALIGN>
1206{
1207}
1208
1209// SAFETY: [`AllocatorClone`] asks that a clone be interchangeable with the
1210// original, so that either can free what the other allocated. [`NtAlloc`] is a
1211// zero-sized type carrying no state: the pool, the tag and the free routine all
1212// come from its const parameters, so two values of the same type forward to the
1213// same pool calls identically. Moving or dropping one frees nothing, so blocks
1214// stay valid while clones live, and the derived `Clone` copies a zero-sized
1215// value and cannot unwind. The type is also `Copy`, which preserves that
1216// equivalence exactly as cloning does.
1217#[cfg(feature = "allocator-api-markers")]
1218unsafe impl<
1219 const ALLOC_POOL_FLAGS: POOL_FLAGS,
1220 const MEMORY_TAG: u32,
1221 const TAGGED_DEALLOC: bool,
1222 const REALIGN: bool,
1223> AllocatorClone for NtAlloc<ALLOC_POOL_FLAGS, MEMORY_TAG, TAGGED_DEALLOC, REALIGN>
1224{
1225}
1226
1227#[cfg(test)]
1228mod tests;