zinnia/arch/x86_64/sched.c

#include <zinnia/alloc.h>
#include <zinnia/assert.h>
#include <zinnia/compiler.h>
#include <zinnia/errno.h>
#include <zinnia/percpu.h>
#include <zinnia/pmap.h>
#include <zinnia/sched.h>
#include <zinnia/vma.h>
#include <arch/sched.h>
#include <bits/sched.h>
#include <stddef.h>
#include <stdint.h>
#include <x86_64/asm.h>
#include <x86_64/defs.h>
#include <x86_64/gdt.h>
#include <x86_64/idt.h>
#include <x86_64/percpu.h>

// The task frame consists of registers that the C ABI marks as callee-saved.
// If we don't save them, these registers are lost during a context switch.
struct task_frame {
    uint64_t rbx, rbp, r12, r13, r14, r15, rip;
};

[[__naked]]
static struct task* perform_switch(uint64_t* old_rsp, uint64_t new_rsp, struct task* prev) {
    asm volatile(
        "sub rsp, 0x30\n" // Make room for all regs (except RIP).
        "mov [rsp + %c[rbx]], rbx\n"
        "mov [rsp + %c[rbp]], rbp\n"
        "mov [rsp + %c[r12]], r12\n"
        "mov [rsp + %c[r13]], r13\n"
        "mov [rsp + %c[r14]], r14\n"
        "mov [rsp + %c[r15]], r15\n"
        "mov [rdi], rsp\n" // rdi = old_rsp
        "mov rsp, rsi\n"   // rsi = new_rsp
        "mov rbx, [rsp + %c[rbx]]\n"
        "mov rbp, [rsp + %c[rbp]]\n"
        "mov r12, [rsp + %c[r12]]\n"
        "mov r13, [rsp + %c[r13]]\n"
        "mov r14, [rsp + %c[r14]]\n"
        "mov r15, [rsp + %c[r15]]\n"
        "add rsp, 0x30\n"
        "mov rax, rdx\n"
        "ret" // This will conveniently move us to the RIP we put at this stack entry.
        :
        : [rbx] "i"(offsetof(struct task_frame, rbx)),
          [rbp] "i"(offsetof(struct task_frame, rbp)),
          [r12] "i"(offsetof(struct task_frame, r12)),
          [r13] "i"(offsetof(struct task_frame, r13)),
          [r14] "i"(offsetof(struct task_frame, r14)),
          [r15] "i"(offsetof(struct task_frame, r15))
        : "memory"
    );
}

struct task* arch_sched_switch(struct task* from, struct task* to) {
    struct arch_percpu* cpu = percpu_ref(arch_data);
    struct arch_task_context* from_ctx = &from->context;
    struct arch_task_context* to_ctx = &to->context;

    cpu->tss.rsp0 = kernel_stack_top(&to->kernel_stack);

    if (from_ctx->fpu_region != nullptr) {
        cpu->fpu_save(from_ctx->fpu_region);
        from_ctx->ds = asm_read_ds();
        from_ctx->es = asm_read_es();
        from_ctx->fs = asm_read_fs();
        from_ctx->gs = asm_read_gs();
        from_ctx->fs_base = asm_rdmsr(MSR_FS_BASE);
        from_ctx->gs_base = asm_rdmsr(MSR_KERNEL_GS_BASE);
    }

    if (to_ctx->fpu_region != nullptr) {
        cpu->fpu_restore(to_ctx->fpu_region);
        asm_write_ds(to_ctx->ds);
        asm_write_es(to_ctx->es);
        asm_write_fs(to_ctx->fs);

        // If we have to change the GS segment we need to reload the MSR, otherwise we lose its value.
        if (to_ctx->gs != asm_read_gs()) {
            uint64_t percpu = asm_rdmsr(MSR_GS_BASE);
            asm_write_gs(to_ctx->gs);
            asm_wrmsr(MSR_GS_BASE, percpu);
        }

        asm_wrmsr(MSR_FS_BASE, to_ctx->fs_base);
        asm_wrmsr(MSR_KERNEL_GS_BASE, to_ctx->gs_base);
    }

    return perform_switch(&from_ctx->rsp, to_ctx->rsp, from);
}

// Calls task_entry with the correct parameters.
[[__naked]]
void task_entry_thunk() {
    asm volatile(
        "mov rdi, rbx\n"
        "mov rsi, r12\n"
        "mov rdx, rax\n"
        "jmp %0"
        :
        : "r"(task_entry)
        : "memory", "rdi", "rsi", "rdx"
    );
}

errno_t arch_task_init(
    struct arch_task_context* context,
    void* entry,
    uintptr_t arg0,
    uintptr_t stack_start,
    bool is_user
) {
    struct arch_percpu* cpu = percpu_ref(arch_data);

    assert((stack_start & 0xF) == 0, "Task stack is not 16-byte aligned!\n");

    // Prepare a dummy stack with an entry point function to return to.
    struct task_frame* frame = (struct task_frame*)(stack_start - sizeof(uint64_t)) - 1;
    frame->rbx = (uint64_t)entry;
    frame->r12 = arg0;
    frame->rip = (uint64_t)task_entry_thunk;
    context->rsp = (uint64_t)frame;

    if (is_user) {
        // Allocate FPU space.
        const size_t fpu_size = cpu->fpu_size;
        uintptr_t fpu_base = vm_alloc(fpu_size);
        if (!fpu_base)
            return ENOMEM;
        phys_t pages;
        errno_t s = mem_phys_alloc(ALIGN_UP(fpu_size, arch_mem_page_size()) / arch_mem_page_size(), 0, &pages);
        if (s) {
            vm_free(fpu_base, fpu_size);
            return s;
        }
        for (size_t i = 0; i < fpu_size; i += arch_mem_page_size())
            pmap_map(&kernel_vma.pmap, fpu_base + i, pages + i, PTE_WRITE | PTE_READ, CACHE_WRITE_BACK);
        context->fpu_region = (void*)fpu_base;
        context->fpu_phys = pages;

        context->ds = asm_read_ds();
        context->es = asm_read_es();
        context->fs = asm_read_fs();
        context->gs = asm_read_gs();
        context->fs_base = asm_rdmsr(MSR_FS_BASE);
        context->gs_base = asm_rdmsr(MSR_KERNEL_GS_BASE);
    }

    return 0;
}

void arch_task_fini(struct arch_task_context* context) {
    if (context->fpu_region == nullptr)
        return;

    const size_t page_size = arch_mem_page_size();
    const size_t fpu_size = percpu_ref(arch_data)->fpu_size;
    const uintptr_t fpu_base = (uintptr_t)context->fpu_region;
    for (size_t i = 0; i < fpu_size; i += page_size)
        pmap_unmap(&kernel_vma.pmap, fpu_base + i);
    mem_phys_free(context->fpu_phys, ALIGN_UP(fpu_size, page_size) / page_size);
    vm_free(fpu_base, fpu_size);

    context->fpu_region = nullptr;
}

[[__naked]]
void arch_sched_preempt_disable() {
    asm volatile(
        "inc qword ptr gs:%c0\n"
        "ret"
        :
        : "i"(offsetof(struct percpu, sched.preempt_level))
    );
}

[[__naked]]
bool arch_sched_preempt_enable() {
    asm volatile(
        "dec qword ptr gs:%c0\n"
        "setz al\n"
        "movzx eax, al\n"
        "ret\n"
        :
        : "i"(offsetof(struct percpu, sched.preempt_level))
        : "cc"
    );
}

[[__naked]]
void arch_sched_jump_to_context(struct arch_context* context) {
    asm volatile(
        "mov rsp, rdi\n"
        "jmp %0\n" ::"r"(interrupt_return)
    );
}

void arch_sched_jump_to_user(uintptr_t ip, uintptr_t sp, uintptr_t arg) {
    struct arch_context context = {
        .rip = ip,
        .rsp = sp,
        .rdi = arg,
        .rflags = 0x202,
        .cs = offsetof(struct gdt, user_code64) | CPL_USER,
        .ss = offsetof(struct gdt, user_data) | CPL_USER,
    };

    uintptr_t percpu = asm_rdmsr(MSR_GS_BASE);
    const uint16_t zero = 0;
    asm volatile(
        "mov ds, %0\n"
        "mov es, %0\n"
        "mov fs, %0\n"
        "mov gs, %0\n"
        :
        : "r"(zero)
    );

    asm_wrmsr(MSR_FS_BASE, 0);
    asm_wrmsr(MSR_GS_BASE, percpu);
    asm_wrmsr(MSR_KERNEL_GS_BASE, 0);

    arch_sched_jump_to_context(&context);
}