tasks.rs

// SPDX-License-Identifier: MIT OR Apache-2.0
//
// Copyright (c) 2022-2023 SUSE LLC
//
// Author: Roy Hopkins <rhopkins@suse.de>

extern crate alloc;

use alloc::collections::btree_map::BTreeMap;
use alloc::string::String;
use alloc::sync::Arc;
use core::fmt;
use core::mem::offset_of;
use core::mem::size_of;
use core::num::NonZeroUsize;
use core::sync::atomic::AtomicBool;
use core::sync::atomic::AtomicU32;
use core::sync::atomic::AtomicUsize;
use core::sync::atomic::Ordering;

use crate::address::{Address, VirtAddr};
use crate::cpu::IrqGuard;
use crate::cpu::ShadowStackInit;
use crate::cpu::X86ExceptionContext;
use crate::cpu::features::{Feature, cpu_has_feat};
use crate::cpu::idt::svsm::return_new_task;
use crate::cpu::irq_state::EFLAGS_IF;
use crate::cpu::irqs_enable;
use crate::cpu::irqs_enabled;
use crate::cpu::percpu::PerCpu;
use crate::cpu::percpu::current_task;
use crate::cpu::percpu::this_cpu;
use crate::cpu::shadow_stack::init_shadow_stack;
use crate::cpu::sse::sse_restore_context;
use crate::cpu::sse::xsave_area_size;
use crate::error::SvsmError;
use crate::fs::{Directory, FileHandle, opendir, stdout_open};
use crate::locking::RWLock;
use crate::locking::SpinLock;
use crate::locking::SpinLockIrqSafe;
use crate::mm::pagetable::{PTEntryFlags, PageTable};
use crate::mm::vm::{Mapping, VMFileMappingFlags, VMKernelStack, VMR};
use crate::mm::{
    PageBox, SVSM_PERTASK_BASE, SVSM_PERTASK_END, USER_MEM_END, USER_MEM_START, VMMappingGuard,
    mappings::create_anon_mapping, mappings::create_file_mapping,
};
use crate::syscall::{Obj, ObjError, ObjHandle};
use crate::types::{SVSM_USER_CS, SVSM_USER_DS};
use crate::utils::{MemoryRegion, is_aligned};
use intrusive_collections::{LinkedList, LinkedListAtomicLink, intrusive_adapter};

use super::WaitQueue;
use super::schedule::complete_task_switch;
use super::schedule::terminate;
use super::task_mm::{TaskKernelMapping, TaskMM};

pub const INITIAL_TASK_ID: u32 = 1;

/// This trait is used to describe a type that can be used as the start
/// parameter for a kernel thread.  The thread start path requires that all
/// start paremeters be able to fit into a single register, which constrains
/// them to `usize`.  Other types can be used as long as there is an
/// appropriate mechanism to covert those types to and from the `usize` that
/// must be passed via the parameter register.  This trait defines such a
/// conversion mechanism.
///
/// # Safety
///
/// Implementations of this thread must ensure that the `to_usize` and
/// `from_usize` methods correctly preserve the safety of the parameter as it
/// transitions from its native type to and from a temporary `usize`
/// representation.
pub unsafe trait KernelThreadStartParameter {
    fn to_usize(parameter: Self) -> usize;

    /// # Safety
    /// This function must only be called with a `usize` parameter that was
    /// previously returned from a call to `to_usize()`.
    unsafe fn from_usize(parameter: usize) -> Self;
}

// SAFETY: types that implement `From` for conversions to and from `usize` are
// trivially convertible to and from `usize` and thus will always preserve
// safety of the parameter during conversion.
unsafe impl<T> KernelThreadStartParameter for T
where
    T: From<usize>,
    usize: From<T>,
{
    fn to_usize(parameter: Self) -> usize {
        usize::from(parameter)
    }

    unsafe fn from_usize(parameter: usize) -> Self {
        Self::from(parameter)
    }
}

#[derive(Debug)]
pub struct KernelThreadStartInfo {
    entry: usize,
    start_parameter: usize,
    start_routine: usize,
}

impl KernelThreadStartInfo {
    pub fn new<T: KernelThreadStartParameter>(entry: fn(T), start_parameter: T) -> Self {
        Self {
            entry: entry as usize,
            start_parameter: T::to_usize(start_parameter),
            start_routine: run_kernel_task::<T> as *const () as usize,
        }
    }

    /// # Safety
    /// The caller is required to provide a start parameter that is appropriate
    /// to the entry point according to the safety requirements of the entry
    /// point.
    pub unsafe fn new_unsafe(entry: unsafe fn(usize), start_parameter: usize) -> Self {
        Self {
            entry: entry as usize,
            start_parameter,
            start_routine: run_kernel_task::<usize> as *const () as usize,
        }
    }
}

#[derive(Debug)]
enum ThreadStartInfo {
    Kernel(KernelThreadStartInfo),
    User(usize),
}

#[derive(PartialEq, Debug, Copy, Clone, Default)]
#[repr(u32)]
pub enum TaskState {
    #[default]
    PENDING,
    RUNNING,
    BLOCKED,
    TERMINATED,
}

impl From<u32> for TaskState {
    fn from(v: u32) -> Self {
        match v {
            x if x == Self::RUNNING as u32 => Self::RUNNING,
            x if x == Self::BLOCKED as u32 => Self::BLOCKED,
            x if x == Self::TERMINATED as u32 => Self::TERMINATED,
            _ => Self::PENDING,
        }
    }
}

impl From<TaskState> for u32 {
    fn from(s: TaskState) -> u32 {
        s as u32
    }
}

#[derive(PartialEq, Debug, Copy, Clone)]
pub enum TaskExitStatus {
    /// Task called sys_exit(code).
    Exited(u16),
    /// Task was killed by an x86 exception.
    Exception,
}

impl From<u32> for TaskExitStatus {
    fn from(v: u32) -> Self {
        let (code, exception) = ((v >> 16) as u16, (v & 0xFFFF) as u16);
        if exception != 0 {
            Self::Exception
        } else {
            Self::Exited(code)
        }
    }
}

impl From<TaskExitStatus> for u32 {
    fn from(s: TaskExitStatus) -> u32 {
        match s {
            TaskExitStatus::Exited(code) => (code as u32) << 16,
            TaskExitStatus::Exception => 1,
        }
    }
}

impl Default for TaskExitStatus {
    fn default() -> Self {
        Self::Exited(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub enum TaskError {
    // Attempt to close or to retrieve the exit status of a non-terminated task
    NotTerminated,
    // A closed task could not be removed from the task list
    CloseFailed,
}

impl From<TaskError> for SvsmError {
    fn from(e: TaskError) -> Self {
        Self::Task(e)
    }
}

pub const TASK_FLAG_SHARE_PT: u16 = 0x01;

#[derive(Debug, Default)]
struct TaskIDAllocator {
    next_id: AtomicU32,
}

impl TaskIDAllocator {
    const fn new() -> Self {
        Self {
            next_id: AtomicU32::new(INITIAL_TASK_ID + 1),
        }
    }

    fn next_id(&self) -> u32 {
        let mut id = self.next_id.fetch_add(1, Ordering::Relaxed);
        // Reserve IDs of 0 and 1
        while (id == 0_u32) || (id == INITIAL_TASK_ID) {
            id = self.next_id.fetch_add(1, Ordering::Relaxed);
        }
        id
    }
}

static TASK_ID_ALLOCATOR: TaskIDAllocator = TaskIDAllocator::new();

#[repr(C)]
#[derive(Default, Debug, Clone, Copy)]
pub struct X86TaskSwitchRegs {
    // The context switch structure only needs to allocate enough space for
    // callee-save registers and any registers that are used as argument
    // registers for task start routines (such as for run_kernel_task).
    //
    // Argument registers come first.
    pub rsi: usize,
    pub rdx: usize,
    pub rcx: usize,

    // Callee-save registers come next.
    pub r12: usize,
    pub r13: usize,
    pub r14: usize,
    pub r15: usize,
    pub rbx: usize,
    pub rbp: usize,
}

#[repr(C)]
#[derive(Default, Debug, Clone, Copy)]
pub struct TaskContext {
    pub regs: X86TaskSwitchRegs,
    pub ret_addr: u64,
}

#[repr(C)]
struct TaskSchedState {
    /// Whether this is an idle task
    idle_task: AtomicBool,

    /// Whether this task is currently active on any CPU
    active: AtomicBool,

    /// Current state of the task
    state: AtomicU32,

    /// CPU this task is currently assigned to
    cpu_index: AtomicUsize,
}

impl TaskSchedState {
    fn new(cpu_index: usize) -> Self {
        Self {
            idle_task: AtomicBool::new(false),
            active: AtomicBool::new(false),
            state: AtomicU32::new(TaskState::PENDING.into()),
            cpu_index: AtomicUsize::new(cpu_index),
        }
    }

    fn panic_on_idle(&self, msg: &str) -> &Self {
        if self.idle_task.load(Ordering::Relaxed) {
            panic!("{}", msg);
        }
        self
    }

    fn set_active(&self) {
        if self.active.swap(true, Ordering::Release) {
            panic!("attempted switch to an active task");
        }
    }

    fn set_state(&self, state: TaskState) {
        self.state.store(state.into(), Ordering::Release);
    }

    fn get_state(&self) -> TaskState {
        self.state.load(Ordering::Acquire).into()
    }
}

pub struct Task {
    pub rsp: u64,

    pub ssp: VirtAddr,

    /// XSave area
    pub xsa: PageBox<[u8]>,

    pub stack_bounds: MemoryRegion<VirtAddr>,

    pub shadow_stack_base: VirtAddr,

    /// Page table that is loaded when the task is scheduled
    pub page_table: SpinLock<PageBox<PageTable>>,

    /// Task kernel stack mapping
    _kernel_stack: TaskKernelMapping,

    /// Task shadow stack mapping
    _shadow_stack: Option<TaskKernelMapping>,

    /// Task memory management state
    mm: Arc<TaskMM>,

    /// State relevant for scheduler
    sched_state: TaskSchedState,

    /// User-visible name of task
    name: String,

    /// ID of the task
    id: u32,

    /// Root directory for this task
    rootdir: Arc<dyn Directory>,

    /// Link to global task list
    list_link: LinkedListAtomicLink,

    /// Link to scheduler run queue
    runlist_link: LinkedListAtomicLink,

    /// Link to wait queue
    waitlist_link: LinkedListAtomicLink,

    /// Objects shared among threads within the same process
    objs: Arc<RWLock<BTreeMap<ObjHandle, Arc<dyn Obj>>>>,

    /// Queue of tasks waiting for completion of this task.
    wait_queue: SpinLockIrqSafe<WaitQueue>,

    /// Exit status of the task
    exit_status: AtomicU32,
}

// Expose the offsets of critical task fields to assembly.
pub const TASK_ACTIVE_OFFSET: usize =
    offset_of!(Task, sched_state) + offset_of!(TaskSchedState, active);
pub const TASK_CUR_CPU_OFFSET: usize =
    offset_of!(Task, sched_state) + offset_of!(TaskSchedState, cpu_index);

// SAFETY: Send + Sync is required for Arc<Task> to implement Send. All members
// of  `Task` are Send + Sync except for the intrusive_collection links, which
// are only Send. The only access to these is via the intrusive_adapter!
// generated code which does not use them concurrently across threads.  The
// kernal address cell is also not Sync, but this is only populated during
// task creation, and can safely be accessed by multiple threads once it has
// been populated.
unsafe impl Sync for Task {}

pub type TaskPointer = Arc<Task>;

intrusive_adapter!(pub TaskRunListAdapter = TaskPointer: Task { runlist_link: LinkedListAtomicLink });
intrusive_adapter!(pub TaskListAdapter = TaskPointer: Task { list_link: LinkedListAtomicLink });
intrusive_adapter!(pub TaskWaitListAdapter = TaskPointer: Task { waitlist_link: LinkedListAtomicLink });

impl PartialEq for Task {
    fn eq(&self, other: &Self) -> bool {
        core::ptr::eq(self, other)
    }
}

impl fmt::Debug for Task {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Task")
            .field("rsp", &self.rsp)
            .field("state", &self.sched_state.get_state())
            .field("id", &self.id)
            .finish()
    }
}

impl Drop for Task {
    fn drop(&mut self) {
        // A task must be inactive to be terminated.  Otherwise, its stack
        // might be freed while it is actively executing.
        assert!(!self.sched_state.active.load(Ordering::Relaxed));
        // Check that the task state is correct and that interrupts are
        // enabled.  These don't affect memory safety so they are debug-only,
        // but they are good sanity checks on the scheduler.
        debug_assert!(self.is_terminated() || self.is_pending());
        debug_assert!(irqs_enabled());
    }
}

struct CreateTaskArguments {
    // The thread start information.
    start_info: ThreadStartInfo,

    // The name of the task.
    name: String,

    // For a user task, supplies the `VMR` that will represent the user-mode
    // address space.
    vm_user_range: Option<VMR>,

    // The root directory that will be associated with this task.
    rootdir: Arc<dyn Directory>,

    // Share state with another task (aka create a thread)
    thread_of: Option<TaskPointer>,
}

impl Task {
    fn create_common(cpu: &PerCpu, args: CreateTaskArguments) -> Result<TaskPointer, SvsmError> {
        let mut pgtable = cpu.get_pgtable().clone_shared()?;

        cpu.populate_page_table(&mut pgtable);

        let (task_mm, objtree) = {
            if let Some(parent_thread) = args.thread_of {
                (parent_thread.mm.clone(), parent_thread.objs.clone())
            } else {
                (
                    Arc::new(TaskMM::create(args.vm_user_range)?),
                    Arc::new(RWLock::new(BTreeMap::new())),
                )
            }
        };

        let xsa = Self::allocate_xsave_area();
        let xsa_addr = u64::from(xsa.vaddr()) as usize;

        // Determine which kernel-mode entry/exit routines will be used for
        // this task.
        let (entry_return, exit_return) = match args.start_info {
            ThreadStartInfo::User(_) => (return_new_task as *const () as usize, None),
            ThreadStartInfo::Kernel(ref info) => {
                (info.start_routine, Some(task_exit as *const () as usize))
            }
        };

        let mut shadow_stack_offset = VirtAddr::null();
        let mut shadow_stack_base = VirtAddr::null();
        let shadow_stack_mapping = if cpu_has_feat(Feature::CetSS) {
            // Allocate shadow stack and safe top_of_stack offset
            let shadow_stack = VMKernelStack::new_shadow()?;
            let offset = shadow_stack.top_of_stack();
            let shadow_stack_page = shadow_stack.shadow_page();

            let base_token_addr;

            // Map shadow stack into virtual address range
            let mapping = TaskKernelMapping::new(task_mm.clone(), Arc::new(shadow_stack))?;
            let stack_base = mapping.virt_addr();

            // Initialize shadow stack
            (base_token_addr, shadow_stack_offset) = init_shadow_stack(
                &shadow_stack_page,
                stack_base + offset,
                ShadowStackInit::Normal {
                    entry_return,
                    exit_return,
                },
            );

            if let Some(base_addr) = base_token_addr {
                shadow_stack_base = base_addr;
            }

            Some(mapping)
        } else {
            None
        };

        // Call the correct stack creation routine for this task.
        let (stack, raw_bounds, rsp_offset) = match args.start_info {
            ThreadStartInfo::User(entry) => Self::allocate_utask_stack(cpu, entry, xsa_addr)?,
            ThreadStartInfo::Kernel(ref info) => Self::allocate_ktask_stack(
                cpu,
                info.entry,
                info.start_routine,
                xsa_addr,
                info.start_parameter,
            )?,
        };
        let kernel_stack_mapping = TaskKernelMapping::new(task_mm.clone(), stack)?;
        let stack_start = kernel_stack_mapping.virt_addr();

        task_mm.kernel_range().populate(&mut pgtable);

        // Remap at the per-task offset
        let bounds = MemoryRegion::new(stack_start + raw_bounds.start().into(), raw_bounds.len());
        // Stack frames should be 16b-aligned
        debug_assert!(bounds.end().is_aligned(16));

        Ok(Arc::new(Task {
            rsp: bounds
                .end()
                .checked_sub(rsp_offset)
                .expect("Invalid stack offset from task stack allocator")
                .bits() as u64,
            ssp: shadow_stack_offset,
            xsa,
            stack_bounds: bounds,
            shadow_stack_base,
            page_table: SpinLock::new(pgtable),
            _kernel_stack: kernel_stack_mapping,
            _shadow_stack: shadow_stack_mapping,
            mm: task_mm,
            sched_state: TaskSchedState::new(cpu.get_cpu_index()),
            name: args.name,
            id: TASK_ID_ALLOCATOR.next_id(),
            rootdir: args.rootdir,
            list_link: LinkedListAtomicLink::default(),
            runlist_link: LinkedListAtomicLink::default(),
            waitlist_link: LinkedListAtomicLink::default(),
            objs: objtree,
            wait_queue: SpinLockIrqSafe::new(WaitQueue::new()),
            exit_status: AtomicU32::new(TaskExitStatus::default().into()),
        }))
    }

    pub fn create(
        cpu: &PerCpu,
        start_info: KernelThreadStartInfo,
        name: String,
    ) -> Result<TaskPointer, SvsmError> {
        let create_args = CreateTaskArguments {
            start_info: ThreadStartInfo::Kernel(start_info),
            name,
            vm_user_range: None,
            rootdir: opendir("/")?,
            thread_of: None,
        };
        Self::create_common(cpu, create_args)
    }

    pub fn create_user(
        cpu: &PerCpu,
        user_entry: usize,
        root: Arc<dyn Directory>,
        name: String,
    ) -> Result<TaskPointer, SvsmError> {
        let vm_user_range = VMR::new(USER_MEM_START, USER_MEM_END, PTEntryFlags::USER)?;
        // SAFETY: the user address range is fully aligned to top-level paging
        // boundaries.
        unsafe {
            vm_user_range.initialize_lazy()?;
        }
        let create_args = CreateTaskArguments {
            start_info: ThreadStartInfo::User(user_entry),
            name,
            vm_user_range: Some(vm_user_range),
            rootdir: root,
            thread_of: None,
        };
        Self::create_common(cpu, create_args)
    }

    /// Create a new thread for an existing task.
    ///
    /// # Arguments
    ///
    /// * `cpu` - Reference to the `[PerCpu]` structure of the local CPU.
    /// * `entry` - Pointer the function to run in the thread.
    /// * `start_parameter` - `usize` to pass as parameter when calling `entry`.
    /// * `name` - User-visible name of the thread.
    /// * `thread_of` - Pointer to existing task for which the new thread will
    ///   be created.
    ///
    /// # Returns
    ///
    /// `Some(TaskPointer)` with the new thread on success, `Err(SvsmError)` on failure.
    pub fn create_thread(
        cpu: &PerCpu,
        start_info: KernelThreadStartInfo,
        name: String,
        thread: TaskPointer,
    ) -> Result<TaskPointer, SvsmError> {
        let create_args = CreateTaskArguments {
            start_info: ThreadStartInfo::Kernel(start_info),
            name,
            vm_user_range: None,
            rootdir: opendir("/")?,
            thread_of: Some(thread),
        };
        Self::create_common(cpu, create_args)
    }

    pub fn stack_bounds(&self) -> MemoryRegion<VirtAddr> {
        self.stack_bounds
    }

    pub fn get_task_name(&self) -> &String {
        &self.name
    }

    pub fn get_task_id(&self) -> u32 {
        self.id
    }

    pub fn rootdir(&self) -> Arc<dyn Directory> {
        self.rootdir.clone()
    }

    pub fn set_task_active(&self) {
        self.sched_state.set_active();
    }

    pub fn set_task_running(&self) {
        self.sched_state.set_state(TaskState::RUNNING);
    }

    pub fn set_task_terminated(&self) -> LinkedList<TaskWaitListAdapter> {
        self.sched_state
            .panic_on_idle("Trying to terminate idle task")
            .set_state(TaskState::TERMINATED);
        self.wait_queue.lock().wakeup(true)
    }

    pub fn set_task_blocked(&self) {
        self.sched_state
            .panic_on_idle("Trying to block idle task")
            .set_state(TaskState::BLOCKED);
    }

    pub fn is_running(&self) -> bool {
        self.sched_state.get_state() == TaskState::RUNNING
    }

    pub fn is_pending(&self) -> bool {
        self.sched_state.get_state() == TaskState::PENDING
    }

    pub fn is_terminated(&self) -> bool {
        self.sched_state.get_state() == TaskState::TERMINATED
    }

    pub fn set_exit_status(&self, exit_status: TaskExitStatus) {
        assert!(
            !self.is_terminated(),
            "The exit status can only be set for a task that is not yet terminated"
        );
        // Use Relaxed order since the exit status is set before the task
        // is marked terminated with Release order
        self.exit_status
            .store(exit_status.into(), Ordering::Relaxed);
    }

    pub fn get_exit_status(&self) -> Result<TaskExitStatus, SvsmError> {
        if self.is_terminated() {
            // Use Relaxed order since the exit status is read after the task
            // state is checked with Acquire order
            Ok(self.exit_status.load(Ordering::Relaxed).into())
        } else {
            Err(TaskError::NotTerminated.into())
        }
    }

    pub fn set_idle_task(&self) {
        self.sched_state.idle_task.store(true, Ordering::Relaxed);
    }

    pub fn is_idle_task(&self) -> bool {
        self.sched_state.idle_task.load(Ordering::Relaxed)
    }

    pub fn fault(&self, vaddr: VirtAddr, write: bool) -> Result<(), SvsmError> {
        let vmr = self
            .mm
            .user_range()
            .filter(|vmr| vmr.virt_range().contains(vaddr))
            .ok_or(SvsmError::Mem)?;
        let mut pgtbl = self.page_table.lock();
        vmr.handle_page_fault(&mut pgtbl, vaddr, write)?;
        Ok(())
    }

    fn allocate_stack_common() -> Result<(Mapping, MemoryRegion<VirtAddr>), SvsmError> {
        let stack = VMKernelStack::new()?;
        let bounds = stack.bounds(VirtAddr::from(0u64));

        let mapping = Arc::new(stack);

        Ok((mapping, bounds))
    }

    fn allocate_ktask_stack(
        cpu: &PerCpu,
        entry: usize,
        start_routine: usize,
        xsa_addr: usize,
        start_parameter: usize,
    ) -> Result<(Mapping, MemoryRegion<VirtAddr>, usize), SvsmError> {
        let (mapping, bounds) = Task::allocate_stack_common()?;

        let percpu_mapping = cpu.new_mapping(mapping.clone())?;

        // We need to setup a context on the stack that matches the stack layout
        // defined in switch_context below.
        let stack_tos = percpu_mapping.virt_addr() + bounds.end().bits();
        // Make space for the task termination handler
        let stack_offset = size_of::<u64>();
        let stack_ptr = stack_tos
            .checked_sub(stack_offset)
            .unwrap()
            .as_mut_ptr::<u8>();
        // To ensure stack frames are 16b-aligned, ret_addr must be 16b-aligned
        // so that (%rsp + 8) is 16b-aligned after the ret instruction in
        // switch_context
        debug_assert!(
            VirtAddr::from(stack_ptr)
                .checked_sub(8)
                .unwrap()
                .is_aligned(16)
        );
        // Make sure there is room for TaskContext
        debug_assert!(
            (percpu_mapping.virt_addr() + bounds.start().bits())
                <= VirtAddr::from(stack_ptr)
                    .checked_sub(size_of::<TaskContext>())
                    .unwrap()
        );

        // 'Push' the task frame onto the stack
        //
        // SAFETY: we ensure that both `TaskContext` and the function pointer
        // can be written to valid memory. The address storing the function
        // pointer is always 8b-aligned.
        unsafe {
            let task_context = stack_ptr
                .sub(size_of::<TaskContext>())
                .cast::<TaskContext>();
            // ret_addr
            (*task_context).regs.rsi = entry;
            // xsave area addr
            (*task_context).regs.rdx = xsa_addr;
            // start argument parameter.
            (*task_context).regs.rcx = start_parameter;
            (*task_context).ret_addr = start_routine as u64;
            // Task termination handler for when entry point returns
            stack_ptr.cast::<u64>().write(task_exit as *const () as u64);
        }

        Ok((mapping, bounds, stack_offset + size_of::<TaskContext>()))
    }

    fn allocate_utask_stack(
        cpu: &PerCpu,
        user_entry: usize,
        xsa_addr: usize,
    ) -> Result<(Mapping, MemoryRegion<VirtAddr>, usize), SvsmError> {
        let (mapping, bounds) = Task::allocate_stack_common()?;
        let iret_rflags: usize = 2 | EFLAGS_IF;

        let percpu_mapping = cpu.new_mapping(mapping.clone())?;

        // We need to setup a context on the stack that matches the stack layout
        // defined in switch_context below.
        let stack_tos = percpu_mapping.virt_addr() + bounds.end().bits();
        // Make space for the IRET frame
        let stack_offset = size_of::<X86ExceptionContext>();
        let stack_ptr = stack_tos
            .checked_sub(stack_offset)
            .unwrap()
            .as_mut_ptr::<u8>();
        // return_new_task is asm code, so %rsp must be 16b-aligned after the
        // ret instruction in switch_context to ensure stack frames are
        // 16b-aligned
        debug_assert!(VirtAddr::from(stack_ptr).is_aligned(16));
        // Make sure there is room for TaskContext
        debug_assert!(
            (percpu_mapping.virt_addr() + bounds.start().bits())
                <= VirtAddr::from(stack_ptr)
                    .checked_sub(size_of::<TaskContext>())
                    .unwrap()
        );

        // 'Push' the task frame onto the stack
        //
        // SAFETY: we ensure that both `X86ExceptionContext` and `TaskContext`
        // can be written to valid memory.
        unsafe {
            // Setup IRQ return frame.  User-mode tasks always run with
            // interrupts enabled.
            let mut iret_frame = X86ExceptionContext::default();
            iret_frame.frame.rip = user_entry;
            iret_frame.frame.cs = (SVSM_USER_CS | 3).into();
            iret_frame.frame.flags = iret_rflags;
            iret_frame.frame.rsp = (USER_MEM_END - 8).into();
            iret_frame.frame.ss = (SVSM_USER_DS | 3).into();
            debug_assert!(is_aligned(iret_frame.frame.rsp + 8, 16));

            // Copy IRET frame to stack
            let stack_iret_frame = stack_ptr.cast::<X86ExceptionContext>();
            *stack_iret_frame = iret_frame;

            let task_context = TaskContext {
                regs: X86TaskSwitchRegs {
                    rsi: xsa_addr, // XSAVE area addr
                    ..Default::default()
                },
                ret_addr: VirtAddr::from(return_new_task as *const ())
                    .bits()
                    .try_into()
                    .unwrap(),
            };

            let stack_task_context = stack_ptr
                .sub(size_of::<TaskContext>())
                .cast::<TaskContext>();
            *stack_task_context = task_context;
        }

        Ok((mapping, bounds, stack_offset + size_of::<TaskContext>()))
    }

    fn allocate_xsave_area() -> PageBox<[u8]> {
        let len = xsave_area_size() as usize;
        let xsa = PageBox::<[u8]>::try_new_slice(0u8, NonZeroUsize::new(len).unwrap());
        if xsa.is_err() {
            panic!("Error while allocating xsave area");
        }
        xsa.unwrap()
    }

    pub fn wait_for_exit(&self) -> Option<IrqGuard> {
        let current_task = this_cpu().current_task();

        // Lock the wait queue before examining the current state.  This must
        // be done with interrupts disabled, since once the current task has
        // joined the wait queue and the wait queue is unlocked, the task is
        // subject to an immediate attempt to wake, and therefore the current
        // task must be descheduled as quickly as possible to prevent
        // contention on the waking CPU.
        let guard = IrqGuard::new();

        // Determine whether this task has already terminated, and only join
        // the wait queue if the task has not already terminated.  Since a task
        // is marked terminated before its wait queue is examined, this
        // sequence guarantees that a task will not block unless it is
        // guaranteed to be woken.
        let mut wait_queue = self.wait_queue.lock();
        if self.is_terminated() {
            return None;
        }

        // Join the wait queue of the target task.
        wait_queue.wait_for_event(current_task);
        drop(wait_queue);

        // Return the IRQ guard as an indication that the caller must wait.
        Some(guard)
    }

    pub fn mmap_common(
        vmr: &VMR,
        addr: VirtAddr,
        file: Option<&FileHandle>,
        offset: usize,
        size: usize,
        flags: VMFileMappingFlags,
    ) -> Result<VirtAddr, SvsmError> {
        let mapping = if let Some(f) = file {
            create_file_mapping(f, offset, size, flags)?
        } else {
            create_anon_mapping(size, flags)?
        };

        if flags.contains(VMFileMappingFlags::Fixed) {
            Ok(vmr.insert_at(addr, mapping)?)
        } else {
            Ok(vmr.insert_hint(addr, mapping)?)
        }
    }

    pub fn mmap_kernel(
        &self,
        addr: VirtAddr,
        file: Option<&FileHandle>,
        offset: usize,
        size: usize,
        flags: VMFileMappingFlags,
    ) -> Result<VirtAddr, SvsmError> {
        Self::mmap_common(self.mm.kernel_range(), addr, file, offset, size, flags)
    }

    pub fn mmap_kernel_guard<'a>(
        &'a self,
        addr: VirtAddr,
        file: Option<&FileHandle>,
        offset: usize,
        size: usize,
        flags: VMFileMappingFlags,
    ) -> Result<VMMappingGuard<'a>, SvsmError> {
        let vaddr = Self::mmap_common(self.mm.kernel_range(), addr, file, offset, size, flags)?;
        Ok(VMMappingGuard::new(self.mm.kernel_range(), vaddr))
    }

    pub fn mmap_user(
        &self,
        addr: VirtAddr,
        file: Option<&FileHandle>,
        offset: usize,
        size: usize,
        flags: VMFileMappingFlags,
    ) -> Result<VirtAddr, SvsmError> {
        let vmr = self.mm.user_range().ok_or(SvsmError::Mem)?;
        Self::mmap_common(vmr, addr, file, offset, size, flags)
    }

    pub fn munmap_kernel(&self, addr: VirtAddr) -> Result<(), SvsmError> {
        self.mm.kernel_range().remove(addr)?;
        Ok(())
    }

    pub fn munmap_user(&self, addr: VirtAddr) -> Result<(), SvsmError> {
        let vmr = self.mm.user_range().ok_or(SvsmError::Mem)?;
        vmr.remove(addr).map(|_| ())
    }

    /// Adds an object to the current task.
    ///
    /// # Arguments
    ///
    /// * `obj` - The object to be added.
    ///
    /// # Returns
    ///
    /// * `Result<ObjHandle, SvsmError>` - Returns the object handle for the object
    ///   to be added if successful, or an `SvsmError` on failure.
    ///
    /// # Errors
    ///
    /// This function will return an error if allocating the object handle fails.
    pub fn add_obj(&self, obj: Arc<dyn Obj>) -> Result<ObjHandle, SvsmError> {
        let mut objs = self.objs.lock_write();
        let last_key = objs
            .keys()
            .last()
            .map_or(Some(0), |k| u32::from(*k).checked_add(1))
            .ok_or(SvsmError::from(ObjError::InvalidHandle))?;
        let id = ObjHandle::new(if last_key != objs.len() as u32 {
            objs.keys()
                .enumerate()
                .find(|&(i, key)| i as u32 != u32::from(*key))
                .unwrap()
                .0 as u32
        } else {
            last_key
        });

        objs.insert(id, obj);

        Ok(id)
    }

    /// Adds an object to the current task and maps it to a given object-id.
    ///
    /// # Arguments
    ///
    /// * `obj` - The object to be added.
    /// * `handle` - Object handle to reference the object.
    ///
    /// # Returns
    ///
    /// * `Result<ObjHandle, SvsmError>` - Returns the object handle for the object
    ///   to be added if successful, or an `SvsmError` on failure.
    ///
    /// # Errors
    ///
    /// This function will return an error if allocating the object handle
    /// fails or the object id is already in use.
    pub fn add_obj_at(&self, obj: Arc<dyn Obj>, handle: ObjHandle) -> Result<ObjHandle, SvsmError> {
        let mut objs = self.objs.lock_write();

        if objs.get(&handle).is_some() {
            return Err(SvsmError::from(ObjError::Busy));
        }

        objs.insert(handle, obj);

        Ok(handle)
    }

    /// Removes an object from the current task.
    ///
    /// # Arguments
    ///
    /// * `id` - The ObjHandle for the object to be removed.
    ///
    /// # Returns
    ///
    /// * `Result<Arc<dyn Obj>>, SvsmError>` - Returns the removed `Arc<dyn Obj>`
    ///   on success, or an `SvsmError` on failure.
    ///
    /// # Errors
    ///
    /// This function will return an error if the object handle id does not
    /// exist in the current task.
    pub fn remove_obj(&self, id: ObjHandle) -> Result<Arc<dyn Obj>, SvsmError> {
        self.objs
            .lock_write()
            .remove(&id)
            .ok_or(ObjError::NotFound.into())
    }

    /// Retrieves an object from the current task.
    ///
    /// # Arguments
    ///
    /// * `id` - The ObjHandle for the object to be retrieved.
    ///
    /// # Returns
    ///
    /// * `Result<Arc<dyn Obj>>, SvsmError>` - Returns the `Arc<dyn Obj>` on
    ///   success, or an `SvsmError` on failure.
    ///
    /// # Errors
    ///
    /// This function will return an error if the object handle id does not exist
    /// in the current task.
    pub fn get_obj(&self, id: ObjHandle) -> Result<Arc<dyn Obj>, SvsmError> {
        self.objs
            .lock_read()
            .get(&id)
            .cloned()
            .ok_or(ObjError::NotFound.into())
    }
}

pub fn is_task_fault(vaddr: VirtAddr) -> bool {
    (vaddr >= USER_MEM_START && vaddr < USER_MEM_END)
        || (vaddr >= SVSM_PERTASK_BASE && vaddr < SVSM_PERTASK_END)
}

fn task_attach_console() {
    let file_handle = stdout_open();
    let obj_handle = ObjHandle::new(0);
    current_task()
        .add_obj_at(file_handle, obj_handle)
        .expect("Failed to attach console");
}

/// Runs the first time a new task is scheduled, in the context of the new
/// task. Any first-time initialization and setup work for a new task that
/// needs to happen in its context must be done here.
/// # Safety
/// The caller is required to verify the correctness of the save area address.
#[unsafe(no_mangle)]
unsafe fn setup_user_task(prev: usize, xsa_addr: u64) {
    // SAFETY: caller needs to make sure that the previous task pointer and
    // xsa_addr are valid.
    unsafe {
        // Needs to be the first function called here.
        setup_new_task_common(xsa_addr, complete_task_switch(prev));
    }
    task_attach_console();
}

unsafe fn setup_new_task_common(xsa_addr: u64, prev_task: Option<TaskPointer>) {
    // Re-enable IRQs here, as they are still disabled from the
    // schedule()/sched_init() functions. After the context switch the IrqGuard
    // from the previous task is not dropped, which causes IRQs to stay
    // disabled in the new task.
    // This only needs to be done for the first time a task runs. Any
    // subsequent task switches will go through schedule() and there the guard
    // is dropped, re-enabling IRQs.
    irqs_enable();

    // Drop the previous task reference now that interrupts are reenabled.
    // This ensures that the task destructor will run with interrupts enabled.
    drop(prev_task);

    // SAFETY: The caller takes responsibility for the correctness of the save
    // area address.
    unsafe {
        sse_restore_context(xsa_addr);
    }
}

/// # Safety
/// This function must only be invoked as a thread start function, and must
/// be invoked according to the start parameter data type expected by the entry
/// point.
unsafe fn run_kernel_task<T: KernelThreadStartParameter>(
    prev: usize,
    entry: fn(T),
    xsa_addr: u64,
    start_parameter: usize,
) {
    // SAFETY: the scheduler passes the correct pointer to the previous task.
    let prev_task = unsafe { complete_task_switch(prev) };

    // SAFETY: the save area address is provided by the context switch assembly
    // code.
    unsafe {
        setup_new_task_common(xsa_addr, prev_task);
    }

    // SAFETY: the `usize` start parameter was generated from an earlier call
    // to `KernelThreadStartParameter::to_usize` when the start argument was
    // prepared for use in the initial stack context, and therefore the safety
    // requirements of `from_usize` are met.  The exception is if this task
    // starts via an unsafe start function, in which case safety is the
    // responsibility of whoever created the thread, and is not required to be
    // assured here.
    entry(unsafe { T::from_usize(start_parameter) });
}

fn task_exit() {
    terminate(None);
}

#[cfg(all(test, test_in_svsm))]
mod tests {
    extern crate alloc;
    use crate::task::{KernelThreadStartInfo, start_kernel_task, wait_for_termination};
    use alloc::string::String;
    use core::arch::asm;
    use core::arch::global_asm;

    #[test]
    #[cfg_attr(not(test_in_svsm), ignore = "Can only be run inside guest")]
    fn test_media_and_x87_instructions() {
        let ret: u64;
        unsafe {
            asm!("call test_fpu", out("rax") ret, options(att_syntax));
        }

        assert_eq!(ret, 0);
    }

    global_asm!(
        r#"
    .section .text
    test_fpu:
        movq $0x3ff, %rax
        shl $52, %rax
        // rax contains 1.0 in Double Precision FP representation
        movq %rax, %xmm1
        movapd %xmm1, %xmm3

        movq $0x400, %rax
        shl $52, %rax
        // rax contains 2.0 in Double Precision FP representation
        movq %rax, %xmm2

        divsd %xmm2, %xmm3
        movq $0, %rax
        ret
        "#,
        options(att_syntax)
    );

    global_asm!(
        r#"
    .section .text
    check_fpu:
        movq $1, %rax
        movq $0x3ff, %rbx
        shl $52, %rbx
        // rbx contains 1.0 in Double Precision FP representation
        movq %rbx, %xmm4
        movapd %xmm4, %xmm6
        comisd %xmm4, %xmm1
        jnz 1f

        movq $0x400, %rbx
        shl $52, %rbx
        // rbx contains 2.0 in Double Precision FP representation
        movq %rbx, %xmm5
        comisd %xmm5, %xmm2
        jnz 1f

        divsd %xmm5, %xmm6
        comisd %xmm6, %xmm3
        jnz 1f
        movq $0, %rax
    1:
        ret
        "#,
        options(att_syntax)
    );

    global_asm!(
        r#"
    .section .text
    alter_fpu:
        movq $0x400, %rax
        shl $52, %rax
        // rax contains 2.0 in Double Precision FP representation
        movq %rax, %xmm1
        movapd %xmm1, %xmm3

        movq $0x3ff, %rax
        shl $52, %rax
        // rax contains 1.0 in Double Precision FP representation
        movq %rax, %xmm2
        divsd %xmm3, %xmm2
        movq $0, %rax
        ret
        "#,
        options(att_syntax)
    );

    #[test]
    #[cfg_attr(not(test_in_svsm), ignore = "Can only be run inside guest")]
    fn test_fpu_context_switch() {
        let task1 = start_kernel_task(KernelThreadStartInfo::new(task1, 1), String::from("task1"))
            .expect("Failed to launch request processing task");

        wait_for_termination(task1);
    }

    fn task1(start_parameter: usize) {
        assert_eq!(start_parameter, 1);

        let ret: u64;
        unsafe {
            asm!("call test_fpu", options(att_syntax));
        }

        let task2 = start_kernel_task(KernelThreadStartInfo::new(task2, 2), String::from("task2"))
            .expect("Failed to launch request processing task");

        wait_for_termination(task2);

        unsafe {
            asm!("call check_fpu", out("rax") ret, options(att_syntax));
        }
        assert_eq!(ret, 0);
    }

    fn task2(start_parameter: usize) {
        assert_eq!(start_parameter, 2);
        unsafe {
            asm!("call alter_fpu", options(att_syntax));
        }
    }
}