From ab58e472642d7626be8362056fe727634383d768 Mon Sep 17 00:00:00 2001
From: flux <49762827+0xflux@users.noreply.github.com>
Date: Fri, 6 Dec 2024 21:05:23 +0000
Subject: [PATCH] Implement `IoGetCurrentIrpStackLocation` function in
 `ntddk.rs`

This commit adds the `IoGetCurrentIrpStackLocation` function, which retrieves a pointer to the caller's I/O stack location from a given IRP.

- The function is implemented to access raw pointers and union fields.
- Documentation has been added to describe the parameters, return value, and safety considerations.
- The function includes an `assert!` to ensure `CurrentLocation` is valid, guarding against invalid IRPs. This was implemented as per the original C implementation.

## Safety

This function requires the caller to guarantee that the provided IRP pointer is valid and properly initialised, as it operates on raw memory.

Future improvements could include replacing the `assert!` with proper error handling or returning an `Option` to handle invalid IRPs gracefully. Given that the original implementation in the C libraries does not implement this behaviour, I have not added it to this commit, but can amend if required.
---
 crates/wdk-sys/src/lib.rs   |  3 +++
 crates/wdk-sys/src/ntddk.rs | 31 +++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/crates/wdk-sys/src/lib.rs b/crates/wdk-sys/src/lib.rs
index e6627c2e..cbd7af31 100644
--- a/crates/wdk-sys/src/lib.rs
+++ b/crates/wdk-sys/src/lib.rs
@@ -150,3 +150,6 @@ macro_rules! PAGED_CODE {
         debug_assert!(unsafe { KeGetCurrentIrql() <= APC_LEVEL as u8 });
     };
 }
+
+#[cfg(any(driver_model__driver_type = "WDM", driver_model__driver_type = "KMDF"))]
+pub use ntddk::IoGetCurrentIrpStackLocation;
\ No newline at end of file
diff --git a/crates/wdk-sys/src/ntddk.rs b/crates/wdk-sys/src/ntddk.rs
index d7a703e9..e3ac22b5 100644
--- a/crates/wdk-sys/src/ntddk.rs
+++ b/crates/wdk-sys/src/ntddk.rs
@@ -8,6 +8,7 @@
 //! module, but are available in the top-level `wdk_sys` module.
 
 pub use bindings::*;
+use crate::{PIRP, PIO_STACK_LOCATION};
 
 #[allow(missing_docs)]
 mod bindings {
@@ -18,3 +19,33 @@ mod bindings {
 
     include!(concat!(env!("OUT_DIR"), "/ntddk.rs"));
 }
+
+/// The IoGetCurrentIrpStackLocation routine returns a pointer to the caller's I/O stack location in 
+/// the specified IRP.
+/// 
+/// # Parameters
+/// - irp: PIRP - A pointer to the IRP.
+/// 
+/// # Returns
+/// IoGetCurrentIrpStackLocation returns a pointer to an IO_STACK_LOCATION structure that contains 
+/// the I/O stack location for the driver.
+///
+/// # Safety
+/// This function directly accesses raw pointers and must only be used
+/// when it is guaranteed that the provided `irp` is valid and properly
+/// initialised. Using an invalid or uninitialised `irp` will result
+/// in undefined behavior.
+#[allow(non_snake_case)]
+pub unsafe fn IoGetCurrentIrpStackLocation(irp: PIRP) -> PIO_STACK_LOCATION {
+    unsafe { 
+        assert!((*irp).CurrentLocation <= (*irp).StackCount + 1);
+    
+        // Access the union fields inside the IRP
+        (*irp)
+        .Tail
+        .Overlay
+        .__bindgen_anon_2
+        .__bindgen_anon_1
+        .CurrentStackLocation
+    }
+}
\ No newline at end of file