tls: cleanup the FS base restoration

Signed-off-by: Andy-Python-Programmer <[email protected]>

Author: Andy-Python-Programmer

src/aero_kernel/src/arch/x86_64/task.asm

@@ -20,15 +20,12 @@ global task_spinup
 global iretq_init
 global fork_init
 
-extern restore_user_tls
-
 jump_userland_exec:
     push rdi ; Param: stack
     push rsi ; Param: RIP
     push rdx ; Param: RFLAGS
 
     cli
-    call restore_user_tls
 
     pop r11
     pop rcx
@@ -38,10 +35,7 @@ jump_userland_exec:
     o64 sysret
 
 fork_init:
-    cli
-    call restore_user_tls
     swapgs
-
     jmp iretq_init
 
 iretq_init:

src/aero_kernel/src/arch/x86_64/task.rs

@@ -45,7 +45,7 @@ use crate::fs::cache::DirCacheItem;
 use crate::mem::paging::*;
 use crate::syscall::ExecArgs;
 use crate::userland::vm::Vm;
-use crate::utils::StackHelper;
+use crate::utils::{io, StackHelper};
 
 use super::controlregs;
 
@@ -327,13 +327,19 @@ impl ArchTask {
         Ok(())
     }
 
-    #[inline]
+    /// Returns the FS base for this instance.
     pub fn get_fs_base(&self) -> VirtAddr {
         self.fs_base
     }
 
-    #[inline]
-    pub fn set_fs_base(&mut self, base: VirtAddr) {
+    /// Sets the FS base to the provided `base`.
+    ///
+    /// ## Safety
+    /// This function **must** be called by the process that this [`ArchTask`] instance
+    /// belongs to. This is required since we also update the FS base register with the
+    /// `base` immediately (not waiting for a switch).
+    pub unsafe fn set_fs_base(&mut self, base: VirtAddr) {
+        io::wrmsr(io::IA32_FS_BASE, base.as_u64());
         self.fs_base = base;
     }
 }
@@ -347,5 +353,10 @@ pub fn arch_task_spinup(from: &mut ArchTask, to: &ArchTask) {
     unsafe {
         super::gdt::get_task_state_segement().rsp[0] = to.context_switch_rsp.as_u64();
         task_spinup(&mut from.context, to.context.as_ref());
+
+        // make a restore point for the current FS base.
+        from.fs_base = VirtAddr::new(io::rdmsr(io::IA32_FS_BASE));
+        // switch to the new FS base.
+        io::wrmsr(io::IA32_FS_BASE, to.fs_base.as_u64());
     }
 }

src/aero_kernel/src/arch/x86_64/tls.rs

@@ -33,7 +33,6 @@ use alloc::vec::Vec;
 
 use super::gdt::*;
 
-use crate::userland::scheduler;
 use crate::utils::io;
 use crate::utils::sync::Mutex;
 
@@ -191,15 +190,3 @@ where
         f(&info);
     }
 }
-
-#[no_mangle]
-extern "C" fn restore_user_tls() {
-    unsafe {
-        let base = scheduler::get_scheduler()
-            .current_task()
-            .arch_task_mut()
-            .get_fs_base();
-
-        io::wrmsr(io::IA32_FS_BASE, base.as_u64());
-    }
-}

src/aero_kernel/src/syscall/handler.asm

@@ -86,9 +86,7 @@ syscall_handler:
 
     cld
     call __inner_syscall
-
     cli
-    call restore_user_tls
 
     ; pop the preserved registers
     pop r15

src/aero_kernel/src/syscall/process.rs

@@ -59,14 +59,14 @@ const ARCH_GET_GS: usize = 0x1004;
 
 pub fn arch_prctl(command: usize, address: usize) -> Result<usize, AeroSyscallError> {
     match command {
-        ARCH_SET_FS => {
+        ARCH_SET_FS => unsafe {
             scheduler::get_scheduler()
                 .current_task()
                 .arch_task_mut()
                 .set_fs_base(VirtAddr::new(address as u64));
 
             Ok(0x00)
-        }
+        },
 
         ARCH_GET_FS => Ok(scheduler::get_scheduler()
             .current_task()