6 files changed, 81 insertions, 38 deletions

diff --git a/kc/core/cpu/cpu.c b/kc/core/cpu/cpu.c
index a979336..e3cff39 100644
--- a/kc/core/cpu/cpu.c
+++ b/kc/core/cpu/cpu.c
@@ -113,48 +113,66 @@ void set_ist(int vec, int ist_index)
     idt[vec].access |= ist_index & 0x7;
 }
 
-static void gdt_init(void)
-{   
-    set_gdt(CODE_SUPER_SEG_INDEX, 0, (uint64_t)-1, CODE64_SUPER_SEG);
-    set_gdt(CODE_USER_SEG_INDEX, 0, (uint64_t)-1, CODE64_USER_SEG);
-    set_gdt(DATA_SEG_INDEX, 0, (uint64_t)-1, DATA_SEG);
-    set_gdt(TASK_SEG_INDEX, (uint64_t)&tss, sizeof(tss) - 1, TASK64_SEG);
-
-    struct dtr64 gdtr = {sizeof(gdt) - 1, (uint64_t)&gdt};
-
+static void gdt_flush(uint16_t code_seg, uint16_t data_seg, uint16_t task_seg)
+{
     __asm__ volatile
-        (   
-         // can't have interrupts during GDT reload
-         // but don't re-enable interrupts if they were already disabled
+        (
          "pushf\n"
          "cli\n"
-         "lgdt %0\n"
-         "mov %w2, %%ds\n"
-         "mov %w2, %%es\n"
-         "mov %w2, %%fs\n"
-         "mov %w2, %%gs\n"
-         "mov %w2, %%ss\n"
+         "mov %w1, %%ds\n"
+         "mov %w1, %%es\n"
+         "mov %w1, %%fs\n"
+         "mov %w1, %%gs\n"
+         "mov %w1, %%ss\n"
          // make a far return frame on the stack
          // [ss] @rsp+8
          // [rip] @rsp
-         "pushq %1\n"
+         "push %q0\n"
          "lea .Lflush(%%rip), %%rax\n"
          "push %%rax\n"
          // far return, now we're in the code segment
          // defined here.
          "lretq\n"
          ".Lflush:\n"
-         "ltr %w3\n"
+         "ltr %w2\n"
+         "popf\n"
+         :
+         :
+            "r"(code_seg << 3),
+            "r"(data_seg << 3),
+            "r"(task_seg << 3)
+        );
+}
+
+static void gdt_load(struct dtr64 gdtr)
+{
+    __asm__ volatile
+        (   
+         // can't have interrupts during GDT reload
+         // but don't re-enable interrupts if they were already disabled
+         "pushf\n"
+         "cli\n"
+         "lgdt %0\n"
          "popf\n"
          // dependency on pre-boot GDT is now gone
          :
-         : "m"(gdtr),
-        "i"(CODE_SUPER_SEG_INDEX << 3),
-        "r"(DATA_SEG_INDEX << 3),
-        "r"(TASK_SEG_INDEX << 3)
-            );
+         : "m"(gdtr)
+        );
+}
+
+static void gdt_init(void)
+{   
+    set_gdt(CODE_SUPER_SEG_INDEX, 0, (uint64_t)-1, CODE64_SUPER_SEG);
+    set_gdt(CODE_USER_SEG_INDEX, 0, (uint64_t)-1, CODE64_USER_SEG);
+    set_gdt(DATA_SEG_INDEX, 0, (uint64_t)-1, DATA_SEG);
+    set_gdt(TASK_SEG_INDEX, (uint64_t)&tss, sizeof(tss) - 1, TASK64_SEG);
+
+    struct dtr64 gdtr = {sizeof(gdt) - 1, (uint64_t)&gdt};
+    gdt_load(gdtr);
+    gdt_flush(CODE_SUPER_SEG_INDEX, DATA_SEG_INDEX, TASK_SEG_INDEX);
 }
 
+
 static void idt_init(void)
 {
     struct dtr64 idtr = {sizeof(idt) - 1, (uint64_t)&idt};
diff --git a/kc/core/cpu/cpu_task.S b/kc/core/cpu/cpu_task.S
index 7216d0e..a411028 100644
--- a/kc/core/cpu/cpu_task.S
+++ b/kc/core/cpu/cpu_task.S
@@ -15,7 +15,7 @@ cpu_set_thread:
     *  the following are the parameters for switching tasks
     *   rdi: pointer to control block of thread being switched to 
     *   rsi: pointer to control block of the current thread
-    *   rdx: pointer to this cpu's task state segment esp0 field
+    *   rdx: pointer to this cpu's task state segment rsp0 field
     */
     // test for attempt to switch to same thread
     cmp %rdi,%rsi
diff --git a/kc/core/cpu/irq.c b/kc/core/cpu/irq.c
index e4dd29f..f42db4c 100644
--- a/kc/core/cpu/irq.c
+++ b/kc/core/cpu/irq.c
@@ -4,13 +4,25 @@
 
 static volatile atomic_uint_fast64_t irq_lock_count = 0;
 
-void irq_lock(void)
+uint64_t irq_lock(void)
 {
+    uint64_t flags;
+
+    __asm__ volatile 
+        ( 
+         "pushf\n"
+         "pop %0\n"
+         "cli\n"
+         : 
+         "=r"(flags)
+        );
+
     irq_lock_count++;
-    __asm__ volatile ("cli;");
+
+    return flags;
 }
 
-void irq_unlock(void)
+void irq_unlock(uint64_t flags)
 {
     if (atomic_load(&irq_lock_count) >= 1)
     {
@@ -19,7 +31,14 @@ void irq_unlock(void)
 
     if (atomic_load(&irq_lock_count) == 0)
     {
-        __asm__ volatile ("sti;");
+        __asm__ volatile
+            (
+             "push %q0\n"
+             "popf"
+             :
+             :
+             "r"(flags)
+            );
     }
 }
 
diff --git a/kc/core/cpu/irq.h b/kc/core/cpu/irq.h
index 529f385..9ee3a98 100644
--- a/kc/core/cpu/irq.h
+++ b/kc/core/cpu/irq.h
@@ -4,8 +4,10 @@
 
 #ifndef __ASSEMBLER__
 
-void irq_lock(void);
-void irq_unlock(void);
+#include <stdint.h>
+
+uint64_t irq_lock(void);
+void irq_unlock(uint64_t flags);
 
 #else
 
diff --git a/kc/core/cpu/mmu.c b/kc/core/cpu/mmu.c
index 9dd2e41..c0b294b 100644
--- a/kc/core/cpu/mmu.c
+++ b/kc/core/cpu/mmu.c
@@ -5,9 +5,9 @@
 
 void mmu_set_map(phys_addr_t map)
 {
-    irq_lock();
+    uint64_t flags = irq_lock();
     __asm__ volatile ("mov %0, %%cr3" :: "r"(page_address(map, 1)));
-    irq_unlock();
+    irq_unlock(flags);
 }
 
 phys_addr_t mmu_get_map(void)
diff --git a/kc/core/task.c b/kc/core/task.c
index 6c73f89..3e0288d 100644
--- a/kc/core/task.c
+++ b/kc/core/task.c
@@ -61,6 +61,8 @@ static void idle_thread_entry(void)
 {
     kprintf("idle thread started\n");
     unlock_scheduler();
+    
+    __asm__ ("sti");
 
     while (true)
     {
@@ -171,20 +173,22 @@ void task_schedule(void)
     }
 }
 
+static uint64_t scheduler_flags;
+
 static void lock_scheduler(void)
 {
-    irq_lock();
+    scheduler_flags = irq_lock();
 }
 
 static void lock_preempt(void)
 {
-    irq_lock();
+    scheduler_flags = irq_lock();
     preempt_switch_count++;
 }
 
 static void unlock_scheduler(void)
 {
-    irq_unlock();
+    irq_unlock(scheduler_flags);
 }
 
 static void unlock_preempt(void)
@@ -201,7 +205,7 @@ static void unlock_preempt(void)
         task_schedule();
     }
 
-    irq_unlock();
+    irq_unlock(scheduler_flags);
 }
 
 void update_time(void)