From 06e755e36ad1077cee1c60fe2f2ffd5f2dfc8305 Mon Sep 17 00:00:00 2001
From: Matheus <matheusuky@gmail.com>
Date: Thu, 27 Apr 2023 23:32:25 -0300
Subject: [PATCH] context: misc useful macros and functions

---
 kernel-src/arch/x86-64/context.asm       | 81 ++++++++++++++++++++++++
 kernel-src/include/x86-64/arch/context.h | 45 ++++++++++++-
 2 files changed, 124 insertions(+), 2 deletions(-)
 create mode 100644 kernel-src/arch/x86-64/context.asm

diff --git a/kernel-src/arch/x86-64/context.asm b/kernel-src/arch/x86-64/context.asm
new file mode 100644
index 00000000..b0985bac
--- /dev/null
+++ b/kernel-src/arch/x86-64/context.asm
@@ -0,0 +1,81 @@
+section .text
+global arch_context_switch
+arch_context_switch:
+	; rdi is the pointer to the context struct
+	mov rsp, rdi
+
+	add rsp, 24 ; cr2 gs and fs are not popped.
+	pop rax
+	mov es,rax
+	pop rax
+	mov ds,rax
+	pop rax
+	pop rbx
+	pop rcx
+	pop rdx
+	pop r8
+	pop r9
+	pop r10
+	pop r11
+	pop r12
+	pop r13
+	pop r14
+	pop r15
+	pop rdi
+	pop rsi
+	pop rbp
+	add rsp,8 ; remove error code
+
+	; check if swapgs is needed
+	cmp qword [rsp+8], 0x18
+	jne .notneeded2
+	swapgs
+	.notneeded2:
+	o64 iret
+
+; saves context on the stack and passes as the argument to a function
+; since this uses the C system V abi, some registers can be clobbered (in this case, rsi and r11)
+; rdi has the pointer to the function to call
+; if the main function returns, just return normally
+global arch_context_saveandcall
+arch_context_saveandcall:
+	mov r11, [rsp] 	; save return address on scratch register
+	mov rsi, rsp
+	add rsi, 8 	; stack pointer before the call instruction
+
+	push qword 0x10 ; SS
+	push rsi 	; RSP
+	pushf 		; movs and pushes don't affect flags
+	push qword 0x08 ; CS
+	push r11 	; return address
+	sub rsp, 8 	; error code
+	push rbp
+	push rsi
+	push rdi
+	push r15
+	push r14
+	push r13
+	push r12
+	push r11
+	push r10
+	push r9
+	push r8
+	push rdx
+	push rcx
+	push rbx
+	push rax
+	push qword 0x10 ; ES
+	push qword 0x10 ; DS
+	push qword 0    ; FS
+	push qword 0	; GS
+	sub rsp, 8 	; cr2
+
+
+	mov r11, rdi 	; save function to call
+	mov rdi, rsp 	; first argument is the context struct
+	push rsi     	; save old stack
+	call r11 	; jump to the desired function
+
+	pop rsp		; restore old stack
+	sub rsp, 8
+	ret
diff --git a/kernel-src/include/x86-64/arch/context.h b/kernel-src/include/x86-64/arch/context.h
index 7f6ba285..e4526a40 100644
--- a/kernel-src/include/x86-64/arch/context.h
+++ b/kernel-src/include/x86-64/arch/context.h
@@ -4,6 +4,10 @@
 #include <stdint.h>
 #include <printf.h>
 
+#include <arch/msr.h>
+
+typedef uint64_t ctxreg_t;
+
 typedef struct {
         uint64_t cr2;
 	uint64_t gs;
@@ -33,11 +37,48 @@ typedef struct {
 	uint64_t ss;
 } __attribute__((packed)) context_t;
 
-#define CTX_SP(x) x->rsp;
-#define CTX_IP(x) x->rip;
+typedef struct {
+	uint64_t gsbase;
+	uint64_t fsbase;
+} extracontext_t;
+
+#define CTX_INIT(x,u) \
+	if (u) { \
+		(x)->cs = 0x18; \
+		(x)->ds = (x)->es = (x)->ss = 0x20; \
+	} else { \
+		(x)->cs = 0x8; \
+		(x)->ds = (x)->es = (x)->ss = 0x10; \
+	} \
+	(x)->rflags = 0x200;
+
+#define CTX_SP(x) (x)->rsp
+#define CTX_IP(x) (x)->rip
 #define PRINT_CTX(x) { \
 	printf("cr2: %016lx  gs: %04x\nrax: %016lx  fs: %04x\nrbx: %016lx  es: %04x\nrcx: %016lx  ds: %04x\nrdx: %016lx  cs: %04x\n r8: %016lx  ss: %04x\n r9: %016lx r10: %016lx\nr11: %016lx r12: %016lx\nr13: %016lx r14: %016lx\nr15: %016lx rdi: %016lx\nrsi: %016lx rbp: %016lx\nrip: %016lx rsp: %016lx\nerr: %016lx rfl: %016lx\n", \
 	x->cr2, x->gs, x->rax, x->fs, x->rbx, x->es, x->rcx, x->ds, x->rdx, x->cs, x->r8, x->ss, x->r9, x->r10, x->r11, x->r12, x->r13, x->r14, x->r15, x->rdi, x->rsi, x->rbp, x->rip, x->rsp, x->error, x->rflags); \
 }
 
+void arch_context_switch(context_t *context);
+void arch_context_saveandcall(void (*fn)(context_t *context));
+
+#include <string.h>
+
+// kernelgsbase is set as user because they'll be swapped in the context switch
+// TODO FPU stuff on all context macros
+#define ARCH_CONTEXT_SWITCHTHREAD(x) \
+	_cpu()->ist.rsp0 = (uint64_t)(x)->kernelstacktop; \
+	wrmsr(MSR_KERNELGSBASE, (x)->extracontext.gsbase); \
+	wrmsr(MSR_FSBASE, (x)->extracontext.fsbase); \
+	arch_context_switch(&thread->context);
+
+#define ARCH_CONTEXT_THREADSAVE(t, c) \
+	memcpy(&(t)->context, c, sizeof(context_t));
+
+#define ARCH_CONTEXT_THREADLOAD(t, c) \
+	memcpy(c, &(t)->context, sizeof(context_t)); \
+	_cpu()->ist.rsp0 = (uint64_t)(t)->kernelstacktop; \
+	wrmsr(MSR_KERNELGSBASE, (t)->extracontext.gsbase); \
+	wrmsr(MSR_FSBASE, (t)->extracontext.fsbase);
+
 #endif