From 8e19da338b264aea2bd99660965b8e4ab32166c6 Mon Sep 17 00:00:00 2001
From: 4lDO2 <4lDO2@protonmail.com>
Date: Thu, 7 Mar 2024 15:03:01 +0100
Subject: [PATCH] Fix kstop and kreset.

---
 src/arch/aarch64/stop.rs    |  6 ++----
 src/arch/x86_shared/stop.rs | 10 ++++------
 src/context/signal.rs       | 30 ++++++++++++++++++++++++++++--
 src/main.rs                 |  7 ++++++-
 src/syscall/process.rs      | 20 --------------------
 5 files changed, 40 insertions(+), 33 deletions(-)

diff --git a/src/arch/aarch64/stop.rs b/src/arch/aarch64/stop.rs
index e3251f334a..cf9a48829d 100644
--- a/src/arch/aarch64/stop.rs
+++ b/src/arch/aarch64/stop.rs
@@ -1,7 +1,6 @@
 use core::arch::asm;
 
-#[no_mangle]
-pub unsafe extern "C" fn kreset() -> ! {
+pub unsafe fn kreset() -> ! {
     println!("kreset");
 
     let val: u32 = 0x8400_0009;
@@ -15,8 +14,7 @@ pub unsafe fn emergency_reset() -> ! {
     asm!("hvc   #0", options(noreturn));
 }
 
-#[no_mangle]
-pub unsafe extern "C" fn kstop() -> ! {
+pub unsafe fn kstop() -> ! {
     println!("kstop");
 
     let val: u32 = 0x8400_0008;
diff --git a/src/arch/x86_shared/stop.rs b/src/arch/x86_shared/stop.rs
index 50ceab3ade..dd1d16065c 100644
--- a/src/arch/x86_shared/stop.rs
+++ b/src/arch/x86_shared/stop.rs
@@ -3,9 +3,8 @@ use crate::{context, scheme::acpi, time};
 
 use crate::syscall::io::{Io, Pio};
 
-#[no_mangle]
-pub unsafe extern "C" fn kreset() -> ! {
-    println!("kreset");
+pub unsafe fn kreset() -> ! {
+    log::info!("kreset");
 
     // 8042 reset
     {
@@ -51,7 +50,7 @@ fn userspace_acpi_shutdown() {
         // TODO: Switch directly to whichever process is handling the kstop pipe. We would add an
         // event flag like EVENT_DIRECT, which has already been suggested for IRQs.
         // TODO: Waitpid with timeout? Because, what if the ACPI driver would crash?
-        let _ = unsafe { context::switch() };
+        let _ = context::switch();
 
         let current = time::monotonic();
         if current - initial > time::NANOS_PER_SEC {
@@ -61,8 +60,7 @@ fn userspace_acpi_shutdown() {
     }
 }
 
-#[no_mangle]
-pub unsafe extern "C" fn kstop() -> ! {
+pub unsafe fn kstop() -> ! {
     log::info!("Running kstop()");
 
     #[cfg(feature = "acpi")]
diff --git a/src/context/signal.rs b/src/context/signal.rs
index 9c9fe5ad83..d0e619502e 100644
--- a/src/context/signal.rs
+++ b/src/context/signal.rs
@@ -5,15 +5,41 @@ use syscall::{
         PTRACE_FLAG_IGNORE, PTRACE_STOP_SIGNAL, SIGCHLD, SIGCONT, SIGKILL, SIGSTOP, SIGTSTP,
         SIGTTIN, SIGTTOU, SIG_DFL, SIG_IGN,
     },
-    ptrace_event, SignalStack, SigActionFlags, IntRegisters, Error, EINTR,
+    ptrace_event, SignalStack, SigActionFlags, IntRegisters, Error, EINTR, SIGTERM,
 };
 
 use crate::{
     context::{self, switch, Status, WaitpidKey},
     ptrace,
-    syscall::usercopy::UserSlice,
+    syscall::usercopy::UserSlice, stop::{kstop, kreset},
 };
 
+use super::ContextId;
+
+pub fn kmain_signal_handler() {
+    if context::context_id() != ContextId::new(1) {
+        log::warn!("kmain signal didn't target PID 1, ignoring");
+        return;
+    }
+
+    let deliverable = context::current().expect("context::kmain_signal_handler not inside of context");
+    let kstop_bit = 1 << (SIGKILL - 1);
+    let kreset_bit = 1 << (SIGTERM - 1);
+    let bits = deliverable.read().sig.deliverable();
+
+    if bits & kstop_bit == kstop_bit {
+        unsafe {
+            kstop();
+        }
+    } else if bits & kreset_bit == kreset_bit {
+        unsafe {
+            kreset();
+        }
+    } else {
+        log::warn!("Spurious kmain signal, bitmask {bits:#0x}.");
+    }
+}
+
 // TODO: Move everything but SIGKILL to userspace. SIGCONT and SIGSTOP do not necessarily need to
 // be done from this current context.
 pub fn signal_handler(eintr: bool) {
diff --git a/src/main.rs b/src/main.rs
index b7b5f0a89c..b08b10f0fb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -243,7 +243,12 @@ fn run_userspace() -> ! {
         unsafe {
             interrupt::disable();
             match context::switch() {
-                SwitchResult::Switched { .. } => interrupt::enable_and_nop(),
+                SwitchResult::Switched { signal } => {
+                    if signal {
+                        crate::context::signal::kmain_signal_handler();
+                    }
+                    interrupt::enable_and_nop();
+                }
                 SwitchResult::AllContextsIdle => {
                     // Enable interrupts, then halt CPU (to save power) until the next interrupt is actually fired.
                     interrupt::enable_and_halt();
diff --git a/src/syscall/process.rs b/src/syscall/process.rs
index aa9a5a0a49..e16620ecc7 100644
--- a/src/syscall/process.rs
+++ b/src/syscall/process.rs
@@ -50,26 +50,6 @@ pub fn exit(status: usize) -> ! {
             context.id
         };
 
-        // TODO: Find a better way to implement this, perhaps when the init process calls exit.
-        if pid == ContextId::from(1) {
-            println!("Main kernel thread exited with status {:X}", status);
-
-            extern "C" {
-                fn kreset() -> !;
-                fn kstop() -> !;
-            }
-
-            if status == SIGTERM {
-                unsafe {
-                    kreset();
-                }
-            } else {
-                unsafe {
-                    kstop();
-                }
-            }
-        }
-
         // Files must be closed while context is valid so that messages can be passed
         for (_fd, file_opt) in close_files.into_iter().enumerate() {
             if let Some(file) = file_opt {