From 86a119ea51dc700923b523d500c2de184a4a14ea Mon Sep 17 00:00:00 2001 From: EnderIce2 Date: Sun, 17 Nov 2024 03:11:20 +0200 Subject: [PATCH] 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 | 67 +++++++++++++++++++++++++++++++++++- core/interrupts_manager.cpp | 47 ++++++++++++++++++++++--- include/cpu.hpp | 35 +++++++++++++++++++ include/ints.hpp | 1 + include/scheduler.hpp | 4 +-- include/signal.hpp | 6 ++-- include/task.hpp | 4 +-- tasking/scheduler/custom.cpp | 16 +++++---- tasking/signal.cpp | 2 +- 9 files changed, 161 insertions(+), 21 deletions(-) diff --git a/arch/amd64/cpu/idt.cpp b/arch/amd64/cpu/idt.cpp index 43ac1ac2..b642361d 100644 --- a/arch/amd64/cpu/idt.cpp +++ b/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) diff --git a/core/interrupts_manager.cpp b/core/interrupts_manager.cpp index 4302894f..cd24689c 100644 --- a/core/interrupts_manager.cpp +++ b/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); } } diff --git a/include/cpu.hpp b/include/cpu.hpp index 5e5f1d70..49090b75 100644 --- a/include/cpu.hpp +++ b/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 } diff --git a/include/ints.hpp b/include/ints.hpp index 21acd247..84fbc452 100644 --- a/include/ints.hpp +++ b/include/ints.hpp @@ -78,6 +78,7 @@ namespace Interrupts public: virtual void OnInterruptReceived(CPU::TrapFrame *Frame); + virtual void OnInterruptReceived(CPU::SchedulerFrame *Frame); }; } diff --git a/include/scheduler.hpp b/include/scheduler.hpp index bd8640f5..0404fb19 100644 --- a/include/scheduler.hpp +++ b/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(); diff --git a/include/signal.hpp b/include/signal.hpp index eecbfd14..9d66277f 100644 --- a/include/signal.hpp +++ b/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); diff --git a/include/task.hpp b/include/task.hpp index b374410b..ec0d382c 100644 --- a/include/task.hpp +++ b/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 diff --git a/tasking/scheduler/custom.cpp b/tasking/scheduler/custom.cpp index fa12cbeb..b88005ed 100644 --- a/tasking/scheduler/custom.cpp +++ b/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); diff --git a/tasking/signal.cpp b/tasking/signal.cpp index ae77a023..c3b0e9d1 100644 --- a/tasking/signal.cpp +++ b/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)