k-os/crates/kernel/src/lib.rs

#![no_std]

pub mod interrupt;
#[cfg(kani)]
pub mod proofs;
pub mod syscall;

pub use syscall::{SyscallOutcome, SyscallResult};

#[cfg(test)]
#[path = "tests/mod.rs"]
mod tests;

use domain::{
    capability::{Capability, CapabilityTable},
    endpoint::EndpointTable,
    error::KernelError,
    irq::IrqRoutingTable,
    memory::{RegionTable, page_pool::PagePool, process_memory::ProcessMemory},
    message::Message,
    process::{MAX_PROCESSES, ProcessTable},
    resource::Resource,
    scheduler::{Scheduler, multi_core::MultiCoreScheduler},
    syscall::WaitSet,
    thread::{MAX_THREADS, ThreadState, ThreadTable},
};
use platform::Platform;
use types::{
    bounded::{CapSlot, Priority},
    generation::Generation,
    id::{ProcessId, ThreadId},
    permission::{Permission, PermissionSet},
    time::TimeSlice,
};

pub struct Kernel<P: Platform> {
    platform: P,
    processes: ProcessTable,
    threads: ThreadTable,
    endpoints: EndpointTable,
    scheduler: MultiCoreScheduler,
    page_pool: PagePool,
    process_memory: [Option<ProcessMemory>; MAX_PROCESSES],
    regions: RegionTable,
    irqs: IrqRoutingTable,
    pending_messages: [Option<Message>; MAX_THREADS],
    pending_results: [Option<SyscallResult>; MAX_THREADS],
    wait_sets: [Option<WaitSet>; MAX_THREADS],
}

impl<P: Platform> Kernel<P> {
    pub fn new(platform: P, page_pool: PagePool) -> Self {
        let num_cores = platform.num_cores();
        Self {
            platform,
            processes: ProcessTable::new(),
            threads: ThreadTable::new(),
            endpoints: EndpointTable::new(),
            scheduler: MultiCoreScheduler::new(num_cores),
            page_pool,
            process_memory: [const { None }; MAX_PROCESSES],
            regions: RegionTable::new(),
            irqs: IrqRoutingTable::new(),
            pending_messages: [const { None }; MAX_THREADS],
            pending_results: [const { None }; MAX_THREADS],
            wait_sets: [const { None }; MAX_THREADS],
        }
    }

    pub fn boot(
        &mut self,
        priority: Priority,
        time_slice: TimeSlice,
    ) -> Result<(ProcessId, ThreadId), KernelError> {
        let pid = self.processes.create(CapabilityTable::new(), None)?;
        let tid = self.threads.spawn(pid, priority, time_slice, None)?;
        self.processes.get_mut(pid)?.add_thread(tid)?;

        self.process_memory[pid.raw() as usize] =
            Some(ProcessMemory::new(pid, ProcessMemory::DEFAULT_PAGE_LIMIT));

        let cap = Capability::new(
            Resource::ProcessCreate,
            PermissionSet::empty().with(Permission::Execute),
            Generation::initial(),
        );
        self.processes
            .get_mut(pid)?
            .capabilities_mut()
            .insert(cap)?;

        let core = self.platform.core_id();
        self.scheduler.add(tid, priority, core);
        self.scheduler.next(core);
        self.threads.get_mut(tid)?.set_running(core)?;
        Ok((pid, tid))
    }

    pub fn platform_mut(&mut self) -> &mut P {
        &mut self.platform
    }

    pub fn pending_result(&self, thread: ThreadId) -> Option<&SyscallResult> {
        self.pending_results[thread.raw() as usize].as_ref()
    }

    pub fn take_pending_result(&mut self, thread: ThreadId) -> Option<SyscallResult> {
        self.pending_results[thread.raw() as usize].take()
    }

    pub fn set_thread_running(
        &mut self,
        tid: ThreadId,
        core: types::id::CoreId,
    ) -> Result<(), domain::error::KernelError> {
        self.threads.get_mut(tid)?.set_running(core)
    }

    pub fn free_page_count(&self) -> usize {
        self.page_pool.free_count()
    }

    pub fn process_exists(&self, pid: ProcessId) -> bool {
        self.processes.get(pid).is_ok()
    }

    pub fn thread_context(&self, tid: ThreadId) -> Option<&domain::thread::SavedContext> {
        self.threads.get(tid).ok().map(|t| t.saved_context())
    }

    pub fn thread_context_mut(
        &mut self,
        tid: ThreadId,
    ) -> Option<&mut domain::thread::SavedContext> {
        self.threads
            .get_mut(tid)
            .ok()
            .map(|t| t.saved_context_mut())
    }

    pub fn thread_process(&self, tid: ThreadId) -> Option<ProcessId> {
        self.threads.get(tid).ok().map(|t| t.process())
    }

    pub fn current_thread(&self) -> Option<ThreadId> {
        let core = self.platform.core_id();
        self.scheduler.current(core)
    }

    pub fn pick_next_thread(&mut self) -> Option<ThreadId> {
        use domain::scheduler::Scheduler;
        let core = self.platform.core_id();
        self.scheduler.next(core)
    }

    pub fn create_process_internal(&mut self) -> Option<ProcessId> {
        let pid = self.processes.create(CapabilityTable::new(), None).ok()?;
        self.process_memory[pid.raw() as usize] =
            Some(domain::memory::process_memory::ProcessMemory::new(
                pid,
                ProcessMemory::DEFAULT_PAGE_LIMIT,
            ));
        Some(pid)
    }

    pub fn spawn_thread_internal(&mut self, pid: ProcessId) -> Option<ThreadId> {
        use domain::scheduler::Scheduler;
        let tid = self
            .threads
            .spawn(pid, Priority::DEFAULT, TimeSlice::DEFAULT, None)
            .ok()?;
        let _ = self.processes.get_mut(pid).map(|p| p.add_thread(tid));
        let core = self.platform.core_id();
        self.scheduler.add(tid, Priority::DEFAULT, core);
        Some(tid)
    }

    pub fn create_endpoint(&mut self) -> Result<types::id::EndpointId, KernelError> {
        self.endpoints.create().ok_or(KernelError::NoFreeSlot)
    }

    pub fn grant_capability(
        &mut self,
        to: ProcessId,
        resource: Resource,
        permissions: PermissionSet,
    ) -> Result<CapSlot, KernelError> {
        self.grant_cap(to, resource, permissions)
    }

    pub fn switch_to_thread(&mut self, tid: ThreadId) {
        use domain::scheduler::Scheduler;
        let core = self.platform.core_id();

        if let Some(current) = self.scheduler.current(core) {
            if current == tid {
                return;
            }
            if let Ok(t) = self.threads.get_mut(current) {
                if matches!(t.state(), ThreadState::Running(_)) {
                    let priority = t.priority();
                    let _ = t.preempt();
                    self.scheduler.add(current, priority, core);
                }
            }
        }

        if let Ok(t) = self.threads.get(tid) {
            if matches!(t.state(), ThreadState::Ready) {
                self.scheduler.next(core);
                let _ = self.threads.get_mut(tid).map(|t| t.set_running(core));
            }
        }
    }

    fn caller(&self) -> Result<(ThreadId, ProcessId), KernelError> {
        let core = self.platform.core_id();
        let tid = self
            .scheduler
            .current(core)
            .ok_or(KernelError::WouldBlock)?;
        let thread = self.threads.get(tid)?;
        match thread.state() {
            ThreadState::Running(_) => Ok((tid, thread.process())),
            _ => Err(KernelError::WouldBlock),
        }
    }

    pub(crate) fn destroy_process(&mut self, pid: ProcessId) {
        use domain::process::MAX_THREADS_PER_PROCESS;

        let thread_ids: [Option<ThreadId>; MAX_THREADS_PER_PROCESS] = match self.processes.get(pid)
        {
            Ok(p) => {
                let mut ids = [None; MAX_THREADS_PER_PROCESS];
                for (i, tid) in p.thread_ids().enumerate() {
                    if i < MAX_THREADS_PER_PROCESS {
                        ids[i] = Some(tid);
                    }
                }
                ids
            }
            Err(_) => return,
        };

        self.cleanup_threads(&thread_ids);

        // Free allocated pages back to pool
        if let Some(pm) = &self.process_memory[pid.raw() as usize] {
            for &(pa, count) in pm.allocations() {
                self.page_pool.free(pa, count);
            }
        }

        // Unblock threads in other processes that were blocked waiting
        // on endpoints where this process's threads were the other side.
        // The cleanup_threads call above already removed this process's
        // threads from all endpoint wait queues. Threads blocked waiting
        // for this process will eventually time out or be rescheduled
        // when the endpoint is revoked.

        let _ = self.processes.destroy(pid);
        self.process_memory[pid.raw() as usize] = None;
    }

    pub(crate) fn cleanup_threads(&mut self, thread_ids: &[Option<ThreadId>]) {
        use domain::scheduler::Scheduler;
        for slot in thread_ids {
            if let Some(tid) = slot {
                self.endpoints.remove_thread_from_all(*tid);
                self.cleanup_wait_any(*tid);
                self.scheduler.remove(*tid);
                self.pending_messages[tid.raw() as usize] = None;
                self.pending_results[tid.raw() as usize] = None;
                let _ = self.threads.destroy(*tid);
            }
        }
    }

    pub(crate) fn cleanup_wait_any(&mut self, tid: ThreadId) {
        if let Some(wait_set) = self.wait_sets[tid.raw() as usize].take() {
            for ep_slot in wait_set.endpoints() {
                if let Some(_slot) = ep_slot {
                    self.endpoints.remove_thread_from_all(tid);
                    break;
                }
            }
        }
    }

    fn check_cap(
        &self,
        owner: ProcessId,
        slot: CapSlot,
        needed: Permission,
    ) -> Result<&Capability, KernelError> {
        let process = self.processes.get(owner)?;
        let cap = process.capabilities().get(slot)?;
        if !cap.permissions().allows(needed) {
            return Err(KernelError::PermissionDenied { slot, needed });
        }
        let current_gen = self.resource_generation(cap.resource());
        if !cap.generation().is_valid_against(&current_gen) {
            return Err(KernelError::RevokedCapability(slot));
        }
        Ok(cap)
    }

    fn resource_generation(&self, resource: &Resource) -> Generation {
        match resource {
            Resource::Endpoint(id) => self
                .endpoints
                .get(*id)
                .map(|ep| *ep.generation())
                .unwrap_or_else(|| Generation::initial()),
            Resource::Memory(id) => self
                .regions
                .get(*id)
                .map(|r| *r.generation())
                .unwrap_or_else(|| Generation::initial()),
            _ => Generation::initial(),
        }
    }

    pub fn revoke_resource(&mut self, resource: &Resource) {
        match resource {
            Resource::Endpoint(id) => {
                if let Some(ep) = self.endpoints.get_mut(*id) {
                    ep.revoke();
                }
            }
            Resource::Memory(id) => {
                if let Some(r) = self.regions.get_mut(*id) {
                    r.revoke();
                }
            }
            _ => {}
        }
    }

    fn grant_cap(
        &mut self,
        to: ProcessId,
        resource: Resource,
        permissions: PermissionSet,
    ) -> Result<CapSlot, KernelError> {
        let generation = self.resource_generation(&resource);
        let cap = Capability::new(resource, permissions, generation);
        self.processes.get_mut(to)?.capabilities_mut().insert(cap)
    }
}