Use stack-based formatter for debug-printing.

crates/wdk-build/src/utils.rs

@@ -344,7 +344,7 @@ fn read_registry_key_string_value(
                 return Some(
                     CStr::from_bytes_with_nul(&buffer[..len as usize])
                         .expect(
-                            "RegGetValueA should always return a null terminated string. The read \
+                            "RegGetValueA should always return a null-terminated string. The read \
                              string (REG_SZ) from the registry should not contain any interior \
                              nulls.",
                         )

crates/wdk/Cargo.toml

@@ -22,6 +22,7 @@ tracing-subscriber = { workspace = true, features = ["env-filter"] }
 wdk-build.workspace = true
 
 [dependencies]
+cfg-if.workspace = true
 wdk-sys.workspace = true
 
 [dev-dependencies]

crates/wdk/src/lib.rs

@@ -5,7 +5,10 @@
 //! built on top of the raw FFI bindings provided by [`wdk-sys`], and provides a
 //! safe, idiomatic rust interface to the WDK.
 
-#![no_std]
+#![cfg_attr(
+    any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"),
+    no_std
+)]
 
 #[cfg(any(
     all(

crates/wdk/src/print.rs

@@ -1,19 +1,71 @@
 // Copyright (c) Microsoft Corporation
 // License: MIT OR Apache-2.0
 
-extern crate alloc;
+use core::fmt;
+#[cfg(driver_model__driver_type = "UMDF")]
+use std::ffi::CString;
 
-use alloc::ffi::CString;
+/// Prints to the debugger.
+///
+/// Equivalent to the println! macro except that a newline is not printed at the
+/// end of the message.
+#[cfg_attr(
+    any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"),
+    doc = r"
+The output is routed to the debugger via [`wdk_sys::ntddk::DbgPrint`], so the `IRQL` 
+requirements of that function apply. In particular, this should only be called at 
+`IRQL` <= `DIRQL`, and calling it at `IRQL` > `DIRQL` can cause deadlocks due to
+the debugger's use of IPIs (Inter-Process Interrupts).
+
+[`wdk_sys::ntddk::DbgPrint`]'s 512 byte limit does not apply to this macro, as it will
+automatically buffer and chunk the output if it exceeds that limit.
+"
+)]
+#[cfg_attr(
+    driver_model__driver_type = "UMDF",
+    doc = r#"
+The output is routed to the debugger via [`wdk_sys::windows::OutputDebugStringA`].
 
-/// print to kernel debugger via [`wdk_sys::ntddk::DbgPrint`]
+If there is no debugger attached to WUDFHost of the driver (i.e., user-mode debugging),
+the output will be routed to the system debugger (i.e., kernel-mode debugging).
+"#
+)]
+/// See the formatting documentation in [`core::fmt`] for details of the macro
+/// argument syntax.
 #[macro_export]
 macro_rules! print {
     ($($arg:tt)*) => {
       ($crate::_print(format_args!($($arg)*)))
     };
 }
 
-/// print with newline to debugger via [`wdk_sys::ntddk::DbgPrint`]
+/// Prints to the debugger, with a newline.
+///
+/// This macro uses the same syntax as [`core::format!`], but writes to the
+/// debugger instead. See [`core::fmt`] for more information.
+#[cfg_attr(
+    any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"),
+    doc = r"
+The output is routed to the debugger via [`wdk_sys::ntddk::DbgPrint`], so the `IRQL` 
+requirements of that function apply. In particular, this should only be called at 
+`IRQL` <= `DIRQL`, and calling it at `IRQL` > `DIRQL` can cause deadlocks due to
+the debugger's use of IPIs (Inter-Process Interrupts).
+
+[`wdk_sys::ntddk::DbgPrint`]'s 512 byte limit does not apply to this macro, as it will
+automatically buffer and chunk the output if it exceeds that limit.
+"
+)]
+#[cfg_attr(
+    driver_model__driver_type = "UMDF",
+    doc = r"
+The output is routed to the debugger via [`wdk_sys::windows::OutputDebugStringA`].
+
+If there is no debugger attached to WUDFHost of the driver (i.e., user-mode debugging),
+the output will be routed to the system debugger (i.e., kernel-mode debugging).
+"
+)]
+/// See the formatting documentation in [`core::fmt`] for details of the macro
+/// argument syntax.
 #[macro_export]
 macro_rules! println {
     () => {
@@ -33,20 +85,220 @@ 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) {
+    cfg_if::cfg_if! {
+        if #[cfg(any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"))] {
+            let mut buffered_writer = dbg_print_buf_writer::DbgPrintBufWriter::new();
+
+            // TODO: handle internal bytes?
+
+            if let Ok(_) = fmt::write(&mut buffered_writer, args) {
+                buffered_writer.flush();
+            } else {
+                unreachable!("DbgPrintBufWriter should never fail write");
+            }
+
+        } else if #[cfg(driver_model__driver_type = "UMDF")] {
+            match CString::new(format!("{args}")) {
+                Ok(c_string) => {
+                    // SAFETY: `CString` guarantees a valid null-terminated string
+                    unsafe {
+                        wdk_sys::windows::OutputDebugStringA(c_string.as_ptr());
+                    }
+                },
+                Err(nul_error) => {
+                    let nul_position = nul_error.nul_position();
+                    let string_vec = nul_error.into_vec();
+                    let c_string = CString::new(&string_vec[..nul_position]).expect("string_vec[..nul_position] should have no internal null bytes");
+                    let remaining_string = str::from_utf8(&string_vec[nul_position+1 ..]).expect("string_vec should always be valid UTF-8 because `format!` returns a String");
+
+                    // SAFETY: `CString` guarantees a valid null-terminated string
+                    unsafe {
+                        wdk_sys::windows::OutputDebugStringA(c_string.as_ptr());
+                    }
+
+                    print!("{remaining_string}");
+                }
+            }
+        }
+    }
+}
+
+#[cfg(any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"))]
+mod dbg_print_buf_writer {
+    use core::fmt;
+
+    /// Max size that can be transmitted by `DbgPrint` in single call:
+    /// <https://learn.microsoft.com/en-us/windows-hardware/drivers/debugger/reading-and-filtering-debugging-messages#dbgprint-buffer-and-the-debugger>
+    const DBG_PRINT_MAX_TXN_SIZE: usize = 512;
+
+    // We will allocate the format buffer on stack instead of heap
+    // so that debug printer won't be subject to DISPATCH_IRQL restriction.
+
+    /// Stack-based format buffer for `DbgPrint`
+    ///
+    /// This buffer is used to format strings via `fmt::write` without needing
+    /// heap allocations. Whenever a new string would cause the buffer to exceed
+    /// its max capacity, it will first empty its buffer via `DbgPrint`.
+    /// The use of a stack-based buffer instead of `alloc::format!` allows for
+    /// printing at IRQL <= DIRQL.
+    pub struct DbgPrintBufWriter {
+        buffer: [u8; DBG_PRINT_MAX_TXN_SIZE],
+        used: usize,
+    }
+
+    impl Default for DbgPrintBufWriter {
+        fn default() -> Self {
+            Self {
+                // buffer is initialized to all null
+                buffer: [0; DBG_PRINT_MAX_TXN_SIZE],
+                used: 0,
+            }
+        }
+    }
+
+    impl fmt::Write for DbgPrintBufWriter {
+        fn write_str(&mut self, s: &str) -> fmt::Result {
+            let mut str_byte_slice = s.as_bytes();
+            let mut remaining_buffer = &mut self.buffer[self.used..Self::USABLE_BUFFER_SIZE];
+            let mut remaining_buffer_size = remaining_buffer.len();
+
+            // If the string is too large for the buffer, keep chunking the string and
+            // flushing the buffer until the entire string is handled
+            while str_byte_slice.len() > remaining_buffer_size {
+                // Fill buffer
+                remaining_buffer[..].copy_from_slice(&str_byte_slice[..remaining_buffer_size]);
+
+                // Flush buffer
+                self.flush();
+
+                // Update remaining string slice to handle and reset remaining buffer
+                str_byte_slice = &str_byte_slice[remaining_buffer_size..];
+                remaining_buffer = &mut self.buffer[self.used..Self::USABLE_BUFFER_SIZE];
+                remaining_buffer_size = remaining_buffer.len();
+            }
+            remaining_buffer[..str_byte_slice.len()].copy_from_slice(str_byte_slice);
+            self.used += str_byte_slice.len();
+
+            Ok(())
+        }
+    }
+
+    impl DbgPrintBufWriter {
+        /// The maximum size of the buffer that can be used for formatting
+        /// strings
+        ///
+        /// The last byte is reserved for the null terminator
+        const USABLE_BUFFER_SIZE: usize = DBG_PRINT_MAX_TXN_SIZE - 1;
+
+        pub fn new() -> Self {
+            Self::default()
+        }
+
+        pub fn flush(&mut self) {
+            // SAFETY: This is safe because:
+            // 1. `self.buffer` contains a valid C-style string with the data placed in
+            //    [0..self.used] by the `write_str` implementation
+            // 2. The `write_str` method ensures `self.used` never exceeds
+            //    `USABLE_BUFFER_SIZE`, leaving the last byte available for null termination
+            // 3. The "%s" format specifier is used as a literal string to prevent
+            //    `DbgPrint` from interpreting format specifiers in the message, which could
+            //    lead to memory corruption or undefined behavior if the buffer contains
+            //    printf-style formatting characters
+            unsafe {
+                wdk_sys::ntddk::DbgPrint(
+                    c"%s".as_ptr().cast(),
+                    self.buffer.as_ptr().cast::<wdk_sys::CHAR>(),
+                );
+            }
+
+            self.used = 0;
         }
+    }
+
+    #[cfg(test)]
+    mod tests {
+        use super::*;
+        use crate::print::dbg_print_buf_writer::DbgPrintBufWriter;
+
+        #[test]
+        fn write_that_fits_buffer() {
+            const TEST_STRING: &str = "Hello, world!";
+            const TEST_STRING_LEN: usize = TEST_STRING.len();
 
-        #[cfg(driver_model__driver_type = "UMDF")]
-        {
-            wdk_sys::windows::OutputDebugStringA(formatted_string.as_ptr());
+            let mut writer = DbgPrintBufWriter::new();
+            fmt::write(&mut writer, format_args!("{TEST_STRING}"))
+                .expect("fmt::write should succeed");
+            assert_eq!(writer.used, TEST_STRING_LEN);
+            assert_eq!(&writer.buffer[..writer.used], TEST_STRING.as_bytes());
+
+            writer.flush();
+            // FIXME: When this test is compiled, rustc automatically links the
+            // usermode-version of DbgPrint. We should either figure out a way to prevent
+            // this in order to stub in a mock implementation via something like `mockall`,
+            // or have `DbgPrintBufWriter` be able to be instantiated with a different
+            // implementation somehow. Ex. `DbgPrintBufWriter::new` can take in a closure
+            // that gets called for flushing (real impl uses Dbgprint and test impl uses a
+            // mock with a counter and some way to validate contents being sent to the flush
+            // closure)
+            assert_eq!(writer.used, 0);
+        }
+
+        #[test]
+        fn write_that_exceeds_buffer() {
+            const TEST_STRING: &str =
+                "This is a test string that exceeds the buffer size limit set for \
+                 DbgPrintBufWriter. It should trigger multiple flushes to handle the overflow \
+                 correctly. The buffer has a limited capacity of 511 bytes (512 minus 1 for null \
+                 terminator), and this string is intentionally much longer. When writing this \
+                 string to the DbgPrintBufWriter, the implementation should automatically chunk \
+                 the content and flush each chunk separately. This ensures large debug messages \
+                 can be properly displayed without being truncated. The current implementation \
+                 handles this by filling the buffer as much as possible, flushing it using \
+                 DbgPrint, then continuing with the remaining content until everything is \
+                 processed. This approach allows debugging messages of arbitrary length without \
+                 requiring heap allocations, which is particularly important in kernel mode where \
+                 memory allocation constraints might be stricter. This test verifies that strings \
+                 larger than the max buffer size are handled correctly, confirming that our \
+                 buffer management logic works as expected. This string is now well over 1000 \
+                 characters long to ensure that the DbgPrintBufWriter's buffer overflow handling \
+                 is thoroughly tested.";
+            const TEST_STRING_LEN: usize = TEST_STRING.len();
+            const UNFLUSHED_STRING_CONTENTS_STARTING_INDEX: usize =
+                TEST_STRING_LEN - (TEST_STRING_LEN % DbgPrintBufWriter::USABLE_BUFFER_SIZE);
+
+            const {
+                assert!(
+                    TEST_STRING_LEN > DbgPrintBufWriter::USABLE_BUFFER_SIZE,
+                    "TEST_STRING_LEN should be greater than buffer size for this test"
+                );
+            }
+
+            let expected_unflushed_string_contents =
+                &TEST_STRING[UNFLUSHED_STRING_CONTENTS_STARTING_INDEX..];
+
+            let mut writer = DbgPrintBufWriter::new();
+            fmt::write(&mut writer, format_args!("{TEST_STRING}"))
+                .expect("fmt::write should succeed");
+            assert_eq!(writer.used, expected_unflushed_string_contents.len());
+            assert_eq!(
+                &writer.buffer[..writer.used],
+                expected_unflushed_string_contents.as_bytes()
+            );
+            // FIXME: When this test is compiled, rustc automatically links the
+            // usermode-version of DbgPrint. We should either figure out a way to prevent
+            // this in order to stub in a mock implementation via something like `mockall`,
+            // or have `DbgPrintBufWriter` be able to be instantiated with a different
+            // implementation somehow. Ex. `DbgPrintBufWriter::new` can take in a closure
+            // that gets called for flushing (real impl uses Dbgprint and test impl uses a
+            // mock with a counter and some way to validate contents being sent to the flush
+            // closure)
+
+            writer.flush();
+            assert_eq!(writer.used, 0);
         }
     }
+
+    // FIXME: add tests for no internal null bytes, string w/ exactly USABLE
+    // Buffer len, string w/ exactly Buffer len
 }