From 08b059bee98543927db5094a40d17e781d785c6f Mon Sep 17 00:00:00 2001
From: donnaskiez <lachiehodges7785@gmail.com>
Date: Sat, 22 Jun 2024 00:22:11 +1000
Subject: [PATCH] session fix

---
 driver/crypt.c   | 22 ----------------------
 driver/crypt.h   | 23 +++++++++++++++++++++++
 driver/io.c      |  9 ++++++---
 driver/session.c |  9 +++++----
 4 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/driver/crypt.c b/driver/crypt.c
index 5c40c4d..b4607e3 100644
--- a/driver/crypt.c
+++ b/driver/crypt.c
@@ -11,8 +11,6 @@
 #include <immintrin.h>
 #include <bcrypt.h>
 
-#define XOR_ROTATION_AMT 13
-
 FORCEINLINE
 STATIC
 UINT64
@@ -41,26 +39,6 @@ CryptXorKeyGenerate_uint64()
     return CryptGenerateRandomKey64(&seed);
 }
 
-VOID
-CryptEncryptPointer64(_Inout_ PUINT64 Pointer, _In_ UINT64 Key)
-{
-    *Pointer = _rotl64(*Pointer ^ Key, XOR_ROTATION_AMT);
-}
-
-VOID
-CryptDecryptPointer64(_Inout_ PUINT64 Pointer, _In_ UINT64 Key)
-{
-    *Pointer = _rotr64(*Pointer, XOR_ROTATION_AMT) ^ Key;
-}
-
-UINT64
-CryptDecryptPointerOutOfPlace64(_In_ PUINT64 Pointer, _In_ UINT64 Key)
-{
-    volatile UINT64 temp = *Pointer;
-    CryptDecryptPointer64(&temp, Key);
-    return temp;
-}
-
 VOID
 CryptEncryptImportsArray(_In_ PUINT64 Array, _In_ UINT32 Entries)
 {
diff --git a/driver/crypt.h b/driver/crypt.h
index 8461b83..6636618 100644
--- a/driver/crypt.h
+++ b/driver/crypt.h
@@ -3,7 +3,30 @@
 
 #include "common.h"
 
+#define XOR_ROTATION_AMT 13
 
+FORCEINLINE
+VOID
+CryptEncryptPointer64(_Inout_ PUINT64 Pointer, _In_ UINT64 Key)
+{
+    *Pointer = _rotl64(*Pointer ^ Key, XOR_ROTATION_AMT);
+}
+
+FORCEINLINE
+VOID
+CryptDecryptPointer64(_Inout_ PUINT64 Pointer, _In_ UINT64 Key)
+{
+    *Pointer = _rotr64(*Pointer, XOR_ROTATION_AMT) ^ Key;
+}
+
+FORCEINLINE
+UINT64
+CryptDecryptPointerOutOfPlace64(_In_ PUINT64 Pointer, _In_ UINT64 Key)
+{
+    volatile UINT64 temp = *Pointer;
+    CryptDecryptPointer64(&temp, Key);
+    return temp;
+}
 
 VOID
 CryptEncryptImportsArray(_In_ PUINT64 Array, _In_ UINT32 Entries);
diff --git a/driver/io.c b/driver/io.c
index c3ed88c..1a5b1d7 100644
--- a/driver/io.c
+++ b/driver/io.c
@@ -1179,9 +1179,12 @@ DeviceClose(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp)
 
     /* This needs to be fixed lol, cos anyone can just open a handle whhich
      * might not begin a session.*/
-    SessionTerminate();
-    UnregisterProcessObCallbacks();
-    SharedMappingTerminate();
+
+    if (GetActiveSession()->is_session_active) {
+        SessionTerminate();
+        UnregisterProcessObCallbacks();
+        SharedMappingTerminate();
+    }
 
     IoCompleteRequest(Irp, IO_NO_INCREMENT);
     return Irp->IoStatus.Status;
diff --git a/driver/session.c b/driver/session.c
index bd02daf..a4ae1fc 100644
--- a/driver/session.c
+++ b/driver/session.c
@@ -148,10 +148,9 @@ SessionInitialise(_In_ PIRP Irp)
         goto end;
     }
 
-    session->km_handle         = ImpPsGetProcessId(process);
-    session->process           = process;
-    session->is_session_active = TRUE;
-    session->cookie            = initiation->cookie;
+    session->km_handle = ImpPsGetProcessId(process);
+    session->process   = process;
+    session->cookie    = initiation->cookie;
 
     RtlCopyMemory(session->aes_key, initiation->aes_key, AES_256_KEY_SIZE);
     RtlCopyMemory(session->iv, initiation->aes_iv, AES_256_IV_SIZE);
@@ -182,6 +181,8 @@ SessionInitialise(_In_ PIRP Irp)
 
     FindOurUserModeModuleEntry(HashOurUserModuleOnEntryCallback, session);
 
+    session->is_session_active = TRUE;
+
 end:
     KeReleaseGuardedMutex(&session->lock);
     return status;