feat: make TSS and GDT per cpu

arch/percpu/macros/src/lib.rs

@@ -18,19 +18,17 @@ pub fn define_percpu(attrs: TokenStream, item: TokenStream) -> TokenStream {
     let ty = &item.ty;
     let expr = &item.expr;
 
-    if !["bool", "u8", "u16", "u32", "u64", "usize"].contains(&quote!(#ty).to_string().as_str())
-        && !quote!(#ty).to_string().contains("NonNull")
-    {
-        panic!("`define_percpu` only supports bool, u8, u16, u32, u64, usize and pointers");
-    }
+    let is_bool = quote!(#ty).to_string().as_str() == "bool";
+    let is_integer =
+        ["u8", "u16", "u32", "u64", "usize"].contains(&quote!(#ty).to_string().as_str());
+
+    let is_atomic_like = is_bool || is_integer || quote!(#ty).to_string().contains("NonNull");
 
     let inner_ident = format_ident!("_percpu_inner_{}", ident);
     let access_ident = format_ident!("_access_{}", ident);
 
-    let integer_methods = match quote!(#ty).to_string().as_str() {
-        "bool" => quote! {},
-        name if name.contains("NonNull") => quote! {},
-        _ => quote! {
+    let integer_methods = if is_integer {
+        quote! {
             pub fn add(&self, value: #ty) {
                 *unsafe { self.as_mut() } += value;
             }
@@ -38,7 +36,21 @@ pub fn define_percpu(attrs: TokenStream, item: TokenStream) -> TokenStream {
             pub fn sub(&self, value: #ty) {
                 *unsafe { self.as_mut() } -= value;
             }
-        },
+        }
+    } else {
+        quote! {}
+    };
+
+    let preempt_disable = if !is_atomic_like {
+        quote! { crate::sync::preempt::disable(); }
+    } else {
+        quote! {}
+    };
+
+    let preempt_enable = if !is_atomic_like {
+        quote! { crate::sync::preempt::enable(); }
+    } else {
+        quote! {}
     };
 
     let as_ptr = arch::get_percpu_pointer(&inner_ident, &ty);
@@ -52,20 +64,31 @@ pub fn define_percpu(attrs: TokenStream, item: TokenStream) -> TokenStream {
         #vis static #ident: #access_ident = #access_ident;
 
         impl #access_ident {
+            /// # Safety
+            /// This function is unsafe because it allows for mutable aliasing of the percpu
+            /// variable.
+            /// Make sure that preempt is disabled when calling this function.
             pub unsafe fn as_ptr(&self) -> *mut #ty {
                 #as_ptr
             }
 
             pub fn get(&self) -> #ty {
-                unsafe { self.as_ptr().read() }
+                #preempt_disable
+                let value = unsafe { self.as_ptr().read() };
+                #preempt_enable
+                value
             }
 
             pub fn set(&self, value: #ty) {
+                #preempt_disable
                 unsafe { self.as_ptr().write(value) }
+                #preempt_enable
             }
 
             pub fn swap(&self, mut value: #ty) -> #ty {
+                #preempt_disable
                 unsafe { self.as_ptr().swap(&mut value) }
+                #preempt_enable
                 value
             }
 

arch/x86_64/src/gdt.rs

@@ -0,0 +1,92 @@
+use core::arch::asm;
+
+use crate::task::TSS;
+
+#[repr(transparent)]
+#[derive(Debug, Clone, Copy)]
+pub struct GDTEntry(u64);
+
+pub struct GDT([GDTEntry; GDT::LEN]);
+
+impl GDTEntry {
+    const NULL: Self = Self(0);
+
+    const KERNEL_CODE64: Self = Self::new(0, 0, 0x9a, 0x2);
+    const KERNEL_DATA64: Self = Self::new(0, 0, 0x92, 0x0);
+
+    const USER_CODE64: Self = Self::new(0, 0, 0xfa, 0x2);
+    const USER_DATA64: Self = Self::new(0, 0, 0xf2, 0x0);
+
+    const USER_CODE32: Self = Self::new(0, 0xfffff, 0xfa, 0xc);
+    const USER_DATA32: Self = Self::new(0, 0xfffff, 0xf2, 0xc);
+
+    pub const fn new(base: u32, limit: u32, access: u8, flags: u8) -> Self {
+        let mut entry = 0u64;
+        entry |= (limit & 0x0000_ffff) as u64;
+        entry |= ((limit & 0x000f_0000) as u64) << 32;
+        entry |= ((base & 0x00ff_ffff) as u64) << 16;
+        entry |= ((base & 0xff00_0000) as u64) << 32;
+        entry |= (access as u64) << 40;
+        entry |= (flags as u64) << 52;
+
+        GDTEntry(entry)
+    }
+
+    pub const fn new_ldt(base: u64, limit: u32) -> [Self; 2] {
+        let first = Self::new(base as u32, limit, 0x82, 0x0);
+        let second = Self(base >> 32);
+        [first, second]
+    }
+
+    pub const fn new_tss(base: u64, limit: u32) -> [Self; 2] {
+        let first = Self::new(base as u32, limit, 0x89, 0x0);
+        let second = Self(base >> 32);
+        [first, second]
+    }
+}
+
+impl GDT {
+    const LEN: usize = 10;
+    const TLS32_INDEX: usize = 7;
+    const TSS_INDEX: usize = 8;
+
+    pub fn new() -> Self {
+        Self([
+            GDTEntry::NULL,
+            GDTEntry::KERNEL_CODE64,
+            GDTEntry::KERNEL_DATA64,
+            GDTEntry::USER_CODE64,
+            GDTEntry::USER_DATA64,
+            GDTEntry::USER_CODE32,
+            GDTEntry::USER_DATA32,
+            GDTEntry::NULL, // User TLS 32bit
+            GDTEntry::NULL, // TSS Descriptor Low
+            GDTEntry::NULL, // TSS Descriptor High
+        ])
+    }
+
+    pub fn set_tss(&mut self, base: u64) {
+        let tss = GDTEntry::new_tss(base, size_of::<TSS>() as u32 - 1);
+        self.0[Self::TSS_INDEX] = tss[0];
+        self.0[Self::TSS_INDEX + 1] = tss[1];
+    }
+
+    pub fn set_tls32(&mut self, desc: GDTEntry) {
+        self.0[7] = desc;
+    }
+
+    pub unsafe fn load(&self) {
+        let len = Self::LEN * 8 - 1;
+        let descriptor: [u64; 2] = [(len as u64) << 48, self.0.as_ptr() as u64];
+        assert!(len < 0x10000, "GDT too large");
+
+        let descriptor_address = &descriptor as *const _ as usize + 6;
+        asm!(
+            "lgdt ({})",
+            "ltr %ax",
+            in(reg) descriptor_address,
+            in("ax") Self::TSS_INDEX as u16 * 8,
+            options(att_syntax)
+        );
+    }
+}

arch/x86_64/src/lib.rs

@@ -64,6 +64,7 @@ pub mod vm {
     }
 }
 
+pub mod gdt;
 pub mod interrupt;
 pub mod io;
 pub mod task;

arch/x86_64/src/task.rs

@@ -2,6 +2,47 @@ use core::arch::{asm, global_asm};
 
 use crate::interrupt;
 
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+struct SP {
+    low: u32,
+    high: u32,
+}
+
+#[repr(C)]
+pub struct TSS {
+    _reserved1: u32,
+    rsp: [SP; 3],
+    _reserved2: u32,
+    _reserved3: u32,
+    ist: [SP; 7],
+    _reserved4: u32,
+    _reserved5: u32,
+    _reserved6: u16,
+    iomap_base: u16,
+}
+
+impl TSS {
+    pub fn new() -> Self {
+        Self {
+            _reserved1: 0,
+            rsp: [SP { low: 0, high: 0 }; 3],
+            _reserved2: 0,
+            _reserved3: 0,
+            ist: [SP { low: 0, high: 0 }; 7],
+            _reserved4: 0,
+            _reserved5: 0,
+            _reserved6: 0,
+            iomap_base: 0,
+        }
+    }
+
+    pub fn set_rsp0(&mut self, rsp: u64) {
+        self.rsp[0].low = rsp as u32;
+        self.rsp[0].high = (rsp >> 32) as u32;
+    }
+}
+
 #[inline(always)]
 pub fn halt() {
     unsafe {

doc/mem_layout.txt

@@ -1,8 +1,8 @@
 physical memory
 
-0x0000 - 0x1000 : GDT, TSS, LDT and some early kernel data
+0x0000 - 0x1000 : GDT for kernel initialization use and some early kernel data
 0x1000 - 0x2000 : kernel stage1
-0x2000 - 0x3000 : kernel PML4
+0x2000 - 0x3000 : kernel space PML4
 0x3000 - 0x4000 : kernel PDPT for physical memory mappings
 0x4000 - 0x5000 : kernel PDPT for kernel space
 0x5000 - 0x6000 : kernel PD for kernel image

include/kernel/mem/phys.hpp

@@ -13,8 +13,7 @@ namespace kernel::mem {
 
 template <typename T, bool Cached = true>
 class physaddr {
-    static constexpr uintptr_t PHYS_OFFSET =
-        Cached ? 0xffffff0000000000ULL : 0xffffff4000000000ULL;
+    static constexpr uintptr_t PHYS_OFFSET = Cached ? 0xffffff0000000000ULL : 0xffffff4000000000ULL;
 
     uintptr_t m_ptr;
 
@@ -33,21 +32,4 @@ class physaddr {
     constexpr uintptr_t phys() const noexcept { return m_ptr; }
 };
 
-//  gdt[0]:  null
-//  gdt[1]:  kernel code
-//  gdt[2]:  kernel data
-//  gdt[3]:  user code
-//  gdt[4]:  user data
-//  gdt[5]:  user code compability mode
-//  gdt[6]:  user data compability mode
-//  gdt[7]:  thread local 32bit
-//  gdt[8]:  tss descriptor low
-//  gdt[9]:  tss descriptor high
-//  gdt[10]: ldt descriptor low
-//  gdt[11]: ldt descriptor high
-//  gdt[12]: null segment(in ldt)
-//  gdt[13]: thread local 64bit(in ldt)
-// &gdt[14]: tss of 0x68 bytes from here
-constexpr physaddr<uint64_t> gdt{0x00000000 + 1 - 1};
-
 } // namespace kernel::mem

src/kernel.rs

@@ -1,3 +1,4 @@
+pub mod arch;
 pub mod block;
 pub mod console;
 pub mod constants;

src/kernel/arch.rs

@@ -0,0 +1,5 @@
+#[cfg(target_arch = "x86_64")]
+pub mod x86_64;
+
+#[cfg(target_arch = "x86_64")]
+pub use x86_64::*;

src/kernel/arch/x86_64.rs

@@ -0,0 +1,108 @@
+use arch::x86_64::{gdt::GDT, task::TSS};
+
+// TODO!!!: This can be stored in the percpu area.
+//          But we need to implement a guard that ensures that preemption is disabled
+//          while we are accessing the percpu variables.
+#[arch::define_percpu]
+static GDT_OBJECT: Option<GDT> = None;
+
+#[arch::define_percpu]
+static TSS_OBJECT: Option<TSS> = None;
+
+pub mod init {
+    use super::{GDT_OBJECT, TSS_OBJECT};
+    use crate::{kernel::smp, sync::preempt};
+    use arch::x86_64::{gdt::GDT, task::TSS};
+
+    unsafe fn init_gdt_tss_thiscpu() {
+        preempt::disable();
+        let gdt_ref = unsafe { GDT_OBJECT.as_mut() };
+        let tss_ref = unsafe { TSS_OBJECT.as_mut() };
+        *gdt_ref = Some(GDT::new());
+        *tss_ref = Some(TSS::new());
+
+        if let Some(gdt) = gdt_ref.as_mut() {
+            if let Some(tss) = tss_ref.as_mut() {
+                gdt.set_tss(tss as *mut _ as u64);
+            } else {
+                panic!("TSS is not initialized");
+            }
+
+            unsafe { gdt.load() };
+        } else {
+            panic!("GDT is not initialized");
+        }
+
+        preempt::enable();
+    }
+
+    pub unsafe fn init_bscpu() {
+        let area = smp::alloc_percpu_area();
+        smp::set_percpu_area(area);
+        init_gdt_tss_thiscpu();
+    }
+}
+
+pub mod user {
+    use crate::sync::preempt;
+    use arch::x86_64::gdt::GDTEntry;
+
+    pub struct InterruptStack(pub u64);
+
+    #[derive(Debug, Clone)]
+    pub enum TLS {
+        TLS64(u64),
+        TLS32 { base: u64, desc: GDTEntry },
+    }
+
+    impl TLS {
+        /// # Return
+        /// Returns the TLS descriptor and the index of the TLS segment.
+        pub fn new32(base: u32, limit: u32, is_limit_in_pages: bool) -> (Self, u32) {
+            let flags = if is_limit_in_pages { 0xc } else { 0x4 };
+
+            (
+                TLS::TLS32 {
+                    base: base as u64,
+                    desc: GDTEntry::new(base, limit, 0xf2, flags),
+                },
+                7,
+            )
+        }
+
+        pub fn load(&self) {
+            match self {
+                TLS::TLS64(base) => {
+                    const IA32_KERNEL_GS_BASE: u32 = 0xc0000102;
+                    arch::x86_64::task::wrmsr(IA32_KERNEL_GS_BASE, *base);
+                }
+                TLS::TLS32 { base, desc } => {
+                    preempt::disable();
+                    let gdt = unsafe {
+                        super::GDT_OBJECT
+                            .as_mut()
+                            .as_mut()
+                            .expect("GDT should be valid")
+                    };
+                    gdt.set_tls32(*desc);
+                    preempt::enable();
+
+                    const IA32_KERNEL_GS_BASE: u32 = 0xc0000102;
+                    arch::x86_64::task::wrmsr(IA32_KERNEL_GS_BASE, *base);
+                }
+            }
+        }
+    }
+
+    pub fn load_interrupt_stack(stack: InterruptStack) {
+        preempt::disable();
+        let tss = unsafe {
+            super::TSS_OBJECT
+                .as_mut()
+                .as_mut()
+                .expect("TSS should be valid")
+        };
+        tss.set_rsp0(stack.0);
+        preempt::enable();
+    }
+}

src/kernel/task/kstack.rs

@@ -1,6 +1,6 @@
-use crate::kernel::mem::{
-    paging::Page,
-    phys::{CachedPP, PhysPtr},
+use crate::kernel::{
+    arch::user::{self},
+    mem::{paging::Page, phys::PhysPtr},
 };
 
 use core::cell::UnsafeCell;
@@ -103,13 +103,7 @@ impl KernelStack {
     }
 
     pub fn load_interrupt_stack(&self) {
-        const TSS_RSP0: CachedPP = CachedPP::new(0x00000074);
-
-        // TODO!!!: Make `TSS` a per cpu struct.
-        // SAFETY: `TSS_RSP0` is always valid.
-        unsafe {
-            TSS_RSP0.as_ptr::<u64>().write_unaligned(self.bottom as u64);
-        }
+        user::load_interrupt_stack(user::InterruptStack(self.bottom as u64));
     }
 
     pub fn get_writer(&mut self) -> KernelStackWriter {

src/kernel/task/thread.rs

@@ -6,10 +6,8 @@ use core::{
 
 use crate::{
     kernel::{
-        mem::{
-            phys::{CachedPP, PhysPtr},
-            MMList,
-        },
+        arch::user::TLS,
+        mem::{phys::PhysPtr, MMList},
         terminal::Terminal,
         user::dataflow::CheckedUserPointer,
         vfs::FsContext,
@@ -191,9 +189,8 @@ struct ThreadInner {
     /// Thread name
     name: Arc<[u8]>,
 
-    /// Thread TLS descriptor 32-bit
-    tls_desc32: Option<u64>,
-    tls_base: Option<u64>,
+    /// Thread TLS
+    tls: Option<TLS>,
 
     /// User pointer
     /// Store child thread's tid when child thread returns to user space.
@@ -812,8 +809,7 @@ impl Thread {
             state: Spin::new(ThreadState::Preparing),
             inner: Spin::new(ThreadInner {
                 name,
-                tls_desc32: None,
-                tls_base: None,
+                tls: None,
                 set_child_tid: 0,
             }),
         });
@@ -842,8 +838,7 @@ impl Thread {
             state: Spin::new(ThreadState::Preparing),
             inner: Spin::new(ThreadInner {
                 name: other_inner.name.clone(),
-                tls_desc32: other_inner.tls_desc32,
-                tls_base: other_inner.tls_base,
+                tls: other_inner.tls.clone(),
                 set_child_tid: other_inner.set_child_tid,
             }),
         });
@@ -898,18 +893,8 @@ impl Thread {
     }
 
     pub fn load_thread_area32(&self) {
-        const IA32_KERNEL_GS_BASE: u32 = 0xc0000102;
-
-        let inner = self.inner.lock();
-
-        if let Some(desc32) = inner.tls_desc32 {
-            // SAFETY: `tls32` should be per cpu.
-            let tls32_addr = CachedPP::new(0x0 + 7 * 8);
-            tls32_addr.as_mut::<u64>().clone_from(&desc32);
-        }
-
-        if let Some(base) = inner.tls_base {
-            arch::x86_64::task::wrmsr(IA32_KERNEL_GS_BASE, base);
+        if let Some(tls) = self.inner.lock().tls.as_ref() {
+            tls.load();
         }
     }
 
@@ -932,24 +917,9 @@ impl Thread {
             return Ok(());
         }
 
-        if desc.entry != u32::MAX || !desc.flags.is_32bit_segment() {
-            return Err(EINVAL);
-        }
-        desc.entry = 7;
-
-        let mut desc32 = desc.limit as u64 & 0xffff;
-        desc32 |= (desc.base as u64 & 0xffffff) << 16;
-        desc32 |= 0x4_0_f2_000000_0000;
-        desc32 |= (desc.limit as u64 & 0xf_0000) << (48 - 16);
-
-        if desc.flags.is_limit_in_pages() {
-            desc32 |= 1 << 55;
-        }
-
-        desc32 |= (desc.base as u64 & 0xff_000000) << (56 - 24);
-
-        inner.tls_desc32 = Some(desc32);
-        inner.tls_base = Some(desc.base as u64);
+        let (tls, entry) = TLS::new32(desc.base, desc.limit, desc.flags.is_limit_in_pages());
+        desc.entry = entry;
+        inner.tls = Some(tls);
         Ok(())
     }
 

src/kinit.cpp

@@ -129,51 +129,12 @@ static inline void save_memory_info(bootloader_data* data) {
            sizeof(kernel::mem::info::e820_entries));
 }
 
-static inline void setup_gdt() {
-    // user code
-    mem::gdt[3] = 0x0020'fa00'0000'0000;
-    // user data
-    mem::gdt[4] = 0x0000'f200'0000'0000;
-    // user code32
-    mem::gdt[5] = 0x00cf'fa00'0000'ffff;
-    // user data32
-    mem::gdt[6] = 0x00cf'f200'0000'ffff;
-    // thread load 32bit
-    mem::gdt[7] = 0x0000'0000'0000'0000;
-
-    // TSS descriptor
-    mem::gdt[8] = 0x0000'8900'0070'0067;
-    mem::gdt[9] = 0x0000'0000'ffff'ff00;
-
-    // LDT descriptor
-    mem::gdt[10] = 0x0000'8200'0060'001f;
-    mem::gdt[11] = 0x0000'0000'ffff'ff00;
-
-    // null segment
-    mem::gdt[12] = 0x0000'0000'0000'0000;
-    // thread local 64bit
-    mem::gdt[13] = 0x0000'0000'0000'0000;
-
-    uint64_t descriptor[] = {0x005f'0000'0000'0000, (uintptr_t)(uint64_t*)mem::gdt};
-
-    asm volatile(
-        "lgdt (%0)\n\t"
-        "mov $0x50, %%ax\n\t"
-        "lldt %%ax\n\t"
-        "mov $0x40, %%ax\n\t"
-        "ltr %%ax\n\t"
-        :
-        : "r"((uintptr_t)descriptor + 6)
-        : "ax", "memory");
-}
-
 extern "C" void rust_kinit(uintptr_t early_kstack_vaddr);
 
 extern "C" void NORETURN kernel_init(bootloader_data* data) {
     enable_sse();
 
     setup_early_kernel_page_table();
-    setup_gdt();
     save_memory_info(data);
 
     uintptr_t addr_max = 0;

src/lib.rs

@@ -147,11 +147,7 @@ fn bootstrap_cpus() {
 pub extern "C" fn rust_kinit(early_kstack_pfn: usize) -> ! {
     // We don't call global constructors.
     // Rust doesn't need that, and we're not going to use global variables in C++.
-
-    unsafe {
-        let area = kernel::smp::alloc_percpu_area();
-        kernel::smp::set_percpu_area(area);
-    }
+    unsafe { kernel::arch::init::init_bscpu() };
 
     kernel::interrupt::init().unwrap();