feat(driver): commit primer

1 files changed, 72 insertions, 0 deletions

diff --git a/crates/windows-kernel-rs/src/allocator.rs b/crates/windows-kernel-rs/src/allocator.rs
new file mode 100644
index 0000000..7b1efc8
--- /dev/null
+++ b/crates/windows-kernel-rs/src/allocator.rs
@@ -0,0 +1,72 @@
+//! This module provides an allocator to use with the [`alloc`] crate. You can
+//! define your own global allocator with the `#[global_allocator]` attribute
+//! when not using the `alloc` feature, in case you want to specify your own tag
+//! to use with [`ExAllocatePool2`] and [`ExAllocatePoolWithTag`].
+
+use core::alloc::{GlobalAlloc, Layout};
+
+use lazy_static::lazy_static;
+use windows_kernel_sys::{
+  base::_POOL_TYPE as POOL_TYPE,
+  ntoskrnl::{ExAllocatePool2, ExAllocatePoolWithTag, ExFreePool},
+};
+
+use crate::version::VersionInfo;
+
+/// See issue #52191.
+#[alloc_error_handler]
+fn alloc_error(_: Layout) -> ! { loop {} }
+
+lazy_static! {
+    /// The version of Microsoft Windows that is currently running. This is used by
+    /// [`KernelAllocator`] to determine whether to use [`ExAllocatePool2`] or
+    /// [`ExAllocatePoolWithTag`].
+    static ref VERSION_INFO: VersionInfo = {
+        VersionInfo::query().unwrap()
+    };
+}
+
+/// Represents a kernel allocator that relies on the `ExAllocatePool` family of
+/// functions to allocate and free memory for the `alloc` crate.
+pub struct KernelAllocator {
+  /// The 32-bit tag to use for the pool, this is usually derived from a
+  /// quadruplet of ASCII bytes, e.g. by invoking
+  /// `u32::from_ne_bytes(*b"rust")`.
+  tag: u32,
+}
+
+impl KernelAllocator {
+  /// Sets up a new kernel allocator with the 32-bit tag specified. The tag is
+  /// usually derived from a quadruplet of ASCII bytes, e.g. by invoking
+  /// `u32::from_ne_bytes(*b"rust")`.
+  pub const fn new(tag: u32) -> Self {
+    Self {
+      tag,
+    }
+  }
+}
+
+unsafe impl GlobalAlloc for KernelAllocator {
+  /// Uses [`ExAllocatePool2`] on Microsoft Windows 10.0.19041 and later, and
+  /// [`ExAllocatePoolWithTag`] on older versions of Microsoft Windows to
+  /// allocate memory.
+  unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+    let use_ex_allocate_pool2 = VERSION_INFO.major() > 10
+      || (VERSION_INFO.major() == 10 && VERSION_INFO.build_number() == 19041);
+
+    let ptr = if use_ex_allocate_pool2 {
+      ExAllocatePool2(POOL_TYPE::NonPagedPool as _, layout.size() as u64, self.tag)
+    } else {
+      ExAllocatePoolWithTag(POOL_TYPE::NonPagedPool, layout.size() as u64, self.tag)
+    };
+
+    if ptr.is_null() {
+      panic!("[kernel-alloc] failed to allocate pool.");
+    }
+
+    ptr as _
+  }
+
+  /// Uses [`ExFreePool`] to free allocated memory.
+  unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) { ExFreePool(ptr as _) }
+}