Better way of getting handle for ZwQueryVirtualMemory for target process

DriverMain.cpp

@@ -84,20 +84,20 @@ VOID ProcessNotifyRoutine(
 
     if (Create)
     {
-       PushLockExclusive lock(&global.ProtectedProcess.SharedLock);
-
-       if (!global.ProtectedProcess.ProcessObject)
-       {
-           KdPrint(("[SimpleAntiCheat] Protected process started: %wZ (PID: %p)\n",
-              procImageName.Get(), ProcessId));
-           
-           global.ProtectedProcess.ProcessId = ProcessId;
-           global.ProtectedProcess.ProcessObject = kernel_std::move(processObject);
-       }
+        PushLockExclusive lock(&global.ProtectedProcess.SharedLock);
+
+        if (!global.ProtectedProcess.ProcessObject)
+        {
+            KdPrint(("[SimpleAntiCheat] Protected process started: %wZ (PID: %p)\n",
+                procImageName.Get(), ProcessId));
+
+            global.ProtectedProcess.ProcessId = ProcessId;
+            global.ProtectedProcess.ProcessObject = kernel_std::move(processObject);
+        }
     }
     else
     {
-        PushLockExclusive lock( &global.ProtectedProcess.SharedLock );
+        PushLockExclusive lock(&global.ProtectedProcess.SharedLock);
 
         if (global.ProtectedProcess.ProcessId == ProcessId)
         {
@@ -316,20 +316,23 @@ static NTSTATUS ScanExecutableMemory()
         targetProcess.AddRef();
     }
 
-    // Attach to target process context
-    KAPC_STATE apcState;
-    KeStackAttachProcess(targetProcess.Get(), &apcState);
-
     KdPrint(("[SimpleAntiCheat] === Scanning executable memory regions ===\n"));
 
+    auto result = targetProcess.GetKernelHandle();
+
+    if (!result.Success())
+    {
+        return result.Status;
+    }
+
     PVOID baseAddress = nullptr;
     while (baseAddress < MmHighestUserAddress)
     {
         MEMORY_BASIC_INFORMATION memInfo;
         SIZE_T returnLength;
 
         NTSTATUS status = ZwQueryVirtualMemory(
-            ZwCurrentProcess(),
+            result.Value.Get(),
             baseAddress,
             static_cast<MEMORY_INFORMATION_CLASS>(MemoryBasicInformation),
             &memInfo,
@@ -397,8 +400,6 @@ static NTSTATUS ScanExecutableMemory()
         baseAddress = reinterpret_cast<PVOID>(reinterpret_cast<ULONG_PTR>(memInfo.BaseAddress) + memInfo.RegionSize);
     }
 
-    KeUnstackDetachProcess(&apcState);
-
     return STATUS_SUCCESS;
 }
 
@@ -516,7 +517,7 @@ VOID DriverUnload(IN PDRIVER_OBJECT DriverObject)
     if (global.ScanThreadObject)
     {
         KeWaitForSingleObject(global.ScanThreadObject, Executive, KernelMode, false, nullptr);
-        ObDereferenceObject(global.ScanThreadObject );
+        ObDereferenceObject(global.ScanThreadObject);
         global.ScanThreadObject = nullptr;
     }
 
@@ -525,7 +526,6 @@ VOID DriverUnload(IN PDRIVER_OBJECT DriverObject)
         ObUnRegisterCallbacks(global.ObCallbackHandle);
         global.ObCallbackHandle = nullptr;
     }
-
 }
 
 extern "C" NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
@@ -597,7 +597,7 @@ extern "C" NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING Reg
     if (!NT_SUCCESS(status))
     {
         KdPrint(("[SimpleAntiCheat] Failed to reference thread object: 0x%X\n", status));
-		KeSetEvent(&global.ThreadStopEvent, IO_NO_INCREMENT, false);
+        KeSetEvent(&global.ThreadStopEvent, IO_NO_INCREMENT, false);
         ObUnRegisterCallbacks(global.ObCallbackHandle);
         PsSetCreateProcessNotifyRoutine(ProcessNotifyRoutine, true);
         return status;

Utilities.hpp

@@ -1,5 +1,82 @@
 #pragma once
 
+template<typename T>
+struct StatusResult
+{
+    NTSTATUS Status;
+    T Value;
+
+    [[nodiscard]] bool Success() const noexcept
+    {
+        return NT_SUCCESS(Status);
+    }
+};
+
+class KernelHandle
+{
+public:
+    explicit KernelHandle(HANDLE Handle = nullptr) noexcept
+        : Handle_(Handle)
+    {
+    }
+
+    ~KernelHandle() noexcept
+    {
+        if (Handle_)
+        {
+            ZwClose(Handle_);
+        }
+    }
+
+    KernelHandle(const KernelHandle&) = delete;
+    KernelHandle& operator=(const KernelHandle&) = delete;
+
+    KernelHandle(KernelHandle&& Other) noexcept
+    {
+        Reset(Other.Release());
+    }
+
+    KernelHandle& operator=(KernelHandle&& Other) noexcept
+    {
+        if (this != &Other)
+        {
+            Reset(Other.Release());
+        }
+
+        return *this;
+    }
+
+    [[nodiscard]] HANDLE Get() const noexcept
+    {
+        return Handle_;
+    }
+
+    explicit operator bool() const noexcept
+    {
+        return Handle_ != nullptr;
+    }
+
+    void Reset(HANDLE NewHandle = nullptr) noexcept
+    {
+        if (Handle_)
+        {
+            ZwClose(Handle_);
+        }
+
+        Handle_ = NewHandle;
+    }
+
+    HANDLE Release() noexcept
+    {
+        HANDLE handle = Handle_;
+        Handle_ = nullptr;
+        return handle;
+    }
+
+private:
+    HANDLE Handle_;
+};
+
 template<typename P>
 class KernelObjectRef
 {
@@ -64,10 +141,23 @@ class KernelObjectRef
 
     void AddRef()
     {
-        if (Obj_)
+        NT_ASSERT(Obj_);
+        ObReferenceObject(Obj_);
+    }
+
+    // Internally this increases object reference count so we don't need to worry about lifetime
+    [[nodiscard]] StatusResult<KernelHandle> GetKernelHandle(ACCESS_MASK DesiredAccess = 0)
+    {
+        if (!Obj_)
         {
-            ObReferenceObject(Obj_);
+            return { STATUS_INVALID_HANDLE, KernelHandle() };
         }
+
+        HANDLE handle;
+
+        NTSTATUS status = ObOpenObjectByPointer(Obj_, OBJ_KERNEL_HANDLE, nullptr, DesiredAccess, nullptr, KernelMode, &handle);
+
+        return { status, NT_SUCCESS(status) ? KernelHandle(handle) : KernelHandle() };
     }
 
 private: