Page fault switching during coroutine implementation

[BigBadE]

I'm preparing my OS for jumping into Userspace in long mode and decided to work on a coroutine system because it'd both give me a testable part of context switching and make IRQ handling easier (so I can yield while waiting for IRQs). I'm using Rust, so I believe I'd need to store the stack pointer and not the registers (however, that may be wrong).

To context switch, I do the following:
1: Allocate a new stack
2: Push the current stack pointer and the function to be called onto the heap
3: Store a reference to that data in the r15 register
4: Set the stack pointer to the new allocated stack and jump to a load function
5: Load the old stack pointer and call the function
6: After the function, set the stack pointer to the old stack pointer and call an empty function to pop the stack

However, when the load function returns, a General Protection Fault is thrown, with the instruction pointer being an obviously wrong value because my kernel is higher half. So that makes me question: Where should I jump back? I can't jump to the instruction after 4 (where my code was before the swap) because I'd have to save the instruction pointer before the swap, so I'd jump right back to the swap. I was hoping the empty function would pop the stack back to where it should be but it looks like it doesn't.

Code:

Code: Select all

const STACK_SIZE: usize = 0x7C00 - 0x500;

pub struct SavedState {
    pub stack: *mut u8,
}

impl SavedState {
    pub fn load(&self) {}

    pub fn swap(&self, target: *mut SavedState) {}
}

impl Drop for SavedState {
    fn drop(&mut self) {
        unsafe {
            ALLOCATOR.dealloc((self.stack as usize - STACK_SIZE) as *mut u8, Layout::from_size_align_unchecked(STACK_SIZE, 8));
        }
    }
}

struct TempData {
    stack: *mut u8,
    target: fn(SavedState, u8),
    arg: u8,
}

pub fn save_state(target: fn(SavedState, u8), arg: u8) {
    unsafe {
        //Get a new stack
        let stack = (ALLOCATOR.alloc_zeroed(Layout::from_size_align_unchecked(STACK_SIZE, 8)) as usize + STACK_SIZE) as *mut u8;

        //Load current stack pointer
        let pointer: *mut u8;
        asm!("mov {}, rsp", out(reg) pointer);

        //Push temp data
        let temp = ALLOCATOR.alloc_zeroed(Layout::from_size_align_unchecked(size_of::<TempData>(), 8));
        ptr::write(temp as *mut TempData, TempData { stack: pointer, target, arg });

        //Send temp data to loaded function
        asm!("mov r15, {}", in(reg) temp, options(nomem, nostack, preserves_flags));

        //Switch to new stack
        switch_state(stack, load_state as *mut fn());
    }
}

pub unsafe fn safe_return() {
    println!("Returned");
}

unsafe fn switch_state(stack: *mut u8, target: *mut fn()) {
    unsafe {
        //Directly loading the register to the jump page faults so I just do this. If anyone knows why, I'm all ears
        asm!("mov r14, {}", in(reg) target);
        asm!("mov rsp, {}", in(reg) stack);
        asm!("jmp r14");
    }
}

pub fn load_state() {
    unsafe {
        //Load the temp data
        let mut position: *mut TempData;
        asm!("mov {}, r15", out(reg) position, options(nomem, nostack, preserves_flags));
        let data = ptr::read(position);

        //Call the passed function
        (data.target)(SavedState { stack: data.stack }, data.arg);

        //If it doesn't load the last state itself, do it here
        let pointer: *mut u8;
        asm!("mov {}, rsp", out(reg) pointer);

        switch_state(data.stack, safe_return as *mut fn());
    }
}

"Returned!" is correctly printed, right before the General Protection Fault.

I'd appriciate any help or advice relating to this topic.

[Octocontrabass]

I can't say I'm very familiar with Rust, but I doubt changing the stack pointer in the middle of a function is allowed! I expect you'll need to write the stack-switching function purely in assembly, like the example here. When you allocate a new stack, fill it with whatever state you want this function to restore.