kernel/task: Align kernel- and user-thread launch process

User-tasks use a different launch path than kernel-tasks. Kernel-tasks
go directly to a Rust function on the `ret` in `switch_context` while
for user-tasks the `ret`-target is an assembly function which calls
the `setup_user_task()` helper and then goes into the `iret` path.

This causes some fraction in how the initial stack layouts of both
task types are set up, so consolidate this to be more similar in both
cases.

Kernel- and User thread now use a common thread entry in assembly code,
which unifies the initial thread setup.

Signed-off-by: Joerg Roedel <joerg.roedel@amd.com>

Author: joergroedel

kernel/src/cpu/idt/svsm.rs

@@ -51,7 +51,7 @@ pub fn load_static_idt() {
 }
 
 unsafe extern "C" {
-    pub fn return_new_task();
+    pub fn thread_entry_asm();
     pub fn default_return();
     fn asm_entry_de();
     fn asm_entry_db();

kernel/src/cpu/idt/svsm_entry.S

@@ -369,10 +369,23 @@ return_user:
 	// Put user-mode specific return code here
 	jmp	return_all_paths
 
-.globl return_new_task
-return_new_task:
-	call	setup_user_task
-	jmp	default_return
+.globl thread_entry_asm
+thread_entry_asm:
+	// Call common thread setup function
+	call complete_new_thread
+
+	// Move thread fn and start parameter
+	movq %r15, %rsi
+	movq %r14, %rdi
+
+	// Fetch thread entry function pointer - do not touch stack pointer to
+	// keep stack alignment
+	movq (%rsp), %rax
+	call *%rax
+
+	// Skip to return addr
+	addq $8, %rsp
+	ret
 
 // #DE Divide-by-Zero-Error Exception (Vector 0)
 default_entry_no_ist	name=de		handler=terminate			error_code=0	vector=0

kernel/src/cpu/shadow_stack.rs

@@ -97,7 +97,9 @@ pub enum ShadowStackInit {
         /// The address of the first instruction that will be executed by the task.
         entry_return: usize,
         /// The address of the function that's executed when the task exits.
-        exit_return: Option<usize>,
+        exit_return: usize,
+        /// Whether there is an iret frame on the bottom of the stack.
+        iret_frame: bool,
     },
     /// A shadow stack to be used during context switches.
     ///
@@ -116,58 +118,63 @@ pub fn init_shadow_stack(
     init: ShadowStackInit,
 ) -> (Option<VirtAddr>, VirtAddr) {
     // Initialize the shadow stack.
-    let mut chunk = [0; 24];
-    let (base_token_addr, ssp) = match init {
+    let mut chunk = [0; 32];
+    let (base_token_addr, ssp, len) = match init {
         ShadowStackInit::Normal {
             entry_return,
             exit_return,
+            iret_frame,
         } => {
-            // If exit return is empty, then this thread will be used as a
-            // user task stack.  In that case, place a busy token at the
-            // base of the shadow stack.
             let base_token_addr = top_of_sstack - 8;
-            let base_token = match exit_return {
-                Some(addr) => addr,
-                None => base_token_addr.bits() + BUSY,
-            };
-
             let (token_bytes, rip_bytes) = chunk.split_at_mut(8);
 
+            let len: usize = if iret_frame { 32 } else { 24 };
+
             // Create a shadow stack restore token.
-            let token_addr = top_of_sstack - 24;
+            let token_addr = top_of_sstack - len;
             let token = (token_addr + 8).bits() + MODE_64BIT;
             token_bytes.copy_from_slice(&token.to_ne_bytes());
 
             let (entry_bytes, base_bytes) = rip_bytes.split_at_mut(8);
             entry_bytes.copy_from_slice(&entry_return.to_ne_bytes());
-            base_bytes.copy_from_slice(&base_token.to_ne_bytes());
 
-            (Some(base_token_addr), token_addr)
+            let (exit_bytes, bottom_bytes) = base_bytes.split_at_mut(8);
+            exit_bytes.copy_from_slice(&exit_return.to_ne_bytes());
+
+            if iret_frame {
+                // If iret_frame is true, then this thread will be used as a
+                // user task stack.  In that case, place a busy token at the
+                // base of the shadow stack.
+                let base_token = base_token_addr.bits() + BUSY;
+                bottom_bytes.copy_from_slice(&base_token.to_ne_bytes());
+            }
+
+            (Some(base_token_addr), token_addr, len)
         }
         ShadowStackInit::ContextSwitch => {
-            let (_, token_bytes) = chunk.split_at_mut(16);
+            let (_, token_bytes) = chunk.split_at_mut(24);
 
             // Create a shadow stack restore token.
             let token_addr = top_of_sstack - 8;
             let token = (token_addr + 8).bits() + MODE_64BIT;
             token_bytes.copy_from_slice(&token.to_ne_bytes());
 
-            (None, token_addr)
+            (None, token_addr, 32)
         }
         ShadowStackInit::Exception => {
-            let (_, token_bytes) = chunk.split_at_mut(16);
+            let (_, token_bytes) = chunk.split_at_mut(24);
 
             // Create a supervisor shadow stack token.
             let token_addr = top_of_sstack - 8;
             let token = token_addr.bits();
             token_bytes.copy_from_slice(&token.to_ne_bytes());
 
-            (None, token_addr)
+            (None, token_addr, 32)
         }
-        ShadowStackInit::Init => (None, top_of_sstack - 8),
+        ShadowStackInit::Init => (None, top_of_sstack - 8, 24),
     };
 
-    page.write(PAGE_SIZE - chunk.len(), &chunk);
+    page.write(PAGE_SIZE - len, &chunk[..len]);
 
     (base_token_addr, ssp)
 }

kernel/src/task/tasks.rs

@@ -22,7 +22,7 @@ use crate::address::{Address, VirtAddr};
 use crate::cpu::IrqGuard;
 use crate::cpu::ShadowStackInit;
 use crate::cpu::features::{Feature, cpu_has_feat};
-use crate::cpu::idt::svsm::return_new_task;
+use crate::cpu::idt::svsm::{default_return, thread_entry_asm};
 use crate::cpu::irq_state::EFLAGS_IF;
 use crate::cpu::irqs_enable;
 use crate::cpu::irqs_enabled;
@@ -231,29 +231,29 @@ pub struct TaskContext {
 #[derive(Default, Debug, Clone, Copy)]
 struct KernelTaskContext {
     pub task_context: TaskContext,
+    pub thread_routine: u64,
     pub exit_addr: u64,
 }
 
 #[repr(C)]
 #[derive(Default, Debug, Clone, Copy)]
 struct UserTaskContext {
     pub task_context: TaskContext,
+    pub thread_routine: u64,
+    pub exit_addr: u64,
     pub excp_context: X86ExceptionContext,
 }
 
 // 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
 const _: () = assert!(
-    (size_of::<KernelTaskContext>() - offset_of!(KernelTaskContext, task_context.ret_addr)) % 16
+    ((size_of::<KernelTaskContext>() - offset_of!(KernelTaskContext, task_context.ret_addr)) % 16)
+        - 8
         == 0
 );
-
-// User-task return to return_new_task which is asm code, so %rsp must be
-// 16b-aligned after the ret instruction in switch_context to ensure stack
-// frames are 16b-aligned
 const _: () = assert!(
-    (size_of::<UserTaskContext>() - offset_of!(UserTaskContext, task_context.ret_addr) - 8) % 16
+    ((size_of::<UserTaskContext>() - offset_of!(UserTaskContext, task_context.ret_addr)) % 16) - 8
         == 0
 );
 
@@ -442,11 +442,17 @@ impl Task {
 
         // 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 (entry_return, exit_return, iret_frame) = match args.start_info {
+            ThreadStartInfo::User(_) => (
+                thread_entry_asm as *const () as usize,
+                default_return as *const () as usize,
+                true,
+            ),
+            ThreadStartInfo::Kernel(_) => (
+                thread_entry_asm as *const () as usize,
+                task_exit as *const () as usize,
+                false,
+            ),
         };
 
         let mut shadow_stack_offset = VirtAddr::null();
@@ -470,6 +476,7 @@ impl Task {
                 ShadowStackInit::Normal {
                     entry_return,
                     exit_return,
+                    iret_frame,
                 },
             );
 
@@ -713,13 +720,15 @@ impl Task {
         let task_context = KernelTaskContext {
             task_context: TaskContext {
                 regs: X86TaskSwitchRegs {
-                    rsi: entry,
-                    rdx: xsa_addr,
-                    rcx: start_parameter,
+                    rsi: xsa_addr,
+                    // Put these into callee-safed registers
+                    r14: entry,
+                    r15: start_parameter,
                     ..Default::default()
                 },
-                ret_addr: start_routine as u64,
+                ret_addr: thread_entry_asm as *const () as u64,
             },
+            thread_routine: start_routine as u64,
             exit_addr: task_exit as *const () as u64,
         };
 
@@ -744,11 +753,13 @@ impl Task {
                     rsi: xsa_addr, // XSAVE area addr
                     ..Default::default()
                 },
-                ret_addr: VirtAddr::from(return_new_task as *const ())
-                    .bits()
-                    .try_into()
-                    .unwrap(),
+                ret_addr: thread_entry_asm as *const () as u64,
             },
+            thread_routine: run_user_task as *const () as u64,
+            exit_addr: VirtAddr::from(default_return as *const ())
+                .bits()
+                .try_into()
+                .unwrap(),
             // Setup IRQ return frame.  User-mode tasks always run with
             // interrupts enabled.
             excp_context: X86ExceptionContext {
@@ -995,23 +1006,8 @@ fn task_attach_console() {
         .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>) {
+unsafe fn complete_new_thread(prev: usize, xsa_addr: u64) {
     // 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
@@ -1021,6 +1017,9 @@ unsafe fn setup_new_task_common(xsa_addr: u64, prev_task: Option<TaskPointer>) {
     // is dropped, re-enabling IRQs.
     irqs_enable();
 
+    // SAFETY: the scheduler passes the correct pointer to the previous task.
+    let prev_task = unsafe { complete_task_switch(prev) };
+
     // Drop the previous task reference now that interrupts are reenabled.
     // This ensures that the task destructor will run with interrupts enabled.
     drop(prev_task);
@@ -1032,25 +1031,22 @@ unsafe fn setup_new_task_common(xsa_addr: u64, prev_task: Option<TaskPointer>) {
     }
 }
 
+/// Runs the first time a new task is scheduled, in the context of the new
+/// task. Any first-time initialization specific to user-tasks need to happen
+/// here.
+/// # Safety
+/// This is called from assembly code. The caller needs to ensure the
+/// parameters are passed correctly.
+#[unsafe(no_mangle)]
+unsafe fn run_user_task() {
+    task_attach_console();
+}
+
 /// # 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);
-    }
-
+unsafe fn run_kernel_task<T: KernelThreadStartParameter>(entry: fn(T), start_parameter: usize) {
     // 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