Rust的Lookaside列表分配器

这个crate提供了一个基于Lookaside列表的Windows内核实验性Rust分配器。

由于Lookaside列表（固定大小的缓冲区）的性质，这个分配器不打算用作全局分配器，因此这个crate没有实现GlobalAlloc trait，另一方面，它实现了Allocator trait（不增长也不缩小），因此这是一个实验性/不稳定的crate。

用法

可以通过allocator_api或直接使用LookasideAlloc结构体来使用这个crate。

Rust驱动程序

如果在一个完全用Rust编写的驱动程序上工作，以下示例展示了我们可以如何通过分配器API使用这个crate。

#![no_std]
#![feature(allocator_api)]
#[macro_use] extern crate win_lookaside;

extern crate alloc;

use alloc::boxed::Box;
use win_lookaside::LookasideAlloc;
use windows_sys::Wdk::Foundation::NonPagedPool;

fn example() {
    // Init Lookaside List allocator with default values//!
    let mut allocator = LookasideAlloc::default();

    // Init Lookaside List with fixed-size to hold a u32
    // Properly handle possible InitError;
    allocator.init(core::mem::size_of::<u32>(), NonPagedPool as i32, None, None, None).unwrap();

    // Allocate from Lookaside & Free to it on Drop
    {
        let Ok(ctx) = Box::try_new_in(10, &allocator) else {
            return; // AllocError
        };
    }

    // Destroy Lookaside List Allocator
    allocator.destroy();
}

上面的示例只是为了示例，通常我们会在DriverEntry中初始化某个结构或全局变量并将Lookaside分配器存储在其中，并在DriverUnload中销毁它。

使用反编译器查看上述示例将生成以下简化的伪代码

void example()
{
  Status = -2;
  memset(&LookasideAlloc.unwrap, 0, sizeof(LookasideAlloc));
  while ( _InterlockedCompareExchange8(&LookasideAlloc.spin, 1, 0) )
  {
    while ( LookasideAlloc.spin )
      _mm_pause();
  }
  if ( ExInitializeLookasideListEx(&LookasideAlloc.lookaside, NULL, NULL, NonPagedPool, 0, 4, 'LLrs', 0) )
  {
    LookasideAlloc.spin = 0;
    return;
  }
  LookasideAlloc.init = 1;
  LookasideAlloc.spin = 0;
  while ( _InterlockedCompareExchange8(&LookasideAlloc.spin, 1, 0) )
  {
    while ( LookasideAlloc.spin )
      _mm_pause();
  }
  ctx = ExAllocateFromLookasideListEx(&LookasideAlloc.lookaside);
  LookasideAlloc.spin = 0;
  if ( !ctx )
    return;
  *ctx = 10;
  while ( _InterlockedCompareExchange8(&LookasideAlloc.spin, 1, 0) )
  {
      while ( LookasideAlloc.spin )
        _mm_pause();
  }
  ExFreeToLookasideListEx(&LookasideAlloc.lookaside, ctx);
  LookasideAlloc.spin = 0;
  while ( _InterlockedCompareExchange8(&LookasideAlloc.spin, 1, 0) )
  {
    while ( LookasideAlloc.spin )
      _mm_pause();
  }
  if ( LookasideAlloc.init )
    ExDeleteLookasideListEx(&LookasideAlloc.lookaside);
  LookasideAlloc.spin = 0;
}

C++驱动程序

另一个选项是如果我们正在使用一个用C++编写的驱动程序，并且我们想在Rust中工作于一个扩展/组件，我们可以在crate之上写一个薄的FFI层以公开其功能。

这个FFI层的一个非常简单的实现如下

#![no_std]
#![feature(allocator_api)]
#[macro_use] extern crate win_lookaside;

extern crate alloc;

use alloc::boxed::Box;
use windows_sys::Wdk::Foundation::PagedPool;
use windows_sys::Win32::Foundation::{NTSTATUS, STATUS_INSUFFICIENT_RESOURCES, STATUS_SUCCESS};
use win_lookaside::LookasideAlloc;

// Interior mutability due to the way the Lookaside API works
static mut LOOKASIDE: LookasideAlloc = LookasideAlloc::default();

struct Context{};

#[no_mangle]
pub unsafe extern "C" fn init_lookaside(tag: u32) -> NTSTATUS {
    LOOKASIDE.init(core::mem::size_of::<Context>(), PagedPool, Some(tag), None, None )?;
    STATUS_SUCCESS
}

#[no_mangle]
pub extern "C" fn create_context(context: *mut *mut Context) -> FfiResult<()> {
    let Ok(ctx) = unsafe { Box::try_new_in(Context {}, &LOOKASIDE) } else {
        return STATUS_INSUFFICIENT_RESOURCES;
    };

    unsafe {
        *context = Box::into_raw(ctx);
    }

    STATUS_SUCCESS
}

#[no_mangle]
pub extern "C" fn remove_context(context: *mut Context) {
    let _ctx = unsafe { Box::from_raw_in(context, &LOOKASIDE) };
}

#[no_mangle]
pub unsafe extern "C" fn free_lookaside() {
    LOOKASIDE.destroy();
}

这里Context只是一个空的struct，但它可以是更复杂的东西，可以提供更多功能，而C++驱动程序只需将其作为不可见指针存储。

然后我们可以以以下方式从我们的C++驱动程序中使用这个FFI层

// No error handling
#define LOOKASIDE_TAG 'aabb'

NTSTATUS
DriverEntry(
    _In_ PDRIVER_OBJECT     DriverObject,
    _In_ PUNICODE_STRING    RegistryPath
    )
{ 
    UNREFERENCED_PARAMETER(RegistryPath);

    PRUST_CTX context1, context2, context3, context4;

    DriverObject->DriverUnload = DriverUnload;

    init_lookaside(LOOKASIDE_TAG);

    create_context(&context1);
    create_context(&context2);
    create_context(&context3);
    create_context(&context4);

    remove_context(context1);
    remove_context(context3);
    remove_context(context4);
    remove_context(context2);
}

VOID
DriverUnload(
    _In_ PDRIVER_OBJECT DriverObject
    )
{
    PAGED_CODE();
    UNREFERENCED_PARAMETER(DriverObject);

    free_lookaside();
}

我们可以在调用free_lookaside之前使用WinDBG扩展!lookaside检查Lookaside列表的状态，以获取如下输出

4: kd> !lookaside fffff80474866030

Lookaside "" @ 0xfffff80474866030  Tag(hex): 0x61616262 "bbaa"
    Type           =       0000  NonPagedPool
    Current Depth  =          0  Max Depth  =          4
    Size           =          8  Max Alloc  =         32
    AllocateMisses =          4  FreeMisses =          0
    TotalAllocates =          4  TotalFrees =          4
    Hit Rate       =          0% Hit Rate   =        100%

备注

• 这个crate是用爱写的，但仍然是实验性的！在使用它之前请记住这一点 😄。
• 这个crate是在22H2 WDK下开发的，意味着某些Lookaside API方法被导出而不是内联。该crate尚未在比22H2更旧的WDK下进行测试。
• 虽然该crate已经在具有标准设置的驱动程序验证器下进行了测试，但尚未进行基准测试或性能测试。
• Allocator API的默认方法实现 grow、grow_zeroed 和 shrink 已被覆盖，如果使用这些方法将触发panic，因为这代表了对Lookaside API的误用。

待办事项

• 当可用时，使用来自 windows-sys 的 Lookaside API 绑定（查看 WDK 元数据）。
• 使用本机 MS 同步原语。
• 使用比22H2更旧的WDK测试该crate。