From 6103e688070fefa9e842815a8807324b92fca635 Mon Sep 17 00:00:00 2001
From: Putta Khunchalee <ultimaweapon@outlook.com>
Date: Sun, 1 Sep 2024 15:41:50 +0700
Subject: [PATCH 1/2] Sets kernel GS

---
 src/obkrnl/src/context/aarch64.rs |  3 ++
 src/obkrnl/src/context/mod.rs     | 46 +++++++++++++++++++++++++++++++
 src/obkrnl/src/context/x86_64.rs  | 35 +++++++++++++++++++++++
 src/obkrnl/src/main.rs            | 30 ++++++++++++++++++--
 src/obkrnl/src/proc/mod.rs        |  3 ++
 src/obkrnl/src/proc/thread.rs     | 12 ++++++++
 6 files changed, 126 insertions(+), 3 deletions(-)
 create mode 100644 src/obkrnl/src/context/aarch64.rs
 create mode 100644 src/obkrnl/src/context/mod.rs
 create mode 100644 src/obkrnl/src/context/x86_64.rs
 create mode 100644 src/obkrnl/src/proc/mod.rs
 create mode 100644 src/obkrnl/src/proc/thread.rs

diff --git a/src/obkrnl/src/context/aarch64.rs b/src/obkrnl/src/context/aarch64.rs
new file mode 100644
index 000000000..3702becf7
--- /dev/null
+++ b/src/obkrnl/src/context/aarch64.rs
@@ -0,0 +1,3 @@
+pub unsafe fn activate(_: *mut Context) {
+    todo!();
+}
diff --git a/src/obkrnl/src/context/mod.rs b/src/obkrnl/src/context/mod.rs
new file mode 100644
index 000000000..f8c8b770d
--- /dev/null
+++ b/src/obkrnl/src/context/mod.rs
@@ -0,0 +1,46 @@
+use crate::proc::Thread;
+use alloc::sync::Arc;
+
+#[cfg_attr(target_arch = "aarch64", path = "aarch64.rs")]
+#[cfg_attr(target_arch = "x86_64", path = "x86_64.rs")]
+mod arch;
+
+/// Implementation of `pcpu` structure.
+///
+/// Access to this structure must be done by **atomic reading or writing its field directly**. It is
+/// not safe to have a temporary a pointer or reference to this struct or its field because the CPU
+/// might get interupted, which mean it is possible for the next instruction to get executed on
+/// a different CPU if the interupt cause the CPU to switch the task.
+///
+/// We don't support `pc_cpuid` field here because it value is 100% unpredictable due to the above
+/// reason. Once we have loaded `pc_cpuid` the next instruction might get executed on a different
+/// CPU, which render the loaded value incorrect. The only way to prevent this issue is to disable
+/// interupt before reading `pc_cpuid`, which can make the CPU missed some events from the other
+/// hardwares.
+pub struct Context {
+    thread: *const Thread, // pc_curthread
+}
+
+impl Context {
+    /// See `pcpu_init` on the PS4 for a reference.
+    pub fn new(td: Arc<Thread>) -> Self {
+        Self {
+            thread: Arc::into_raw(td),
+        }
+    }
+
+    /// # Safety
+    /// The only place this method is safe to call is in the CPU entry point. Once this method
+    /// return this instance must outlive the CPU lifetime and it must never be accessed via this
+    /// variable again. The simple way to achieve this is keep the activated [`Context`] as a local
+    /// variable then move all code after it to a dedicated no-return function.
+    pub unsafe fn activate(&mut self) {
+        self::arch::activate(self);
+    }
+}
+
+impl Drop for Context {
+    fn drop(&mut self) {
+        unsafe { drop(Arc::from_raw(self.thread)) };
+    }
+}
diff --git a/src/obkrnl/src/context/x86_64.rs b/src/obkrnl/src/context/x86_64.rs
new file mode 100644
index 000000000..fa4859793
--- /dev/null
+++ b/src/obkrnl/src/context/x86_64.rs
@@ -0,0 +1,35 @@
+use super::Context;
+use core::arch::asm;
+
+/// Set kernel `GS` segment register to `cx`.
+///
+/// This also set user-mode `FS` and `GS` to null.
+pub unsafe fn activate(cx: *mut Context) {
+    // Set GS for kernel mode.
+    let cx = cx as usize;
+
+    asm!(
+        "wrmsr",
+        in("ecx") 0xc0000101u32,
+        in("edx") cx >> 32,
+        in("eax") cx,
+        options(preserves_flags, nostack)
+    );
+
+    // Clear FS and GS for user mode.
+    asm!(
+        "wrmsr",
+        in("ecx") 0xc0000100u32,
+        in("edx") 0,
+        in("eax") 0,
+        options(preserves_flags, nostack)
+    );
+
+    asm!(
+        "wrmsr",
+        in("ecx") 0xc0000102u32,
+        in("edx") 0,
+        in("eax") 0,
+        options(preserves_flags, nostack)
+    );
+}
diff --git a/src/obkrnl/src/main.rs b/src/obkrnl/src/main.rs
index f0b3f90a2..4edfb2df4 100644
--- a/src/obkrnl/src/main.rs
+++ b/src/obkrnl/src/main.rs
@@ -2,8 +2,11 @@
 #![cfg_attr(not(test), no_main)]
 
 use crate::config::set_boot_env;
+use crate::context::Context;
 use crate::malloc::KernelHeap;
+use crate::proc::Thread;
 use alloc::string::String;
+use alloc::sync::Arc;
 use core::arch::asm;
 use core::mem::zeroed;
 use core::panic::PanicInfo;
@@ -11,19 +14,24 @@ use obconf::BootEnv;
 
 mod config;
 mod console;
+mod context;
 mod imgfmt;
 mod malloc;
 mod panic;
+mod proc;
 
 extern crate alloc;
 
 /// Entry point of the kernel.
 ///
-/// This will be called by a bootloader or a hypervisor. The following are requirements before
-/// transfer a control to this function:
+/// This will be called by a bootloader or a hypervisor. The following are requirements to call this
+/// function:
 ///
 /// 1. The kernel does not remap itself so it must be mapped at a desired virtual address and all
-///    relocations must be applied.
+///    relocations must be applied. This imply that the kernel can only be run in a virtual address
+///    space.
+/// 2. Interrupt is disabled.
+/// 3. Only main CPU can execute this function.
 ///
 /// See PS4 kernel entry point for a reference.
 #[allow(dead_code)]
@@ -34,6 +42,22 @@ extern "C" fn _start(env: &'static BootEnv) -> ! {
 
     info!("Starting Obliteration Kernel.");
 
+    // Setup thread0 to represent this thread.
+    let thread0 = unsafe { Thread::new_bare() };
+
+    // Setup CPU context. We use a different mechanism here. The PS4 put all of pcpu at a global
+    // level but we put it on each CPU stack instead.
+    let thread0 = Arc::new(thread0);
+    let mut cx = Context::new(thread0);
+
+    // SAFETY: We are in the main CPU entry point and we move all the remaining code after this into
+    // a dedicated no-return function.
+    unsafe { cx.activate() };
+
+    main();
+}
+
+fn main() -> ! {
     loop {
         #[cfg(target_arch = "x86_64")]
         unsafe {
diff --git a/src/obkrnl/src/proc/mod.rs b/src/obkrnl/src/proc/mod.rs
new file mode 100644
index 000000000..85bc2233a
--- /dev/null
+++ b/src/obkrnl/src/proc/mod.rs
@@ -0,0 +1,3 @@
+pub use self::thread::*;
+
+mod thread;
diff --git a/src/obkrnl/src/proc/thread.rs b/src/obkrnl/src/proc/thread.rs
new file mode 100644
index 000000000..975d04ebf
--- /dev/null
+++ b/src/obkrnl/src/proc/thread.rs
@@ -0,0 +1,12 @@
+/// Implementation of `thread` structure.
+pub struct Thread {}
+
+impl Thread {
+    /// # Safety
+    /// This function does not do anything except initialize the struct memory. It is the caller
+    /// responsibility to configure the thread after this so it have a proper states and trigger
+    /// necessary events.
+    pub unsafe fn new_bare() -> Self {
+        Self {}
+    }
+}

From 4709b7089dcffb8fae35fa96d8c8ab71d9272077 Mon Sep 17 00:00:00 2001
From: Putta Khunchalee <ultimaweapon@outlook.com>
Date: Sun, 1 Sep 2024 15:50:40 +0700
Subject: [PATCH 2/2] Fixes build error on AArch64

---
 src/obkrnl/src/context/aarch64.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/obkrnl/src/context/aarch64.rs b/src/obkrnl/src/context/aarch64.rs
index 3702becf7..c3bfa024b 100644
--- a/src/obkrnl/src/context/aarch64.rs
+++ b/src/obkrnl/src/context/aarch64.rs
@@ -1,3 +1,5 @@
+use super::Context;
+
 pub unsafe fn activate(_: *mut Context) {
     todo!();
 }