From 29777c1ada7616d6b5996f6d804b2731e7c5ac55 Mon Sep 17 00:00:00 2001
From: Loic Guegan <manzerbredes@mailbox.org>
Date: Sat, 17 Apr 2021 19:35:56 +0200
Subject: [PATCH] Allow syscall to be interrupted!

---
 src/core/idt.c           |  2 +-
 src/core/scheduler.c     | 35 +++++++++++++++++++++++------------
 src/core/scheduler.h     |  1 +
 src/core/scheduler_asm.S | 28 ++++++++++++++++++++++++----
 4 files changed, 49 insertions(+), 17 deletions(-)

diff --git a/src/core/idt.c b/src/core/idt.c
index 861927a..30bb113 100644
--- a/src/core/idt.c
+++ b/src/core/idt.c
@@ -26,7 +26,7 @@ void idt_init(){
         if(i==33)
             idt_write_entry((IDT_ENTRY){0x08,(u32)&INT_KEYPRESS,IDT_INT_GATE|IDT_P},i);
         if(i==48)
-            idt_write_entry((IDT_ENTRY){0x08,(u32)&INT_SYSCALL,IDT_INT_GATE|IDT_P|IDT_PRVL_3},i);
+            idt_write_entry((IDT_ENTRY){0x08,(u32)&INT_SYSCALL,IDT_TRAP_GATE|IDT_P|IDT_PRVL_3},i);
     }
     // Load IDT
     asm("lidtl (IDTR)");
diff --git a/src/core/scheduler.c b/src/core/scheduler.c
index 59f3b9d..0c1d255 100644
--- a/src/core/scheduler.c
+++ b/src/core/scheduler.c
@@ -37,8 +37,19 @@ void schedule(u32 *stack){
   p->regs.eip=stack[12];
   p->regs.cs=stack[13];
   p->regs.eflags=stack[14];
-  p->regs.esp=stack[15];
-  p->regs.ss=stack[16];
+
+  // If clock occurs during a syscall (another interrupt)
+  if(p->regs.cs==0x08){
+    // Since we where already in kernel mode
+    // the CPU did not push ss and esp
+    // thus esp point to 
+    p->regs.esp=&stack[15]; // Since we came from another interrupt this value will be restore by the last interrupt
+    p->regs.ss=TSS.ss0;
+  }
+  else {
+    p->regs.esp=stack[15];
+    p->regs.ss=stack[16];
+  }
 
   // Get the next task to run
   current_id++;
@@ -46,9 +57,9 @@ void schedule(u32 *stack){
     current_id=0;
   p=&procs[current_id];
  
-  // Have a clean stack on next interrupt
-  TSS.esp0=(u32)stack+17;
-  asm("mov %%ss, %0": "=m" (TSS.ss0));
+  // Use kernel stack of the next task
+  TSS.ss0=p->regs.ss0;
+  TSS.esp0=p->regs.esp0;
 
   // Ensure interrupts are activated and NT flag is clear
   p->regs.eflags|=0x200;
@@ -90,6 +101,8 @@ void task_create(void *task, int task_size){
     void *entry_point=PAGING_ENTRY_POINT_VIRT;
     // User stack start at the end of last allocated page (since addresses are going down)
     void *ustack=(void*)(PAGING_ENTRY_POINT_VIRT+allocated*4096);
+    // Kernel stack start at the end of last allocated page +1 (see paging_allocate() implementation)
+    void *kstack=(void*)(PAGING_ENTRY_POINT_VIRT+(allocated+1)*4096);
 
     // Load the task into memory
     memcpy(task,PAGING_ENTRY_POINT_PHY(page_dir), task_size);
@@ -105,6 +118,8 @@ void task_create(void *task, int task_size){
     p->regs.ebx=0;
     p->regs.ecx=0;
     p->regs.edx=0;
+    p->regs.ss0=0x18;
+    p->regs.esp0=kstack;
 
     // Manage eflags
     u32 eflags;
@@ -135,19 +150,15 @@ void scheduler_start(){
     // Enable scheduling
     scheduler_on=1;
 
-    // Save kernel stack state
-    u32 *stack;
-    asm("mov %%ebp, %0":"=r" (stack));
-    // Remove ebp from the (c call convention) and return address (call)
-    TSS.esp0=(u32)stack+2;
-    asm("mov %%ss, %0": "=m" (TSS.ss0));
-
     // Get first stack
     current_id=0;
     PROC *p=&procs[current_id];
     // Ensure interrupts are activated and NT flag is clear
     p->regs.eflags|=0x200;
     p->regs.eflags&=0xffffbfff;
+    // Use task kernel stack
+    TSS.ss0=p->regs.ss0;
+    TSS.esp0=p->regs.esp0;
     // Switch to user task
     asm(
       "push %0    \n\t"
diff --git a/src/core/scheduler.h b/src/core/scheduler.h
index 3410ba1..9678470 100644
--- a/src/core/scheduler.h
+++ b/src/core/scheduler.h
@@ -14,6 +14,7 @@ typedef struct REGISTERS {
     u32 esi, edi;
     u32 ds, es, fs, gs;
     u32 eflags;
+    u32 ss0, esp0;
 } __attribute__((packed)) REGISTERS;
 
 // If you change the following struct
diff --git a/src/core/scheduler_asm.S b/src/core/scheduler_asm.S
index 1f916ec..2fdf3e5 100644
--- a/src/core/scheduler_asm.S
+++ b/src/core/scheduler_asm.S
@@ -5,6 +5,24 @@
 // to this process
 task_switch:
     pop %esi
+
+    // Switch to the new task kernel stack
+    mov 72(%esi), %ss
+    mov 76(%esi), %esp
+
+    // If we where in kernel mode we should keep using the same
+    // stack which is stored in regs.ss and regs.esp by schedule
+    mov 24(%esi), %eax
+    cmp $0x08, %eax # Check if we were in kernel mode
+    jne skip_kernel_stack
+    mov 32(%esi), %ss
+    mov 36(%esi), %esp
+    skip_kernel_stack:
+
+    // Setup process page directory to be able to use task stack
+    mov 4(%esi), %eax
+    mov %eax, %cr3
+
     push 8(%esi)  # eax (cf PROC struct)
     push 12(%esi)  # ebx
     push 16(%esi)  # ecx
@@ -20,10 +38,6 @@ task_switch:
     movb $0x20, %al
     outb %al, $0x20
 
-    // Setup process page directory
-    mov 4(%esi), %eax
-    mov %eax, %cr3
-
     // Setup registers
     pop %gs
     pop %fs
@@ -37,8 +51,14 @@ task_switch:
     pop %eax
 
     // Perform the task switch
+    push %eax
+    mov 24(%esi), %eax
+    cmp $0x08, %eax # Check if we were in kernel mode
+    pop %eax
+    je skip_stack
     push 32(%esi)  # ss
     push 36(%esi)  # esp
+    skip_stack:
     push 68(%esi)  # eflags
     push 24(%esi)  # cs
     push 28(%esi)  # eip