ptos/amd64: Add task state segment

Signed-off-by: Chloe M <chloe@mensia.org>

service/ptos/arch/amd64/cpu/init.c

@@ -12,6 +12,10 @@
 #include <machine/trap.h>
 #include <machine/msr.h>
 #include <machine/lapic.h>
+#include <machine/gdt.h>
+
+/* Externs */
+extern GDT_ENTRY g_GDT[GDT_ENTRY_COUNT];
 
 /*
  * Initialize interrupts for the current processor
@@ -66,9 +70,22 @@ HalKpcrP1Init(KPCR *Kpcr)
 VOID
 HalKpcrP2Init(KPCR *Kpcr)
 {
+    TSS_DESCRIPTOR *TssDesc;
+
     /* Initialize per-processor pools */
     ExPoolRegionInit(&Kpcr->PoolRegion);
 
     /* Initialize the Local APIC unit */
     MdLapicInit(Kpcr);
+
+    /* Load the Task State Segment */
+    TssDesc = (TSS_DESCRIPTOR *)&g_GDT[GDT_TSS_INDEX];
+    MdTssInit(TssDesc);
+    ASMV(
+        "mov %0, %%ax\n"
+        "ltr %%ax\n"
+        :
+        : "i" (GDT_TSS)
+        : "memory"
+    );
 }

service/ptos/arch/amd64/cpu/tss.c

@@ -0,0 +1,139 @@
+/*
+ * Copyright (c) 2026, Chloe M.
+ * Provided under the BSD-3 clause.
+ *
+ * Description: Task state segment
+ * Author:      Chloe M.
+ */
+
+#include <machine/tss.h>
+#include <hal/kpcr.h>
+#include <ex/pool.h>
+#include <ke/bugcheck.h>
+
+#define TSS_POOL_TAG 'TSS'
+
+VOID
+MdTssAllocIst(UCHAR IstIndex, USIZE StackSize)
+{
+    KPCR *Kpcr;
+    MCB *Mcb;
+    TSS_ENTRY *TssEntry;
+    UPTR StackBase, *Ptr;
+    UPTR StackTop;
+    ULONG BaseLow, BaseHigh;
+
+    if (StackSize == 0 || IstIndex > 7) {
+        return;
+    }
+
+    Kpcr = HalKpcrCurrent();
+    Mcb = &Kpcr->Mcb;
+    TssEntry = Mcb->Tss;
+
+    /* Allocate the stack */
+    Ptr = ExAllocatePoolWithTag(
+        POOL_NON_PAGED,
+        StackSize,
+        TSS_POOL_TAG
+    );
+
+    if (Ptr == NULL) {
+        KeBugCheck(
+            BUGCHECK_OOM,
+            "failed to allocate interrupt stack %d\n",
+            IstIndex
+        );
+    }
+
+    /* Obtain the stack base */
+    StackBase = (UPTR)Ptr;
+    StackTop = StackBase + StackSize;
+    BaseLow = StackTop & 0xFFFFFFFF;
+    BaseHigh = (StackTop >> 32) & 0xFFFFFFFF;
+
+    /* Set the target index */
+    switch (IstIndex) {
+    case 0:
+        KeBugCheck(
+            BUGCHECK_MISC,
+            "cannot set IST index 0\n"
+        );
+
+        break;
+    case 1:
+        TssEntry->Ist1Low = BaseLow;
+        TssEntry->Ist1High = BaseHigh;
+        break;
+    case 2:
+        TssEntry->Ist2Low = BaseLow;
+        TssEntry->Ist2High = BaseHigh;
+        break;
+    case 3:
+        TssEntry->Ist3Low = BaseLow;
+        TssEntry->Ist3High = BaseHigh;
+        break;
+    case 4:
+        TssEntry->Ist4Low = BaseLow;
+        TssEntry->Ist4High = BaseHigh;
+        break;
+    case 5:
+        TssEntry->Ist5Low = BaseLow;
+        TssEntry->Ist5High = BaseHigh;
+        break;
+    case 6:
+        TssEntry->Ist6Low = BaseLow;
+        TssEntry->Ist6High = BaseHigh;
+        break;
+    case 7:
+        TssEntry->Ist7Low = BaseLow;
+        TssEntry->Ist7High = BaseHigh;
+        break;
+    }
+}
+
+VOID
+MdTssInit(TSS_DESCRIPTOR *Descriptor)
+{
+    TSS_ENTRY *TssEntry;
+    KPCR *KpcrCurrent;
+    MCB *Mcb;
+    UPTR TssBase;
+
+    if (Descriptor == NULL) {
+        KeBugCheck(
+            BUGCHECK_MISC,
+            "bad descriptor passed to MdTssInit()\n"
+        );
+    }
+
+    TssEntry = ExAllocatePoolWithTag(
+        POOL_NON_PAGED,
+        sizeof(*TssEntry),
+        TSS_POOL_TAG
+    );
+
+    if (TssEntry == NULL) {
+        KeBugCheck(
+            BUGCHECK_OOM,
+            "out of memory when allocating tss\n"
+        );
+    }
+
+    TssBase = (UPTR)TssEntry;
+    Descriptor->SegmentLimit = sizeof(*TssEntry);
+    Descriptor->Present = 1;
+    Descriptor->Gran = 1;
+    Descriptor->Type = 0x9;
+    Descriptor->Avl = 0;
+    Descriptor->Dpl = 0;
+    Descriptor->BaseLow16 = TssBase & 0xFFFF;
+    Descriptor->BaseMid8 = (TssBase >> 16) & 0xFF;
+    Descriptor->BaseHighMid8 = (TssBase >> 24) & 0xFF;
+    Descriptor->BaseHigh32 = (TssBase >> 32) & 0xFFFFFFFF;
+    TssEntry->IoBitmap = 0xFF;
+
+    KpcrCurrent = HalKpcrCurrent();
+    Mcb = &KpcrCurrent->Mcb;
+    Mcb->Tss = TssEntry;
+}

service/ptos/head/arch/amd64/mcb.h

@@ -10,6 +10,7 @@
 #define _MACHINE_MCB_H_ 1
 
 #include <ptdef.h>
+#include <machine/tss.h>
 
 /*
  * The machine core block contains MD processor information
@@ -17,11 +18,13 @@
  * @LapicBase:      Local APIC MMIO base
  * @HasX2Apic:      Has an x2APIC unit
  * @LapicTmrFreq:   Local APIC timer frequency
+ * @Tss:            Task state segment
  */
 typedef struct {
     VOID *LapicBase;
     UCHAR HasX2Apic : 1;
     USIZE LapicTmrFreq;
+    TSS_ENTRY *Tss;
 } MCB;
 
 #endif  /* !_MACHINE_MCB_H_ */

service/ptos/head/arch/amd64/tss.h

@@ -0,0 +1,99 @@
+/*
+ * Copyright (c) 2026, Chloe M.
+ * Provided under the BSD-3 clause.
+ *
+ * Description: Task state segment
+ * Author:      Chloe M.
+ */
+
+#ifndef _MACHINE_TSS_H_
+#define _MACHINE_TSS_H_ 1
+
+#include <ptdef.h>
+
+/*
+ * Represents used interrupt stack table indices
+ *
+ * @IST_RESERVED: This is not used
+ * @IST_NMI:      Stack for NMIs
+ * @IST_SCHED:    Stack for scheduler
+ */
+typedef enum {
+    IST_RESERVED,
+    IST_NMI,
+    IST_SCHED
+} IST_INDEX;
+
+/*
+ * Represents a platform task state segment
+ *
+ * Refer to section 8.7 of the Intel SDM
+ */
+typedef struct PACKED {
+    ULONG Reserved;
+    ULONG Rsp0Low;
+    ULONG Rsp0High;
+    ULONG Rsp1Low;
+    ULONG Rsp1High;
+    ULONG Rsp2Low;
+    ULONG Rsp2High;
+    ULONG Reserved1;
+    ULONG Reserved2;
+    ULONG Ist1Low;
+    ULONG Ist1High;
+    ULONG Ist2Low;
+    ULONG Ist2High;
+    ULONG Ist3Low;
+    ULONG Ist3High;
+    ULONG Ist4Low;
+    ULONG Ist4High;
+    ULONG Ist5Low;
+    ULONG Ist5High;
+    ULONG Ist6Low;
+    ULONG Ist6High;
+    ULONG Ist7Low;
+    ULONG Ist7High;
+    ULONG Reserved3;
+    ULONG Reserved4;
+    USHORT Reserved5;
+    USHORT IoBitmap;
+} TSS_ENTRY;
+
+/*
+ * Task state segment descriptor within GDT
+ *
+ * Refer to section 8.2.3 of the Intel SDM
+ */
+typedef struct PACKED {
+    USHORT SegmentLimit;
+    USHORT BaseLow16;
+    UCHAR BaseMid8;
+    UCHAR Type : 4;
+    UCHAR Zero : 1;
+    UCHAR Dpl  : 2;
+    UCHAR Present : 1;
+    UCHAR SegLimitHigh : 4;
+    UCHAR Avl : 1;
+    UCHAR Unused : 2;
+    UCHAR Gran : 1;
+    UCHAR BaseHighMid8;
+    ULONG BaseHigh32;
+    ULONG Reserved;
+} TSS_DESCRIPTOR;
+
+/*
+ * Allocate an interrupt stack table entry
+ *
+ * @IstIndex:  Interrupt stack table index
+ * @StackSize: Size of stack to allocate
+ */
+VOID MdTssAllocIst(UCHAR IstIndex, USIZE StackSize);
+
+/*
+ * Initialize the task state segment
+ *
+ * @Descriptor:   Task state segment desciptor
+ */
+VOID MdTssInit(TSS_DESCRIPTOR *Descriptor);
+
+#endif  /* !_MACHINE_TSS_H_ */