refactor: tasks code cleanup

This commit is contained in:
Zoe
2026-09-03 07:13:35 -05:00
parent 4c32bcfb35
commit b0f804e412
13 changed files with 607 additions and 114 deletions
+4 -1
View File
@@ -202,7 +202,6 @@ pub unsafe fn enter_user(
"mov ds, {user_data_selector:x}",
"mov es, {user_data_selector:x}",
"mov fs, {user_data_selector:x}",
"mov gs, {user_data_selector:x}",
"push {user_data_selector}",
"push {user_stack_pointer}",
@@ -227,6 +226,10 @@ pub unsafe fn enter_user(
"xor r14, r14",
"xor r15, r15",
// Kernel GS is CpuLocal; leave it in IA32_KERNEL_GS_BASE so
// syscall_entry can recover it with SWAPGS.
"swapgs",
"mov gs, {user_data_selector:x}",
"iretq",
user_data_selector = in(reg) gdt::USER_DATA_SELECTOR as usize,
user_code_selector = in(reg) gdt::USER_CODE_SELECTOR as usize,
+5 -1
View File
@@ -44,6 +44,10 @@ pub fn init(cpu_local: *const CpuLocal) {
write_msr(IA32_FMASK, RFLAGS_MASK);
write_msr(IA32_GS_BASE, 0);
write_msr(IA32_KERNEL_GS_BASE, cpu_local as u64);
// Kernel code always runs with GS pointing at CpuLocal. User entry
// swaps this into IA32_KERNEL_GS_BASE before transitioning to ring 3.
asm!("swapgs", options(nostack, preserves_flags));
}
}
@@ -103,7 +107,7 @@ unsafe extern "C" fn syscall_entry() {
}
extern "C" fn syscall_dispatch(frame: &mut SyscallFrame) {
let ret = crate::task::syscall::handle(
let ret = crate::syscall::handle(
frame.rax, frame.rdi, frame.rsi, frame.rdx, frame.r10, frame.r8, frame.r9,
);
+6
View File
@@ -160,6 +160,12 @@ pub fn init() -> Result<(), SerialPortError> {
com1().init()
}
pub fn write_bytes(bytes: &[u8]) {
for byte in bytes {
com1().write_byte(*byte);
}
}
pub fn print(args: core::fmt::Arguments) {
use core::fmt::Write;
+192 -60
View File
@@ -8,12 +8,16 @@ mod boot;
mod debug;
mod memory;
mod platform;
mod syscall;
mod task;
use core::arch::global_asm;
use crate::{
debug::serial,
memory::{
AddressSpace, KernelStackPool, MemoryRegionKind, PagePermissions, UserStack, VirtualAddr,
AddressSpace, DirectMap, FrameAllocator, KernelStackPool, MemoryRegionKind,
PagePermissions, UserStack, VirtualAddr,
},
task::tcb::Tcb,
};
@@ -88,6 +92,112 @@ pub extern "C" fn _start() -> ! {
}
}
unsafe extern "C" {
static task_a_start: u8;
static task_a_end: u8;
static task_b_start: u8;
static task_b_end: u8;
static task_c_start: u8;
static task_c_end: u8;
}
unsafe fn embedded_code(start: *const u8, end: *const u8) -> &'static [u8] {
let length = unsafe { end.offset_from(start) as usize };
unsafe { core::slice::from_raw_parts(start, length) }
}
global_asm!(
r#"
.global task_a_start
task_a_start:
mov r12d, 100
.Ltask_a_loop:
mov eax, 3
mov edi, 1
lea rsi, [rip + .Ltask_a_message]
mov edx, 7
xor r10d, r10d
syscall
mov eax, 1
syscall
dec r12d
jnz .Ltask_a_loop
.Ltask_a_done:
mov eax, 2
syscall
jmp .Ltask_a_done
.Ltask_a_message:
.ascii "task A\n"
.global task_a_end
task_a_end:
.global task_b_start
task_b_start:
mov r12d, 100
.Ltask_b_loop:
mov eax, 3
mov edi, 1
lea rsi, [rip + .Ltask_b_message]
mov edx, 7
xor r10d, r10d
syscall
mov eax, 1
syscall
dec r12d
jnz .Ltask_b_loop
.Ltask_b_done:
mov eax, 2
syscall
jmp .Ltask_b_done
.Ltask_b_message:
.ascii "task B\n"
.global task_b_end
task_b_end:
.global task_c_start
task_c_start:
mov r12d, 100
.Ltask_c_loop:
mov eax, 3
mov edi, 1
lea rsi, [rip + .Ltask_c_message]
mov edx, 7
xor r10d, r10d
syscall
mov eax, 1
syscall
dec r12d
jnz .Ltask_c_loop
.Ltask_c_done:
mov eax, 2
syscall
jmp .Ltask_c_done
.Ltask_c_message:
.ascii "task C\n"
.global task_c_end
task_c_end:
"#
);
pub unsafe extern "C" fn kernel_main(handoff: *mut KernelHandoff) -> ! {
let (
mut allocator,
@@ -134,73 +244,95 @@ pub unsafe extern "C" fn kernel_main(handoff: *mut KernelHandoff) -> ! {
arch::init_interrupt_controller(&madt, &mut allocator, &mut address_space)
.expect("failed to initialize interrupt controller");
println!("Creating user address space...");
let task_a_code = unsafe { embedded_code(&task_a_start, &task_a_end) };
let task_b_code = unsafe { embedded_code(&task_b_start, &task_b_end) };
let task_c_code = unsafe { embedded_code(&task_c_start, &task_c_end) };
let task_kernel_stack = kernel_stack_pool
.allocate(&mut address_space, &mut allocator)
.expect("failed to allocate user task stack");
let task_a = create_test_task(
task_a_code,
&mut address_space,
&mut kernel_stack_pool,
&mut allocator,
direct_map,
);
let mut user_addr_space = address_space
.new_user(&mut allocator)
.expect("failed to create user address space");
let task_b = create_test_task(
task_b_code,
&mut address_space,
&mut kernel_stack_pool,
&mut allocator,
direct_map,
);
println!("Allocating user stack...");
let task_c = create_test_task(
task_c_code,
&mut address_space,
&mut kernel_stack_pool,
&mut allocator,
direct_map,
);
let user_stack = UserStack::allocate(&mut user_addr_space, &mut allocator)
.expect("failed to allocate user stack");
let user_instruction_pointer = VirtualAddr::new(0x8000);
let user_code_page = allocator
.alloc()
.expect("failed to allocate user code page")
.frame_address();
user_addr_space
.map(
user_code_page.start_address(),
user_instruction_pointer,
PagePermissions::new(true, true, true),
&mut allocator,
memory::CachePolicy::WriteBack,
)
.expect("failed to map user code page");
unsafe {
let user_code: &[u8] = &[
0xCC, // INT3
0x0F, 0x05, // SYSCALL
0xEB, 0xFE, // JMP -2 (loop forever if exit returns)
];
let user_code_virtual = direct_map
.translate(user_code_page.start_address())
.unwrap();
println!("Copying user code to {:#X}", user_code_virtual.as_usize());
core::ptr::copy_nonoverlapping(
user_code.as_ptr(),
user_code_virtual.as_mut_ptr::<u8>(),
user_code.len(),
);
println!("Activating user address space...");
let user_tcb = Tcb::new_user(
0,
user_addr_space,
task_kernel_stack,
user_instruction_pointer,
user_stack.top(),
);
println!("Running first user task {user_tcb:?}...");
task::tcb::run_first_user(&user_tcb);
}
task::scheduler::add_task(task_a).expect("scheduler is full");
task::scheduler::add_task(task_b).expect("scheduler is full");
task::scheduler::add_task(task_c).expect("scheduler is full");
task::scheduler::start();
hcf();
}
fn create_test_task(
code: &[u8],
kernel_address_space: &mut AddressSpace,
kernel_stack_pool: &mut KernelStackPool,
allocator: &mut FrameAllocator,
direct_map: DirectMap,
) -> Tcb {
assert!(code.len() <= memory::FRAME_SIZE);
let kernel_stack = kernel_stack_pool
.allocate(kernel_address_space, allocator)
.expect("failed to allocate task kernel stack");
let mut user_address_space = kernel_address_space
.new_user(allocator)
.expect("failed to create user address space");
let user_stack = UserStack::allocate(&mut user_address_space, allocator)
.expect("failed to allocate user stack");
let entry = VirtualAddr::new(0x8000);
let code_frame = allocator
.alloc()
.expect("failed to allocate code frame")
.into_raw();
user_address_space
.map(
code_frame.start_address(),
entry,
PagePermissions::new(false, true, true),
allocator,
memory::CachePolicy::WriteBack,
)
.expect("failed to map user code");
let destination = direct_map
.translate(code_frame.start_address())
.expect("code frame outside direct map");
unsafe {
core::ptr::copy_nonoverlapping(code.as_ptr(), destination.as_mut_ptr(), code.len());
}
Tcb::new_user(
0, // overwritten by add_task for now
user_address_space,
kernel_stack,
entry,
user_stack.top(),
)
}
#[panic_handler]
fn panic(info: &core::panic::PanicInfo) -> ! {
println!("Uh oh, something went wrong!");
+2 -2
View File
@@ -2,7 +2,7 @@ use crate::{
arch::{PageTable, PageTableCreateError, PageTableMapError, PageTableUnmapError, PagingConfig},
memory::{
CachePolicy, DirectMap, FRAME_SIZE, FrameAddr, FrameAllocator, KernelMemoryLayout,
MemoryRegion, MemoryRegionKind, PagePermissions, PhysicalAddr, VirtualAddr,
MemoryRegion, MemoryRegionKind, PagePermissions, PhysicalAddr, USER_SPACE_END, VirtualAddr,
},
};
@@ -195,7 +195,7 @@ impl AddressSpace {
let global = self.kind == AddressSpaceKind::Kernel;
if self.kind == AddressSpaceKind::User {
if virtual_addr.as_usize() >= 0x0000_8000_0000_0000 {
if virtual_addr.as_usize() >= USER_SPACE_END.as_usize() {
return Err(MapError::InvalidUserAddress);
}
+3
View File
@@ -1,12 +1,15 @@
mod address_space;
mod frame;
mod stack;
mod user;
#[allow(unused)]
pub use address_space::{AddressSpace, AddressSpaceCreateError, MapError, UnmapError};
pub use frame::{FRAME_SIZE, FrameAddr, FrameAllocator, OwnedFrame};
#[allow(unused)]
pub use stack::{KernelStack, KernelStackPool, StackCreateError, UserStack};
#[allow(unused)]
pub use user::*;
pub struct KernelSegment {
pub physical_base: PhysicalAddr,
+71
View File
@@ -0,0 +1,71 @@
use crate::{memory::VirtualAddr, syscall::Status};
pub const USER_SPACE_END: VirtualAddr = VirtualAddr::new(0x0000_8000_0000_0000);
pub fn copy_from_user(src: VirtualAddr, dst: &mut [u8]) -> Result<(), Status> {
// TODO: guard against unmapped pages
let end = src
.as_usize()
.checked_add(dst.len())
.ok_or(Status::BadAddress)?;
if end > USER_SPACE_END.as_usize() {
return Err(Status::BadAddress);
}
unsafe {
core::ptr::copy_nonoverlapping(src.as_ptr(), dst.as_mut_ptr(), dst.len());
}
Ok(())
}
pub fn copy_to_user(dst: VirtualAddr, src: &[u8]) -> Result<(), Status> {
let end = dst
.as_usize()
.checked_add(src.len())
.ok_or(Status::BadAddress)?;
if end > USER_SPACE_END.as_usize() {
return Err(Status::BadAddress);
}
unsafe {
core::ptr::copy_nonoverlapping(src.as_ptr(), dst.as_mut_ptr::<u8>(), src.len());
}
Ok(())
}
pub fn copy_val_to_user<T: Copy>(dst: VirtualAddr, val: &T) -> Result<(), Status> {
if dst.as_usize() % core::mem::align_of::<T>() != 0 {
return Err(Status::InvalidArgument);
}
let end = dst
.as_usize()
.checked_add(core::mem::size_of::<T>())
.ok_or(Status::BadAddress)?;
if end > USER_SPACE_END.as_usize() {
return Err(Status::BadAddress);
}
unsafe {
(dst.as_mut_ptr::<T>()).write(*val);
}
Ok(())
}
pub fn copy_val_from_user<T: Copy>(src: VirtualAddr) -> Result<T, Status> {
if src.as_usize() % core::mem::align_of::<T>() != 0 {
return Err(Status::InvalidArgument);
}
let end = src
.as_usize()
.checked_add(core::mem::size_of::<T>())
.ok_or(Status::BadAddress)?;
if end > USER_SPACE_END.as_usize() {
return Err(Status::BadAddress);
}
let val = unsafe { src.as_ptr::<T>().read() };
Ok(val)
}
+52
View File
@@ -0,0 +1,52 @@
mod table;
use table::*;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u64)]
pub enum Status {
// Status::Success = 0
InvalidArgument = 1, // EINVAL
BadAddress = 2, // EFAULT
BadFileDescriptor = 3, // EBADF
NoSuchTask = 4, // ESRCH
OutOfMemory = 5, // ENOMEM
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u64)]
pub enum SyscallNumber {
Yield = 1,
Exit = 2,
Write = 3,
}
impl TryFrom<u64> for SyscallNumber {
type Error = Status;
fn try_from(val: u64) -> Result<Self, Self::Error> {
match val {
1 => Ok(Self::Yield),
2 => Ok(Self::Exit),
3 => Ok(Self::Write),
_ => Err(Status::InvalidArgument),
}
}
}
pub fn handle(num: u64, arg0: u64, arg1: u64, arg2: u64, arg3: u64, _arg4: u64, _arg5: u64) -> u64 {
let result = (|| -> Result<(), Status> {
let syscall = SyscallNumber::try_from(num)?;
match syscall {
SyscallNumber::Yield => sys_yield(),
SyscallNumber::Exit => sys_exit(arg0 as usize),
SyscallNumber::Write => {
sys_write(arg0 as usize, arg1 as usize, arg2 as usize, arg3 as usize)
}
}
})();
match result {
Ok(()) => 0,
Err(err) => err as u64,
}
}
+33
View File
@@ -0,0 +1,33 @@
use crate::memory::{VirtualAddr, copy_from_user, copy_val_to_user};
use super::Status;
pub fn sys_yield() -> Result<(), Status> {
crate::task::scheduler::yield_current();
Ok(())
}
pub fn sys_exit(exit_code: usize) -> ! {
crate::task::scheduler::exit_current(exit_code);
}
pub fn sys_write(fd: usize, buf_ptr: usize, len: usize, out_ptr: usize) -> Result<(), Status> {
if fd != 1 && fd != 2 {
return Err(Status::BadFileDescriptor);
}
let mut chunk = [0u8; 128];
let mut written = 0;
while written < len {
let n = (len - written).min(chunk.len());
copy_from_user(VirtualAddr::new(buf_ptr + written), &mut chunk[..n])?;
crate::debug::serial::write_bytes(&chunk[..n]);
written += n;
}
if out_ptr != 0 {
copy_val_to_user(VirtualAddr::new(out_ptr), &written)?;
}
Ok(())
}
+1 -1
View File
@@ -1,2 +1,2 @@
pub mod syscall;
pub mod scheduler;
pub mod tcb;
+230
View File
@@ -0,0 +1,230 @@
use core::cell::UnsafeCell;
use crate::{
arch::ThreadContext,
memory::{AddressSpace, VirtualAddr},
println,
task::tcb::{ExitReason, Tcb, ThreadState},
};
const MAX_TASKS: usize = 32;
type TaskId = usize;
struct Scheduler {
current: Option<TaskId>,
tasks: [Option<Tcb>; MAX_TASKS],
ready: ReadyQueue,
}
impl Scheduler {
const fn new() -> Self {
Self {
current: None,
tasks: [const { None }; MAX_TASKS],
ready: ReadyQueue::new(),
}
}
fn make_switch(&mut self, current_id: TaskId, next_id: TaskId) -> Switch {
assert_ne!(current_id, next_id);
let current = self.tasks[current_id].as_mut().unwrap();
let prev_ctx = &mut current.context as *mut ThreadContext;
let prev_addr_space = &current.address_space as *const AddressSpace;
let next = self.tasks[next_id].as_ref().unwrap();
let next_ctx = &next.context as *const ThreadContext;
let next_addr_space = &next.address_space as *const AddressSpace;
let next_kernel_stack = next.kernel_stack.top();
Switch {
previous_context: prev_ctx,
next_context: next_ctx,
next_address_space: next_addr_space,
next_kernel_stack: next_kernel_stack,
activate_address_space: unsafe { *next_addr_space != *prev_addr_space },
}
}
}
struct Switch {
previous_context: *mut ThreadContext,
next_context: *const ThreadContext,
next_address_space: *const AddressSpace,
next_kernel_stack: VirtualAddr,
activate_address_space: bool,
}
impl Switch {
unsafe fn perform(self) {
if self.activate_address_space {
unsafe {
(&*self.next_address_space).activate();
}
}
crate::arch::set_kernel_stack(self.next_kernel_stack);
unsafe {
crate::arch::switch_context(self.previous_context, self.next_context);
}
}
}
struct ReadyQueue {
entries: [TaskId; MAX_TASKS],
head: usize,
len: usize,
}
impl ReadyQueue {
pub const fn new() -> Self {
Self {
entries: [0; MAX_TASKS],
head: 0,
len: 0,
}
}
pub fn push_back(&mut self, task: TaskId) -> bool {
if self.len == MAX_TASKS {
return false;
}
let tail = (self.head + self.len) % MAX_TASKS;
self.entries[tail] = task;
self.len += 1;
true
}
pub fn pop_front(&mut self) -> Option<TaskId> {
if self.len == 0 {
return None;
}
let task = self.entries[self.head];
self.head = (self.head + 1) % MAX_TASKS;
self.len -= 1;
Some(task)
}
}
struct GlobalScheduler(UnsafeCell<Scheduler>);
unsafe impl Sync for GlobalScheduler {}
static SCHEDULER: GlobalScheduler = GlobalScheduler(UnsafeCell::new(Scheduler::new()));
pub fn add_task(mut task: Tcb) -> Result<TaskId, Tcb> {
let interrupt_state = crate::arch::disable_interrupts_and_save();
let result = {
let scheduler = unsafe { &mut *SCHEDULER.0.get() };
match scheduler.tasks.iter().position(Option::is_none) {
Some(id) => {
task.id = id;
task.state = ThreadState::Ready;
scheduler.tasks[id] = Some(task);
assert!(scheduler.ready.push_back(id));
Ok(id)
}
None => Err(task),
}
};
crate::arch::restore_interrupts(interrupt_state);
result
}
pub fn start() -> ! {
crate::arch::disable_interrupts();
let mut bootstrap_context = ThreadContext::empty();
let switch = {
let scheduler = unsafe { &mut *SCHEDULER.0.get() };
let next_id = scheduler.ready.pop_front().expect("no tasks to run");
let next = scheduler.tasks[next_id]
.as_mut()
.expect("ready task is missing");
next.state = ThreadState::Running;
scheduler.current = Some(next_id);
Switch {
previous_context: &mut bootstrap_context,
next_context: &next.context,
next_address_space: &next.address_space,
next_kernel_stack: next.kernel_stack.top(),
activate_address_space: true,
}
};
unsafe {
switch.perform();
}
panic!("scheduler returned to bootstrap context");
}
pub fn yield_current() {
let interrupt_state = crate::arch::disable_interrupts_and_save();
let switch = {
let scheduler = unsafe { &mut *SCHEDULER.0.get() };
let Some(next_id) = scheduler.ready.pop_front() else {
crate::arch::restore_interrupts(interrupt_state);
return;
};
let current_id = scheduler.current.expect("no current task");
scheduler.tasks[current_id].as_mut().unwrap().state = ThreadState::Ready;
assert!(scheduler.ready.push_back(current_id));
scheduler.tasks[next_id].as_mut().unwrap().state = ThreadState::Running;
scheduler.current = Some(next_id);
scheduler.make_switch(current_id, next_id)
};
unsafe {
switch.perform();
}
// this runs when this task is selected to run again
crate::arch::restore_interrupts(interrupt_state);
}
pub fn exit_current(exit_code: usize) -> ! {
crate::arch::disable_interrupts();
let switch = {
let scheduler = unsafe { &mut *SCHEDULER.0.get() };
let current_id = scheduler.current.expect("no current task");
let Some(next_id) = scheduler.ready.pop_front() else {
println!("All tasks exited");
crate::hcf();
};
let current = scheduler.tasks[current_id].as_mut().unwrap();
current.state = ThreadState::Dead(ExitReason::Exited(exit_code));
scheduler.tasks[next_id].as_mut().unwrap().state = ThreadState::Running;
scheduler.current = Some(next_id);
scheduler.make_switch(current_id, next_id)
};
unsafe {
switch.perform();
}
panic!("dead task was scheduled again");
}
-18
View File
@@ -1,18 +0,0 @@
use crate::{hcf, println};
pub fn handle(
syscall_num: u64,
arg0: u64,
arg1: u64,
arg2: u64,
arg3: u64,
arg4: u64,
arg5: u64,
) -> u64 {
println!(
"Syscall nr={:#X} args=({:#X}, {:#X}, {:#X}, {:#X}, {:#X}, {:#X})",
syscall_num, arg0, arg1, arg2, arg3, arg4, arg5
);
hcf();
0
}
+8 -31
View File
@@ -1,16 +1,22 @@
use crate::{
arch::ThreadContext,
memory::{AddressSpace, KernelStack, VirtualAddr},
println,
};
// Thread Control Block
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExitReason {
Exited(usize),
Killed,
Fault,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ThreadState {
Ready,
Running,
Blocked,
Dead,
Dead(ExitReason),
}
#[derive(Debug)]
@@ -41,32 +47,3 @@ impl Tcb {
}
}
}
pub unsafe fn run_first_user(user_tcb: &Tcb) {
let previous_interrupts = crate::arch::disable_interrupts_and_save();
let mut boot_thread_ctx = ThreadContext::empty();
unsafe {
user_tcb.address_space.activate();
crate::arch::set_kernel_stack(user_tcb.kernel_stack.top());
crate::arch::switch_context(&mut boot_thread_ctx, &user_tcb.context);
}
crate::arch::restore_interrupts(previous_interrupts);
}
pub unsafe fn switch(prev: &mut Tcb, next: &Tcb) {
let previous_interrupts = crate::arch::disable_interrupts_and_save();
if prev.address_space != next.address_space {
unsafe { next.address_space.activate() };
}
crate::arch::set_kernel_stack(next.kernel_stack.top());
unsafe {
crate::arch::switch_context(&mut prev.context, &next.context);
}
crate::arch::restore_interrupts(previous_interrupts);
}