From 02f35dbd8cc698dfe9595679b4382af5aba4590a Mon Sep 17 00:00:00 2001
From: belliash <belliash@codingworkshop.eu.org>
Date: Fri, 3 Nov 2023 16:04:10 +0100
Subject: [PATCH] Try to allocate new stack for a thread if needed

---
 xtoskrnl/ke/kthread.c | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/xtoskrnl/ke/kthread.c b/xtoskrnl/ke/kthread.c
index 2989790..2e4afe6 100644
--- a/xtoskrnl/ke/kthread.c
+++ b/xtoskrnl/ke/kthread.c
@@ -52,8 +52,13 @@ KeInitializeThread(IN PKPROCESS Process,
                    IN PVOID Stack)
 {
     PKWAIT_BLOCK TimerWaitBlock;
+    BOOLEAN Allocation;
+    XTSTATUS Status;
     ULONG Index;
 
+    /* No stack allocation was done yet */
+    Allocation = FALSE;
+
     /* Initialize thread dispatcher header */
     Thread->Header.Type = ThreadObject;
     Thread->Header.SignalState = 0;
@@ -118,6 +123,21 @@ KeInitializeThread(IN PKPROCESS Process,
     /* Initialize Thread Environment Block*/
     Thread->EnvironmentBlock = EnvironmentBlock;
 
+    /* Make sure there is a valid stack available */
+    if(!Stack)
+    {
+        /* Allocate new stack */
+        Status = MmAllocateKernelStack(&Stack, FALSE, 0);
+        if(Status != STATUS_SUCCESS || !Stack)
+        {
+            /* Stack allocation failed */
+            return STATUS_INSUFFICIENT_RESOURCES;
+        }
+
+        /* Mark allocation as successful */
+        Allocation = TRUE;
+    }
+
     Thread->InitialStack = Stack;
     Thread->StackBase = Stack;
     Thread->StackLimit = Stack - KERNEL_STACK_SIZE;
@@ -129,7 +149,16 @@ KeInitializeThread(IN PKPROCESS Process,
     }
     __except(EXCEPTION_EXECUTE_HANDLER)
     {
-        /* Failed to initialize thread context */
+        /* Failed to initialize thread context, check stack allocation */
+        if(Allocation)
+        {
+            /* Deallocate stack */
+            MmFreeKernelStack(Stack, FALSE);
+            Thread->InitialStack = NULL;
+            Thread->StackBase = NULL;
+        }
+
+        /* Thread initialization failed */
         return STATUS_UNSUCCESSFUL;
     }