Skip to content

Commit

Permalink
Use stack-based formatter.
Browse files Browse the repository at this point in the history
This will enable print!/println! macros to be used in <=DIRQL environments.
  • Loading branch information
Zero-Tang committed Oct 8, 2024
1 parent 45a2299 commit 0402d08
Showing 1 changed file with 52 additions and 15 deletions.
67 changes: 52 additions & 15 deletions crates/wdk/src/print.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,40 @@
// Copyright (c) Microsoft Corporation
// License: MIT OR Apache-2.0

extern crate alloc;
// Note: DbgPrint can work in <= DIRQL, so there is no reason using alloc
// crate which may limit the debug printer to work in <= DISPATCH_IRQL.
use core::fmt;

use alloc::ffi::CString;
// We will allocate the format buffer on stack instead of heap
// so that debug printer won't be subject to DISPATCH_IRQL restriction.
struct DebugPrintFormatBuffer {
// Limit buffer to 512 bytes because DbgPrint can only transport 512 bytes per call.
// See: https://learn.microsoft.com/en-us/windows-hardware/drivers/debugger/reading-and-filtering-debugging-messages
buffer: [u8; 512],
used: usize,
}

impl DebugPrintFormatBuffer {
fn new() -> Self {
DebugPrintFormatBuffer {
buffer: [0; 512],
used: 0,
}
}
}

impl fmt::Write for DebugPrintFormatBuffer {
fn write_str(&mut self, s: &str) -> fmt::Result {
let remainder = &mut self.buffer[self.used..];
let current = s.as_bytes();
if remainder.len() < current.len() {
return Err(fmt::Error);
}
remainder[..current.len()].copy_from_slice(current);
self.used += current.len();
return Ok(());
}
}

/// print to kernel debugger via [`wdk_sys::ntddk::DbgPrint`]
#[macro_export]
Expand Down Expand Up @@ -33,20 +64,26 @@ macro_rules! println {
///
/// Panics if an internal null byte is passed in
#[doc(hidden)]
pub fn _print(args: core::fmt::Arguments) {
let formatted_string = CString::new(alloc::format!("{args}"))
.expect("CString should be able to be created from a String.");

// SAFETY: `formatted_string` is a valid null terminated string
unsafe {
#[cfg(any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"))]
{
wdk_sys::ntddk::DbgPrint(formatted_string.as_ptr());
}
pub fn _print(args: fmt::Arguments) {
// Use stack-based formatter. Avoid heap allocation.
let mut w = DebugPrintFormatBuffer::new();
let r = fmt::write(&mut w, args);
if let Ok(_) = r {
let formatted_string = &mut w.buffer;
let formatted_string_pointer = formatted_string.as_ptr() as *const i8;
// No need to append a null-terminator in that the formatted string buffer was zero-initialized.
unsafe {
#[cfg(any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"))]
{
// Use "%s" to prevent the system from reformatting our message.
// It's possible the message can contain keywords like "%s" "%d" etc.
wdk_sys::ntddk::DbgPrint("%s\0".as_ptr() as *const i8, formatted_string_pointer);
}

#[cfg(driver_model__driver_type = "UMDF")]
{
wdk_sys::windows::OutputDebugStringA(formatted_string.as_ptr());
#[cfg(driver_model__driver_type = "UMDF")]
{
wdk_sys::windows::OutputDebugStringA(formatted_string_pointer);
}
}
}
}

0 comments on commit 0402d08

Please sign in to comment.