From 9b5fb222f5b89a3f50a8e98e0200bbd380cb8c3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Thomas=20Pr=C3=A9vost?= <7011145@gmail.com> Date: Mon, 2 Jan 2023 14:59:20 +0100 Subject: [PATCH] Add cooperative multitasking support --- src/dos.rs | 1 + .../cooperative_task_switching.S | 60 +++++++++ src/dos/cooperative_multitasking/mod.rs | 123 ++++++++++++++++++ src/dos/cooperative_multitasking/task.rs | 44 +++++++ src/dos_tests/allocator_test.rs | 1 + .../cooperative_multitasking_test.rs | 31 +++++ src/dos_tests/mod.rs | 3 +- src/lib.rs | 2 + src/main.rs | 13 +- 9 files changed, 271 insertions(+), 7 deletions(-) create mode 100644 src/dos/cooperative_multitasking/cooperative_task_switching.S create mode 100644 src/dos/cooperative_multitasking/mod.rs create mode 100644 src/dos/cooperative_multitasking/task.rs create mode 100644 src/dos_tests/cooperative_multitasking_test.rs diff --git a/src/dos.rs b/src/dos.rs index e4271b1..211dbb3 100644 --- a/src/dos.rs +++ b/src/dos.rs @@ -7,6 +7,7 @@ pub mod file; pub mod error_code; pub mod panic; pub mod math; +pub mod cooperative_multitasking; use core::arch::asm; pub use alloc::string::String as String; diff --git a/src/dos/cooperative_multitasking/cooperative_task_switching.S b/src/dos/cooperative_multitasking/cooperative_task_switching.S new file mode 100644 index 0000000..8c4b77f --- /dev/null +++ b/src/dos/cooperative_multitasking/cooperative_task_switching.S @@ -0,0 +1,60 @@ +# inspired from https://wiki.osdev.org/Cooperative_Multitasking + + ; .section .text + ; .global cooperative_task_switching_assembly +cooperative_task_switching_assembly: + pushad # 32 bytes + pushfd # 4 bytes # pushf + mov eax, cr3 #Push CR3 + push eax + + mov eax, [44+esp] #The first argument, where to save + mov [4+eax], ebx + mov [8+eax], ecx + mov [12+eax], edx + mov [16+eax], esi + mov [20+eax], edi + + mov ebx, [36+esp] # EAX + mov ecx, [40+esp] # IP + mov edx, [20+esp] # ESP + add edx, 4 # Remove the return value + mov esi, [16+esp] # EBP + mov edi, [4+esp] # EFLAGS + + mov [eax], ebx + + mov [24+eax], edx + mov [28+eax], esi + mov [32+eax], ecx + mov [36+eax], edi + pop ebx # CR3 + mov [40+eax], ebx + push ebx # Goodbye again + mov eax, [48+esp] # The second argument, where to load + + mov ebx, [4+eax] # EBX + mov ecx, [8+eax] # ECX + mov edx, [12+eax] # EDX + mov esi, [16+eax] # ESI + mov edi, [20+eax] # EDI + mov ebp, [28+eax] # EBP + + push eax + mov eax, [36+eax] # EFLAGS + push eax + popfd + pop eax + + mov esp, [24+eax] # ESP ## error ? + push eax + + mov eax, [40+eax] # CR3 + mov cr3, eax + pop eax + + push eax + mov eax, [32+eax] # EIP + xchg eax, [esp] # Cannot use a tmp storage + mov eax, [eax] # EAX + ret diff --git a/src/dos/cooperative_multitasking/mod.rs b/src/dos/cooperative_multitasking/mod.rs new file mode 100644 index 0000000..d7d0bc2 --- /dev/null +++ b/src/dos/cooperative_multitasking/mod.rs @@ -0,0 +1,123 @@ +use alloc::collections::VecDeque; +use core::arch::{asm, global_asm}; +use crate::dos::cooperative_multitasking::task::{Registers, Task}; + +mod task; + +global_asm!(include_str!("cooperative_task_switching.S")); + +extern "C" { + fn cooperative_task_switching_assembly(from: *mut Registers, to: *mut Registers) -> (); +} + +pub struct Tasking{ + task_list: Option>, + current_task_id: u8, + eflags_register: u32, + cr3_register: u32, + initialized: bool, +} + +impl Tasking { + const MAX_TASKS: usize = 10; + + pub fn init(&mut self) { + if self.task_list.is_some() { + self.task_list = None; + } + let (eflags, cr3) = Self::get_eflags_and_cr3_registers(); + + // Create main task + self.task_list = Some(VecDeque::with_capacity(Self::MAX_TASKS)); + self.current_task_id = 0; + self.eflags_register = eflags; + self.cr3_register = cr3; + self.initialized = true; + self.task_list.as_mut().unwrap().push_back(Task { + registers: Registers { + eax: 0, + ebx: 0, + ecx: 0, + edx: 0, + esi: 0, + edi: 0, + esp: 0, + ebp: 0, + eip: 0, + eflags, + cr3, + }, + }); + } + + pub fn add_task(&mut self, main_function: *mut fn()) -> Result<(), &'static str> { + if !self.initialized { + return Err("Cooperative tasking manager is not initialized"); + } + if self.task_list.as_ref().unwrap().len() >= Self::MAX_TASKS { + return Err("Maximum number of tasks reached"); + } + let task_list = self.task_list.as_mut().unwrap(); + task_list.push_back(Task::new(main_function, self.eflags_register, self.cr3_register as *mut u32, task_list.len() as u8)); + Ok(()) + } + + pub fn yield_task(&mut self) { + if !self.initialized { + panic!("Cooperative tasking manager is not initialized"); + } + + let task_list = self.task_list.as_mut().unwrap(); + + let current_task_registers_ptr = &mut task_list[self.current_task_id as usize].registers as *mut Registers; + + self.current_task_id += 1; + if self.current_task_id >= task_list.len() as u8 { + self.current_task_id = 0; + } + + let next_task_registers_ptr = &mut task_list[self.current_task_id as usize].registers as *mut Registers; + + unsafe { + cooperative_task_switching_assembly(current_task_registers_ptr, next_task_registers_ptr); + } + } + + fn get_eflags_and_cr3_registers() -> (u32, u32) { + let mut eflags: u32; + let mut cr3: u32; + unsafe { + // Read CR3 + asm!("mov {}, cr3", out(reg) cr3); + // Read EFLAGS + asm!("pushfd; mov eax, [esp]; mov {}, eax; popfd;", out(reg) eflags); + } + (eflags, cr3) + } +} + +pub static mut TASKING: Tasking = Tasking { + task_list: None, + current_task_id: 0, + eflags_register: 0, + cr3_register: 0, + initialized: false, +}; + +#[macro_export] +macro_rules! yield_cooperative_task { + () => { + unsafe { + $crate::dos::cooperative_multitasking::TASKING.yield_task(); + } + }; +} + +#[macro_export] +macro_rules! add_cooperative_task { + ($main_function: expr) => { + unsafe { + $crate::dos::cooperative_multitasking::TASKING.add_task($main_function as *mut fn()) + } + }; +} \ No newline at end of file diff --git a/src/dos/cooperative_multitasking/task.rs b/src/dos/cooperative_multitasking/task.rs new file mode 100644 index 0000000..6aaa441 --- /dev/null +++ b/src/dos/cooperative_multitasking/task.rs @@ -0,0 +1,44 @@ +#[repr(C)] // To ensure that the struct is laid out in the same way as the assembly code expects +#[derive(Copy, Clone, Debug)] +pub(crate)struct Registers { + pub eax: u32, + pub ebx: u32, + pub ecx: u32, + pub edx: u32, + pub esi: u32, + pub edi: u32, + pub esp: u32, + pub ebp: u32, + pub eip: u32, + pub eflags: u32, + pub cr3: u32, +} + +#[derive(Debug)] +pub(crate) struct Task { + pub(crate) registers: Registers, +} + +// In order to use heap as stack, we need to change ss stack segment register +impl Task { + const TASK_STACK_SIZE: usize = 4096; + + /// Max stack for each task, including the main task, is 4KB + pub fn new(main_function: *mut fn(), flags: u32, pagedir: *mut u32, task_index: u8) -> Task { + Task { + registers: Registers { + eax: 0, + ebx: 0, + ecx: 0, + edx: 0, + esi: 0, + edi: 0, + esp: 0xffff as u32 - (Self::TASK_STACK_SIZE as u32 * task_index as u32), + ebp: 0, + eip: main_function as u32, + eflags: flags, + cr3: pagedir as u32, + } + } + } +} \ No newline at end of file diff --git a/src/dos_tests/allocator_test.rs b/src/dos_tests/allocator_test.rs index 6cf632e..78c0294 100644 --- a/src/dos_tests/allocator_test.rs +++ b/src/dos_tests/allocator_test.rs @@ -1,6 +1,7 @@ use rust_dos::*; use dos::*; +#[allow(dead_code)] pub(crate) fn allocator_test() { let mut box1 = Box::new(5); assert_eq!(*box1, 5); diff --git a/src/dos_tests/cooperative_multitasking_test.rs b/src/dos_tests/cooperative_multitasking_test.rs new file mode 100644 index 0000000..d77e011 --- /dev/null +++ b/src/dos_tests/cooperative_multitasking_test.rs @@ -0,0 +1,31 @@ +use dos::*; +use rust_dos::*; + +#[allow(dead_code)] +pub(crate) fn cooperative_multitasking_test() { + add_cooperative_task!(task_2_main).unwrap(); + println!("Hello from main task!"); + yield_cooperative_task!(); + println!("Hello from main task! (bis)"); + yield_cooperative_task!(); + println!("Hello from main task! (tris)"); +} + +fn task_2_main() { + add_cooperative_task!(task_3_main).unwrap(); + for _ in 0..2 { + let task_number = 2; + let task2_string = String::from("Hello from task 2!"); + println!("Message from task{}: {}", task_number, task2_string); + yield_cooperative_task!(); + } +} + +fn task_3_main() { + for _ in 0..2 { + let task_number = 3; + let task2_string = String::from("Hello from task 3!"); + println!("Message from task{}: {}", task_number, task2_string); + yield_cooperative_task!(); + } +} diff --git a/src/dos_tests/mod.rs b/src/dos_tests/mod.rs index 990c346..b3b49e4 100644 --- a/src/dos_tests/mod.rs +++ b/src/dos_tests/mod.rs @@ -1,2 +1,3 @@ pub(crate) mod file; -pub(crate) mod allocator_test; \ No newline at end of file +pub(crate) mod allocator_test; +pub(crate) mod cooperative_multitasking_test; \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 4be4f11..237bb28 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,12 +8,14 @@ extern crate rlibc; extern crate alloc; use crate::dos::allocator::GLOBAL_ALLOCATOR; +use crate::dos::cooperative_multitasking::TASKING; #[link_section = ".startup"] #[no_mangle] fn _start() -> ! { unsafe { GLOBAL_ALLOCATOR.init(); + TASKING.init(); // Relies on the allocator } extern "Rust" { fn main() -> (); diff --git a/src/main.rs b/src/main.rs index 3b2da1f..aab100d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,18 +1,19 @@ #![no_std] #![no_main] - extern crate alloc; mod dos_tests; use rust_dos::*; -use crate::dos_tests::allocator_test::allocator_test; -use crate::dos_tests::file::file_read_test; +use crate::dos_tests::cooperative_multitasking_test::cooperative_multitasking_test; +//use crate::dos_tests::allocator_test::allocator_test; +//use crate::dos_tests::file::file_read_test; entry!(main); fn main() { - allocator_test(); - file_read_test(); -} + /*allocator_test(); + file_read_test();*/ + cooperative_multitasking_test(); +} \ No newline at end of file