From 5677719038a9745145ef3d2d70fb8045f6ad4298 Mon Sep 17 00:00:00 2001
From: belliash <belliash@codingworkshop.eu.org>
Date: Tue, 14 Mar 2023 17:37:20 +0100
Subject: [PATCH] Initialize thread context for AMD64

---
 sdk/xtdk/amd64/ketypes.h    | 29 ++++++++++++++
 sdk/xtdk/ketypes.h          |  1 +
 xtoskrnl/ke/amd64/kthread.c | 80 ++++++++++++++++++++++++++++++++++++-
 3 files changed, 109 insertions(+), 1 deletion(-)

diff --git a/sdk/xtdk/amd64/ketypes.h b/sdk/xtdk/amd64/ketypes.h
index ef54530..fe83506 100644
--- a/sdk/xtdk/amd64/ketypes.h
+++ b/sdk/xtdk/amd64/ketypes.h
@@ -88,6 +88,14 @@
 #define EFLAGS_ID_MASK                            0x00200000L
 #define EFLAGS_USER_SANITIZE                      0x003F4DD7L
 
+/* Context control flags */
+#define CONTEXT_ARCHITECTURE                      0x00100000
+#define CONTEXT_CONTROL                           (CONTEXT_ARCHITECTURE | 0x01)
+#define CONTEXT_INTEGER                           (CONTEXT_ARCHITECTURE | 0x02)
+#define CONTEXT_SEGMENTS                          (CONTEXT_ARCHITECTURE | 0x04)
+#define CONTEXT_FLOATING_POINT                    (CONTEXT_ARCHITECTURE | 0x08)
+#define CONTEXT_DEBUG_REGISTERS                   (CONTEXT_ARCHITECTURE | 0x10)
+
 /* Interrupt request levels definitions */
 #define PASSIVE_LEVEL                             0
 #define LOW_LEVEL                                 0
@@ -321,6 +329,17 @@ typedef struct _KEXCEPTION_FRAME
     ULONG64 Return;
 } KEXCEPTION_FRAME, *PKEXCEPTION_FRAME;
 
+/* Thread start frame definition */
+typedef struct _KSTART_FRAME
+{
+    ULONG64 P1Home;
+    ULONG64 P2Home;
+    ULONG64 P3Home;
+    ULONG64 P4Home;
+    ULONG64 Reserved;
+    ULONG64 Return;
+} KSTART_FRAME, *PKSTART_FRAME;
+
 /* Switch frame definition */
 typedef struct _KSWITCH_FRAME
 {
@@ -422,6 +441,16 @@ typedef struct _KTRAP_FRAME
     ULONG CodePatchCycle;
 } KTRAP_FRAME, *PKTRAP_FRAME;
 
+/* Thread initialization frame definition */
+typedef struct _KTHREAD_INIT_FRAME
+{
+    KSWITCH_FRAME SwitchFrame;
+    KSTART_FRAME StartFrame;
+    KEXCEPTION_FRAME ExceptionFrame;
+    KTRAP_FRAME TrapFrame;
+    FLOATING_SAVE_AREA NpxFrame;
+} KTHREAD_INIT_FRAME, *PKTHREAD_INIT_FRAME;
+
 /* Special kernel registers structure definition */
 typedef struct _KSPECIAL_REGISTERS
 {
diff --git a/sdk/xtdk/ketypes.h b/sdk/xtdk/ketypes.h
index aa60161..0d1393c 100644
--- a/sdk/xtdk/ketypes.h
+++ b/sdk/xtdk/ketypes.h
@@ -388,6 +388,7 @@ typedef struct _KTHREAD
         SINGLE_LIST_ENTRY SwapListEntry;
     };
     PKQUEUE Queue;
+    CHAR PreviousMode;
     SHORT SpecialApcDisable;
     PTHREAD_ENVIRONMENT_BLOCK EnvironmentBlock;
     union
diff --git a/xtoskrnl/ke/amd64/kthread.c b/xtoskrnl/ke/amd64/kthread.c
index 1743739..c983ee4 100644
--- a/xtoskrnl/ke/amd64/kthread.c
+++ b/xtoskrnl/ke/amd64/kthread.c
@@ -39,5 +39,83 @@ KepInitializeThreadContext(IN PKTHREAD Thread,
                            IN PVOID StartContext,
                            IN PCONTEXT ContextRecord)
 {
-    UNIMPLEMENTED;
+    PKTHREAD_INIT_FRAME ThreadFrame;
+
+    /* Set initial thread frame */
+    ThreadFrame = (PKTHREAD_INIT_FRAME)Thread->InitialStack - sizeof(KTHREAD_INIT_FRAME);
+
+    /* Fill floating point save area with zeroes */
+    RtlZeroMemory(&ThreadFrame->NpxFrame, sizeof(FLOATING_SAVE_AREA));
+
+    /* Check if context provided for this thread */
+    if(ContextRecord)
+    {
+        /* Fill exception and trap frames with zeroes */
+        RtlZeroMemory(&ThreadFrame->ExceptionFrame, sizeof(KEXCEPTION_FRAME));
+        RtlZeroMemory(&ThreadFrame->TrapFrame, sizeof(KTRAP_FRAME));
+
+        /* Disable debug registers and enable context registers */
+        ContextRecord->ContextFlags &= ~CONTEXT_DEBUG_REGISTERS | CONTEXT_CONTROL;
+
+        /* Align the stack and reserve space for 4 parameters and return value */
+        ContextRecord->Rsp = (ContextRecord->Rsp & ~15) - 40;
+
+        /* Set CS and SS segments for user mode */
+        ContextRecord->SegCs = KGDT_R3_CODE | RPL_MASK;
+        ContextRecord->SegSs = KGDT_R3_DATA | RPL_MASK;
+
+        /* This is user mode thread */
+        Thread->PreviousMode = UserMode;
+
+        /* Enable floating point state */
+        Thread->NpxState = 1;
+
+        /* Set initial floating point state */
+        ThreadFrame->NpxFrame.ControlWord = 0x27F;
+        ThreadFrame->NpxFrame.StatusWord = 0;
+        ThreadFrame->NpxFrame.TagWord = 0xFFFF;
+        ThreadFrame->NpxFrame.ErrorOffset = 0;
+        ThreadFrame->NpxFrame.ErrorSelector = 0;
+        ThreadFrame->NpxFrame.ErrorOpcode = 0;
+        ThreadFrame->NpxFrame.DataOffset = 0;
+        ThreadFrame->NpxFrame.DataSelector = 0;
+
+        /* Clear DR6 and DR7 registers */
+        ThreadFrame->TrapFrame.Dr6 = 0;
+        ThreadFrame->TrapFrame.Dr7 = 0;
+
+        /* Set initial MXCSR register value */
+        ThreadFrame->TrapFrame.MxCsr = INITIAL_MXCSR;
+
+        /* Initialize exception frame */
+        ThreadFrame->ExceptionFrame.P1Home = (ULONG64)StartContext;
+        ThreadFrame->ExceptionFrame.P2Home = (ULONG64)StartRoutine;
+        ThreadFrame->ExceptionFrame.P3Home = (ULONG64)SystemRoutine;
+        ThreadFrame->ExceptionFrame.P4Home = (ULONG64)SystemRoutine;
+    }
+    else
+    {
+        /* This is kernel mode thread */
+        Thread->PreviousMode = KernelMode;
+
+        /* Disable floating point state */
+        Thread->NpxState = 0;
+
+        /* Set thread startup frame return information */
+        ThreadFrame->StartFrame.Return = (ULONG64)NULL;
+    }
+
+    /* Initialize thread startup information */
+    ThreadFrame->StartFrame.P1Home = (ULONG64)StartContext;
+    ThreadFrame->StartFrame.P2Home = (ULONG64)StartRoutine;
+    ThreadFrame->StartFrame.P3Home = (ULONG64)SystemRoutine;
+    ThreadFrame->StartFrame.P4Home = (ULONG64)SystemRoutine;
+
+    /* Initialize switch frame */
+    ThreadFrame->SwitchFrame.Rbp = (ULONG64)&ThreadFrame->TrapFrame + 128;
+    ThreadFrame->SwitchFrame.MxCsr = INITIAL_MXCSR;
+
+    /* Set thread stack */
+    Thread->InitialStack = &ThreadFrame->NpxFrame;
+    Thread->KernelStack = &ThreadFrame->SwitchFrame;
 }