scheduler: Fix page table switch for scheduler

The userspace process may map pages where the kernel has allocated data and cause a crash.

This patch fixes this issue by having a separate IRQ handler which sets the kernel page table at the start of SchedulerInterruptHandler() and restores it in SchedulerHandlerStub() function.

arch/amd64/cpu/idt.cpp

@@ -29,6 +29,7 @@
 #pragma GCC diagnostic ignored "-Wconversion"
 
 extern "C" void MainInterruptHandler(void *Data);
+extern "C" void SchedulerInterruptHandler(void *Data);
 extern "C" void ExceptionHandler(void *Data);
 
 #define __stub_handler \
@@ -247,6 +248,62 @@ namespace InterruptDescriptorTable
 			"iretq"); // pop CS RIP RFLAGS SS RSP
 	}
 
+	extern "C" __stub_handler void SchedulerHandlerStub()
+	{
+		asm("cld\n"
+			"cli\n"
+
+			"pushq %rax\n"
+			"pushq %rbx\n"
+			"pushq %rcx\n"
+			"pushq %rdx\n"
+			"pushq %rsi\n"
+			"pushq %rdi\n"
+			"pushq %rbp\n"
+
+			"pushq %r8\n"
+			"pushq %r9\n"
+			"pushq %r10\n"
+			"pushq %r11\n"
+			"pushq %r12\n"
+			"pushq %r13\n"
+			"pushq %r14\n"
+			"pushq %r15\n"
+
+			/* TODO: Add advanced check so we won't update the cr3 when not needed */
+
+			"movq %cr3, %rax\n pushq %rax\n" /* Push opt */
+			"pushq %rax\n" /* Push ppt */
+
+			"movq %rsp, %rdi\n"
+			"call SchedulerInterruptHandler\n"
+
+			"popq %rax\n movq %rax, %cr3\n" /* Restore to ppt */
+			"popq %rax\n" /* Pop opt */
+
+			"popq %r15\n"
+			"popq %r14\n"
+			"popq %r13\n"
+			"popq %r12\n"
+			"popq %r11\n"
+			"popq %r10\n"
+			"popq %r9\n"
+			"popq %r8\n"
+
+			"popq %rbp\n"
+			"popq %rdi\n"
+			"popq %rsi\n"
+			"popq %rdx\n"
+			"popq %rcx\n"
+			"popq %rbx\n"
+			"popq %rax\n"
+
+			"addq $16, %rsp\n"
+
+			"sti\n"
+			"iretq"); // pop CS RIP RFLAGS SS RSP
+	}
+
 #pragma region Interrupt Macros
 
 #define EXCEPTION_HANDLER(num)                          \
@@ -272,6 +329,14 @@ namespace InterruptDescriptorTable
 			"jmp InterruptHandlerStub\n");       \
 	}
 
+#define SCHEDULER_HANDLER(num)                   \
+	__stub_handler void InterruptHandler_##num() \
+	{                                            \
+		asm("pushq $0\n"                         \
+			"pushq $" #num "\n"                  \
+			"jmp SchedulerHandlerStub\n");       \
+	}
+
 	/* ISR */
 
 	EXCEPTION_HANDLER(0x0);
@@ -328,7 +393,7 @@ namespace InterruptDescriptorTable
 
 	/* Reserved by OS */
 
-	INTERRUPT_HANDLER(0x30)
+	SCHEDULER_HANDLER(0x30)
 	INTERRUPT_HANDLER(0x31)
 	INTERRUPT_HANDLER(0x32)
 	INTERRUPT_HANDLER(0x33)

core/interrupts_manager.cpp

@@ -280,7 +280,7 @@ namespace Interrupts
 		warn("IRQ%d not found.", InterruptNumber);
 	}
 
-	nsa inline void ReturnFromInterrupt(CPU::TrapFrame *)
+	nsa inline void ReturnFromInterrupt()
 	{
 		CPUData *CoreData = GetCurrentCPU();
 		int Core = CoreData->ID;
@@ -348,7 +348,7 @@ namespace Interrupts
 		if (it == RegisteredEvents.end())
 		{
 			warn("IRQ%d is not registered.", Frame->InterruptNumber - 32);
-			ReturnFromInterrupt(Frame);
+			ReturnFromInterrupt();
 			return;
 		}
 
@@ -367,7 +367,40 @@ namespace Interrupts
 				it->Callback(Frame);
 		}
 
-		ReturnFromInterrupt(Frame);
+		ReturnFromInterrupt();
+	}
+
+	extern "C" nsa void SchedulerInterruptHandler(void *Data)
+	{
+		KernelPageTable->Update();
+		CPU::SchedulerFrame *Frame = (CPU::SchedulerFrame *)Data;
+#if defined(a86)
+		assert(Frame->InterruptNumber == CPU::x86::IRQ16);
+#else
+		assert(Frame->InterruptNumber == 16);
+#endif
+
+		auto it = std::find_if(RegisteredEvents.begin(), RegisteredEvents.end(),
+							   [](const Event &ev)
+							   {
+								   return ev.IRQ == 16;
+							   });
+
+		if (it == RegisteredEvents.end())
+		{
+			warn("Scheduler interrupt is not registered.");
+			ReturnFromInterrupt();
+			Frame->ppt = Frame->opt;
+			debug("opt = %#lx", Frame->opt);
+			return;
+		}
+		assert(it->IsHandler);
+
+		it->Priority++;
+
+		Handler *hnd = (Handler *)it->Data;
+		hnd->OnInterruptReceived(Frame);
+		ReturnFromInterrupt();
 	}
 
 	Handler::Handler(int InterruptNumber, bool Critical)
@@ -426,7 +459,11 @@ namespace Interrupts
 
 	void Handler::OnInterruptReceived(CPU::TrapFrame *Frame)
 	{
-		trace("Unhandled interrupt %d",
-			  Frame->InterruptNumber);
+		debug("Unhandled interrupt %d", Frame->InterruptNumber);
+	}
+
+	void Handler::OnInterruptReceived(CPU::SchedulerFrame *Frame)
+	{
+		debug("Unhandled scheduler interrupt %d", Frame->InterruptNumber);
 	}
 }

include/cpu.hpp

@@ -641,6 +641,38 @@ namespace CPU
 			uint64_t ss;   /* Stack Segment */
 		};
 
+		struct SchedulerFrame
+		{
+			uint64_t ppt; /* Process Page Table */
+			uint64_t opt; /* Original Page Table */
+
+			uint64_t r15; /* General purpose */
+			uint64_t r14; /* General purpose */
+			uint64_t r13; /* General purpose */
+			uint64_t r12; /* General purpose */
+			uint64_t r11; /* General purpose */
+			uint64_t r10; /* General purpose */
+			uint64_t r9;  /* General purpose */
+			uint64_t r8;  /* General purpose */
+
+			uint64_t rbp; /* Base Pointer (meant for stack frames) */
+			uint64_t rdi; /* Destination index for string operations */
+			uint64_t rsi; /* Source index for string operations */
+			uint64_t rdx; /* Data (commonly extends the A register) */
+			uint64_t rcx; /* Counter */
+			uint64_t rbx; /* Base */
+			uint64_t rax; /* Accumulator */
+
+			uint64_t InterruptNumber; /* Interrupt Number */
+			uint64_t ErrorCode;		  /* Error code */
+
+			uint64_t rip;  /* Instruction Pointer */
+			uint64_t cs;   /* Code Segment */
+			RFLAGS rflags; /* Register Flags */
+			uint64_t rsp;  /* Stack Pointer */
+			uint64_t ss;   /* Stack Segment */
+		};
+
 		struct ExceptionFrame
 		{
 			uint64_t cr0; /* Control Register 0 (system control) */
@@ -1026,6 +1058,7 @@ namespace CPU
 	 * @note This is for x86_64
 	 */
 	typedef x64::TrapFrame TrapFrame;
+	typedef x64::SchedulerFrame SchedulerFrame;
 	typedef x64::ExceptionFrame ExceptionFrame;
 #elif defined(a32)
 	/**
@@ -1034,6 +1067,7 @@ namespace CPU
 	 * @note This is for x86_32
 	 */
 	typedef x32::TrapFrame TrapFrame;
+	typedef x32::SchedulerFrame SchedulerFrame;
 	typedef x32::ExceptionFrame ExceptionFrame;
 #elif defined(aa64)
 	/**
@@ -1042,6 +1076,7 @@ namespace CPU
 	 * @note This is for aarch64
 	 */
 	typedef aarch64::TrapFrame TrapFrame;
+	typedef aarch64::SchedulerFrame SchedulerFrame;
 	typedef aarch64::TrapFrame ExceptionFrame;
 #endif
 }

include/ints.hpp

@@ -78,6 +78,7 @@ namespace Interrupts
 
 	public:
 		virtual void OnInterruptReceived(CPU::TrapFrame *Frame);
+		virtual void OnInterruptReceived(CPU::SchedulerFrame *Frame);
 	};
 }
 

include/scheduler.hpp

@@ -142,8 +142,8 @@ namespace Tasking::Scheduler
 		void WakeUpThreads();
 		void CleanupTerminated();
 
-		void Schedule(CPU::TrapFrame *Frame);
-		void OnInterruptReceived(CPU::TrapFrame *Frame) final;
+		void Schedule(CPU::SchedulerFrame *Frame);
+		void OnInterruptReceived(CPU::SchedulerFrame *Frame) final;
 
 		Custom(Task *ctx);
 		virtual ~Custom();

include/signal.hpp

@@ -312,11 +312,11 @@ namespace Tasking
 		{
 #ifdef a64
 			CPU::x64::FXState fx;
-			CPU::x64::TrapFrame tf;
+			CPU::x64::SchedulerFrame tf;
 			uintptr_t GSBase, FSBase, ShadowGSBase;
 #else
 			CPU::x32::FXState fx;
-			CPU::x32::TrapFrame tf;
+			CPU::x32::SchedulerFrame tf;
 			uintptr_t GSBase, FSBase;
 #endif
 			sigset_t SignalMask;
@@ -422,7 +422,7 @@ namespace Tasking
 		int AddSignal(Signals sig, union sigval val = {0}, pid_t tid = -1);
 		int RemoveSignal(Signals sig);
 
-		bool HandleSignal(CPU::TrapFrame *tf, void *thread);
+		bool HandleSignal(CPU::SchedulerFrame *tf, void *thread);
 		void RestoreHandleSignal(SyscallsFrame *tf, void *thread);
 
 		int SetAction(Signals sig, const SignalAction *act);

include/task.hpp

@@ -356,10 +356,10 @@ namespace Tasking
 
 		/* CPU state */
 #if defined(a64)
-		CPU::x64::TrapFrame Registers{};
+		CPU::x64::SchedulerFrame Registers{};
 		uintptr_t ShadowGSBase, GSBase, FSBase;
 #elif defined(a32)
-		CPU::x32::TrapFrame Registers{};
+		CPU::x32::SchedulerFrame Registers{};
 		uintptr_t ShadowGSBase, GSBase, FSBase;
 #elif defined(aa64)
 		uintptr_t Registers; // TODO

tasking/scheduler/custom.cpp

@@ -556,7 +556,7 @@ namespace Tasking::Scheduler
 		}
 	}
 
-	nsa NIF void Custom::Schedule(CPU::TrapFrame *Frame)
+	nsa NIF void Custom::Schedule(CPU::SchedulerFrame *Frame)
 	{
 		if (unlikely(StopScheduler))
 		{
@@ -564,9 +564,6 @@ namespace Tasking::Scheduler
 			return;
 		}
 		bool ProcessNotChanged = false;
-		/* Restore kernel page table for safety reasons. */
-		if (!SchedulerUpdateTrapFrame)
-			KernelPageTable->Update();
 		uint64_t SchedTmpTicks = TimeManager->GetCounter();
 		this->LastTaskTicks.store(size_t(SchedTmpTicks - this->SchedulerTicks.load()));
 		CPUData *CurrentCPU = GetCurrentCPU();
@@ -671,6 +668,13 @@ namespace Tasking::Scheduler
 		CurrentCPU->CurrentProcess->State.store(TaskState::Running);
 		CurrentCPU->CurrentThread->State.store(TaskState::Running);
 
+		if (CurrentCPU->CurrentThread->Registers.cs != GDT_KERNEL_CODE)
+			CurrentCPU->CurrentThread->Registers.ppt = (uint64_t)(void *)CurrentCPU->CurrentProcess->PageTable;
+		else
+			CurrentCPU->CurrentThread->Registers.ppt = (uint64_t)(void *)KernelPageTable;
+
+		// if (!SchedulerUpdateTrapFrame) {} // TODO
+
 		*Frame = CurrentCPU->CurrentThread->Registers;
 
 #ifdef a64
@@ -713,11 +717,9 @@ namespace Tasking::Scheduler
 		}
 
 		this->SchedulerTicks.store(size_t(TimeManager->GetCounter() - SchedTmpTicks));
-		if (CurrentCPU->CurrentThread->Registers.cs != GDT_KERNEL_CODE)
-			CurrentCPU->CurrentProcess->PageTable->Update();
 	}
 
-	nsa NIF void Custom::OnInterruptReceived(CPU::TrapFrame *Frame)
+	nsa NIF void Custom::OnInterruptReceived(CPU::SchedulerFrame *Frame)
 	{
 		SmartCriticalSection(SchedulerLock);
 		this->Schedule(Frame);

tasking/signal.cpp

@@ -344,7 +344,7 @@ namespace Tasking
 		return {};
 	}
 
-	bool Signal::HandleSignal(CPU::TrapFrame *tf, void *thread)
+	bool Signal::HandleSignal(CPU::SchedulerFrame *tf, void *thread)
 	{
 		/* We don't want to do this in kernel mode */
 		if (tf->cs != GDT_USER_CODE)