Fix stack discipline for recursive traps

6 files changed, 25 insertions, 19 deletions

diff --git a/pk/entry.S b/pk/entry.S
index cdf076f..3d4a587 100644
--- a/pk/entry.S
+++ b/pk/entry.S
@@ -37,13 +37,12 @@
   STORE  x31,31*REGBYTES(x2)
 
   # get sr, epc, badvaddr, cause
-  addi   t0,sp,320
-  csrrw  t0,sscratch,t0
-  csrr   t1,sstatus
+  csrrw  t0,sscratch,x0
+  csrr   s0,sstatus
   csrr   t2,sepc
   csrr   t3,scause
   STORE  t0,2*REGBYTES(x2)
-  STORE  t1,32*REGBYTES(x2)
+  STORE  s0,32*REGBYTES(x2)
   STORE  t2,33*REGBYTES(x2)
   STORE  t3,35*REGBYTES(x2)
 
@@ -57,15 +56,22 @@
   .global  trap_entry
 trap_entry:
   csrrw sp, sscratch, sp
-
+  bnez sp, 1f
+  csrr sp, sscratch
 1:addi sp,sp,-320
   save_tf
   move  a0,sp
-  j   handle_trap
+  jal handle_trap
+
+  mv a0,sp
+  # don't restore sstatus if trap came from kernel
+  andi s0,s0,SSTATUS_PS
+  bnez s0,start_user
+  addi sp,sp,320
+  csrw sscratch,sp
   
-  .globl  pop_tf
-pop_tf:  # write the trap frame onto the stack
-  # restore sstatus and epc
+  .globl start_user
+start_user:
   csrc sstatus, SSTATUS_IE
   li t0, SSTATUS_PS
   LOAD t1, 32*REGBYTES(a0)
diff --git a/pk/handlers.c b/pk/handlers.c
index 8678327..3ceac1e 100644
--- a/pk/handlers.c
+++ b/pk/handlers.c
@@ -81,8 +81,6 @@ static void handle_syscall(trapframe_t* tf)
 static void handle_interrupt(trapframe_t* tf)
 {
   clear_csr(sstatus, SSTATUS_SIP);
-
-  pop_tf(tf);
 }
 
 void handle_trap(trapframe_t* tf)
@@ -107,6 +105,4 @@ void handle_trap(trapframe_t* tf)
   kassert(tf->cause < ARRAY_SIZE(trap_handlers) && trap_handlers[tf->cause]);
 
   trap_handlers[tf->cause](tf);
-
-  pop_tf(tf);
 }
diff --git a/pk/init.c b/pk/init.c
index 3058d43..fd9c819 100644
--- a/pk/init.c
+++ b/pk/init.c
@@ -98,7 +98,7 @@ uintptr_t boot_loader(struct mainvars* args)
     __builtin_unreachable();
   }
 
-  pk_vm_init();
+  uintptr_t kernel_stack_top = pk_vm_init();
   asm volatile("la t0, 1f; csrw mepc, t0; eret; 1:" ::: "t0");
 
   // copy phdrs to user stack
@@ -179,5 +179,6 @@ uintptr_t boot_loader(struct mainvars* args)
   trapframe_t tf;
   init_tf(&tf, current.entry, stack_top, current.elf64);
   __clear_cache(0, 0);
-  pop_tf(&tf);
+  write_csr(sscratch, kernel_stack_top);
+  start_user(&tf);
 }
diff --git a/pk/pk.h b/pk/pk.h
index d21a890..2df80e4 100644
--- a/pk/pk.h
+++ b/pk/pk.h
@@ -50,7 +50,7 @@ struct mainvars* parse_args(struct mainvars*);
 void printk(const char* s, ...);
 int snprintf(char* out, size_t n, const char* s, ...);
 void init_tf(trapframe_t*, long pc, long sp, int user64);
-void pop_tf(trapframe_t*) __attribute__((noreturn));
+void start_user(trapframe_t* tf) __attribute__((noreturn));
 void dump_tf(trapframe_t*);
 
 void unhandled_trap(trapframe_t*);
diff --git a/pk/vm.c b/pk/vm.c
index 64d7956..0719a4b 100644
--- a/pk/vm.c
+++ b/pk/vm.c
@@ -483,17 +483,20 @@ void supervisor_vm_init()
   flush_tlb();
 }
 
-void pk_vm_init()
+uintptr_t pk_vm_init()
 {
   __map_kernel_range(0, 0, current.first_free_paddr, PROT_READ|PROT_WRITE|PROT_EXEC);
   __map_kernel_range(first_free_page, first_free_page, free_pages * RISCV_PGSIZE, PROT_READ|PROT_WRITE);
 
   extern char trap_entry;
   write_csr(stvec, &trap_entry);
-  write_csr(sscratch, __page_alloc() + RISCV_PGSIZE);
+  write_csr(sscratch, 0);
 
   size_t stack_size = RISCV_PGSIZE * CLAMP(mem_size/(RISCV_PGSIZE*32), 1, 256);
   current.stack_bottom = __do_mmap(current.mmap_max - stack_size, stack_size, PROT_READ|PROT_WRITE|PROT_EXEC, MAP_PRIVATE|MAP_ANONYMOUS|MAP_FIXED, 0, 0);
   current.stack_top = current.stack_bottom + stack_size;
   kassert(current.stack_bottom != (uintptr_t)-1);
+
+  uintptr_t kernel_stack_top = __page_alloc() + RISCV_PGSIZE;
+  return kernel_stack_top;
 }
diff --git a/pk/vm.h b/pk/vm.h
index 6031af0..06e0e71 100644
--- a/pk/vm.h
+++ b/pk/vm.h
@@ -34,7 +34,7 @@
 
 void vm_init();
 void supervisor_vm_init();
-void pk_vm_init();
+uintptr_t pk_vm_init();
 int handle_page_fault(uintptr_t vaddr, int prot);
 void populate_mapping(const void* start, size_t size, int prot);
 void __map_kernel_range(uintptr_t va, uintptr_t pa, size_t len, int prot);