From 80d9d730343c3ddd12a6eb784cbdc0ec2a1edd0b Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Wed, 16 Mar 2022 20:07:15 +0100
Subject: [PATCH] Register power callback

---
 CMakeLists.txt          | 13 +++++++
 external/CMakeLists.txt |  1 +
 src/driver/main.cpp     | 84 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index f374ae1..e37013d 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -14,6 +14,19 @@ set_property(GLOBAL PROPERTY USE_FOLDERS ON)
 
 ##########################################
 
+set(CMAKE_MSVC_RUNTIME_LIBRARY MultiThreaded$<$<CONFIG:Debug>:Debug>)
+
+if(MSVC)
+  add_link_options(
+    $<$<NOT:$<STREQUAL:${CMAKE_MSVC_RUNTIME_LIBRARY},MultiThreaded>>:/NODEFAULTLIB:libcmt.lib>
+    $<$<NOT:$<STREQUAL:${CMAKE_MSVC_RUNTIME_LIBRARY},MultiThreadedDLL>>:/NODEFAULTLIB:msvcrt.lib>
+    $<$<NOT:$<STREQUAL:${CMAKE_MSVC_RUNTIME_LIBRARY},MultiThreadedDebug>>:/NODEFAULTLIB:libcmtd.lib>
+    $<$<NOT:$<STREQUAL:${CMAKE_MSVC_RUNTIME_LIBRARY},MultiThreadedDebugDLL>>:/NODEFAULTLIB:msvcrtd.lib>
+  )
+endif()
+
+##########################################
+
 include(cmake/utils.cmake)
 
 ##########################################
diff --git a/external/CMakeLists.txt b/external/CMakeLists.txt
index 82217e9..4d181ec 100644
--- a/external/CMakeLists.txt
+++ b/external/CMakeLists.txt
@@ -30,6 +30,7 @@ ExternalProject_Add(
     "-DCMAKE_CXX_FLAGS_DEBUG=${CMAKE_CXX_FLAGS_DEBUG}"
     "-DCMAKE_CXX_FLAGS_RELEASE=${CMAKE_CXX_FLAGS_RELEASE}"
 
+    "-DCMAKE_MSVC_RUNTIME_LIBRARY=${CMAKE_MSVC_RUNTIME_LIBRARY}"
     
     "-DCMAKE_TOOLCHAIN_FILE=${CMAKE_TOOLCHAIN_FILE}"
 )
diff --git a/src/driver/main.cpp b/src/driver/main.cpp
index eaddd85..7ca416e 100644
--- a/src/driver/main.cpp
+++ b/src/driver/main.cpp
@@ -157,12 +157,93 @@ NTSTATUS create_io_device(const PDRIVER_OBJECT DriverObject)
 	return Status;
 }
 
+_Function_class_(CALLBACK_FUNCTION)
+VOID
+PowerCallback(
+	_In_opt_ PVOID CallbackContext,
+	_In_opt_ PVOID Argument1,
+	_In_opt_ PVOID Argument2
+)
+{
+	UNREFERENCED_PARAMETER(CallbackContext);
+
+	//
+	// Ignore non-Sx changes
+	//
+	if (Argument1 != (PVOID)PO_CB_SYSTEM_STATE_LOCK)
+	{
+		return;
+	}
+
+	//
+	// Check if this is S0->Sx, or Sx->S0
+	//
+	if (ARGUMENT_PRESENT(Argument2))
+	{
+		//
+		// Reload the hypervisor
+		//
+		debug_log("Waking up!\n");
+	}
+	else
+	{
+		//
+		// Unload the hypervisor
+		//
+		debug_log("Going to sleep!\n");
+	}
+}
+
+PVOID g_PowerCallbackRegistration{nullptr};
+
+NTSTATUS register_sleep_callback()
+{
+	PCALLBACK_OBJECT callbackObject;
+	UNICODE_STRING callbackName =
+		RTL_CONSTANT_STRING(L"\\Callback\\PowerState");
+	OBJECT_ATTRIBUTES objectAttributes =
+		RTL_CONSTANT_OBJECT_ATTRIBUTES(&callbackName,
+		                               OBJ_CASE_INSENSITIVE |
+		                               OBJ_KERNEL_HANDLE);
+
+	auto status = ExCreateCallback(&callbackObject, &objectAttributes, FALSE, TRUE);
+	if (!NT_SUCCESS(status))
+	{
+		return status;
+	}
+
+	//
+	// Now register our routine with this callback
+	//
+	g_PowerCallbackRegistration = ExRegisterCallback(callbackObject,
+	                                                 PowerCallback,
+	                                                 NULL);
+
+	//
+	// Dereference it in both cases -- either it's registered, so that is now
+	// taking a reference, and we'll unregister later, or it failed to register
+	// so we failing now, and it's gone.
+	//
+	ObDereferenceObject(callbackObject);
+
+	//
+	// Fail if we couldn't register the power callback
+	//
+	if (g_PowerCallbackRegistration == NULL)
+	{
+		return STATUS_INSUFFICIENT_RESOURCES;
+	}
+
+	return STATUS_SUCCESS;
+}
+
 _Function_class_(DRIVER_UNLOAD)
 
 void unload(PDRIVER_OBJECT DriverObject)
 {
 	debug_log("Leaving World\n");
 	IrpUnloadHandler(DriverObject);
+	ExUnregisterCallback(g_PowerCallbackRegistration);
 }
 
 void throw_test()
@@ -198,7 +279,8 @@ extern "C" NTSTATUS DriverEntry(const PDRIVER_OBJECT DriverObject, PUNICODE_STRI
 	debug_log("Final i = %i\n", i);
 
 	throw_test();
-	
+	register_sleep_callback();
+
 	return create_io_device(DriverObject);
 
 	//return STATUS_SUCCESS;