From efc4be8a978cc8d86017bbbd6294d1ca6b0852cb Mon Sep 17 00:00:00 2001
From: Dave Thaler <dthaler@microsoft.com>
Date: Tue, 1 Aug 2023 17:24:53 -0700
Subject: [PATCH] Add KfRaiseIrql

The DDK just defines KeRaiseIrql as inline wrapper around KfRaiseIrql
so we need to expose KfRaiseIrql when linking.

Addresses part of #69

Signed-off-by: Dave Thaler <dthaler@microsoft.com>
---
 inc/usersim/ke.h  |  4 ++++
 src/ke.cpp        | 12 ++++++++++--
 tests/ke_test.cpp | 17 +++++++++++++++++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/inc/usersim/ke.h b/inc/usersim/ke.h
index 3f4ca95..ea36440 100644
--- a/inc/usersim/ke.h
+++ b/inc/usersim/ke.h
@@ -67,6 +67,10 @@ extern "C"
     VOID
     KeRaiseIrql(_In_ KIRQL new_irql, _Out_ PKIRQL old_irql);
 
+    USERSIM_API
+    _IRQL_requires_max_(HIGH_LEVEL) _IRQL_raises_(new_irql) _IRQL_saves_ KIRQL
+    KfRaiseIrql(_In_ KIRQL new_irql);
+
     USERSIM_API
     KIRQL
     KeRaiseIrqlToDpcLevel();
diff --git a/src/ke.cpp b/src/ke.cpp
index a3d8ce0..0767864 100644
--- a/src/ke.cpp
+++ b/src/ke.cpp
@@ -68,12 +68,18 @@ KeGetCurrentIrql() { return _usersim_current_irql; }
 VOID
 KeRaiseIrql(_In_ KIRQL new_irql, _Out_ PKIRQL old_irql)
 {
-    *old_irql = KeGetCurrentIrql();
+    *old_irql = KfRaiseIrql(new_irql);
+}
+
+KIRQL
+KfRaiseIrql(_In_ KIRQL new_irql)
+{
+    KIRQL old_irql = KeGetCurrentIrql();
     _usersim_current_irql = new_irql;
     BOOL result = SetThreadPriority(GetCurrentThread(), new_irql);
     ASSERT(result);
 
-    if (new_irql >= DISPATCH_LEVEL && *old_irql < DISPATCH_LEVEL) {
+    if (new_irql >= DISPATCH_LEVEL && old_irql < DISPATCH_LEVEL) {
         PROCESSOR_NUMBER processor;
         uint32_t processor_index = KeGetCurrentProcessorNumberEx(&processor);
 
@@ -85,6 +91,8 @@ KeRaiseIrql(_In_ KIRQL new_irql, _Out_ PKIRQL old_irql)
 
         _usersim_dispatch_locks[processor_index].lock();
     }
+
+    return old_irql;
 }
 
 KIRQL
diff --git a/tests/ke_test.cpp b/tests/ke_test.cpp
index d4d6bb7..81dc260 100644
--- a/tests/ke_test.cpp
+++ b/tests/ke_test.cpp
@@ -36,6 +36,23 @@ TEST_CASE("irql", "[ke]")
     REQUIRE(KeGetCurrentIrql() == PASSIVE_LEVEL);
 }
 
+TEST_CASE("KfRaiseIrql", "[ke]")
+{
+    REQUIRE(KeGetCurrentIrql() == PASSIVE_LEVEL);
+
+    KIRQL old_irql;
+    old_irql = KfRaiseIrql(DISPATCH_LEVEL);
+    REQUIRE(old_irql == PASSIVE_LEVEL);
+    REQUIRE(KeGetCurrentIrql() == DISPATCH_LEVEL);
+
+    old_irql = KfRaiseIrql(DISPATCH_LEVEL);
+    REQUIRE(old_irql == DISPATCH_LEVEL);
+    REQUIRE(KeGetCurrentIrql() == DISPATCH_LEVEL);
+
+    KeLowerIrql(PASSIVE_LEVEL);
+    REQUIRE(KeGetCurrentIrql() == PASSIVE_LEVEL);
+}
+
 TEST_CASE("spin lock", "[ke]")
 {
     KSPIN_LOCK lock;