feat(sched): made tasks reference counted and closed use-after-free races in the wake and kill paths

kernel/arch/aarch64/sched/sched.cpp

@@ -191,7 +191,10 @@ __PRIVILEGED_CODE void on_yield(aarch64::trap_frame* tf) {
     save_cpu_context(tf, &prev->exec.cpu_ctx);
     prev->exec.tls_base = cpu::read_tls_base();
 
-    if (!(prev->exec.flags & TASK_FLAG_KERNEL) && prev->state.load_relaxed() != TASK_STATE_DEAD) {
+    // A task inside a syscall still owns kernel state such as a linked wait
+    // node, so it dies at the syscall exit fatal check instead of here
+    if (!(prev->exec.flags & (TASK_FLAG_KERNEL | TASK_FLAG_IN_SYSCALL)) &&
+        prev->state.load_relaxed() != TASK_STATE_DEAD) {
         uint32_t fsig = signals::fatal_pending(prev);
         if (fsig) {
             signals::die_from_signal(fsig);
@@ -243,7 +246,10 @@ __PRIVILEGED_CODE void on_tick(aarch64::trap_frame* tf) {
     save_cpu_context(tf, &prev->exec.cpu_ctx);
     prev->exec.tls_base = cpu::read_tls_base();
 
-    if (!(prev->exec.flags & TASK_FLAG_KERNEL) && prev->state.load_relaxed() != TASK_STATE_DEAD) {
+    // A task inside a syscall still owns kernel state such as a linked wait
+    // node, so it dies at the syscall exit fatal check instead of here
+    if (!(prev->exec.flags & (TASK_FLAG_KERNEL | TASK_FLAG_IN_SYSCALL)) &&
+        prev->state.load_relaxed() != TASK_STATE_DEAD) {
         uint32_t fsig = signals::fatal_pending(prev);
         if (fsig) {
             signals::die_from_signal(fsig);

kernel/arch/aarch64/timer/timer.cpp

@@ -157,6 +157,9 @@ __PRIVILEGED_CODE bool on_interrupt() {
         return true;
     }
 
+    // Waking raw pointers under the queue lock satisfies the wake pin
+    // contract: a sleeper re-takes this lock in cancel_sleep before it
+    // can exit, and it slept on this CPU, so wake never spins off-CPU
     while (!state.sleep_queue.empty()) {
         sched::task* t = state.sleep_queue.front();
         if (t->timer_deadline > now) break;

kernel/arch/x86_64/sched/sched.cpp

@@ -179,7 +179,10 @@ __PRIVILEGED_CODE void on_yield(x86::trap_frame* tf) {
     save_cpu_context(tf, &prev->exec.cpu_ctx);
     prev->exec.tls_base = cpu::read_tls_base();
 
-    if (!(prev->exec.flags & TASK_FLAG_KERNEL) && prev->state.load_relaxed() != TASK_STATE_DEAD) {
+    // A task inside a syscall still owns kernel state such as a linked wait
+    // node, so it dies at the syscall exit fatal check instead of here
+    if (!(prev->exec.flags & (TASK_FLAG_KERNEL | TASK_FLAG_IN_SYSCALL)) &&
+        prev->state.load_relaxed() != TASK_STATE_DEAD) {
         uint32_t fsig = signals::fatal_pending(prev);
         if (fsig) {
             signals::die_from_signal(fsig);
@@ -230,7 +233,10 @@ __PRIVILEGED_CODE void on_tick(x86::trap_frame* tf) {
     save_cpu_context(tf, &prev->exec.cpu_ctx);
     prev->exec.tls_base = cpu::read_tls_base();
 
-    if (!(prev->exec.flags & TASK_FLAG_KERNEL) && prev->state.load_relaxed() != TASK_STATE_DEAD) {
+    // A task inside a syscall still owns kernel state such as a linked wait
+    // node, so it dies at the syscall exit fatal check instead of here
+    if (!(prev->exec.flags & (TASK_FLAG_KERNEL | TASK_FLAG_IN_SYSCALL)) &&
+        prev->state.load_relaxed() != TASK_STATE_DEAD) {
         uint32_t fsig = signals::fatal_pending(prev);
         if (fsig) {
             signals::die_from_signal(fsig);

kernel/arch/x86_64/timer/timer.cpp

@@ -223,6 +223,9 @@ __PRIVILEGED_CODE bool on_interrupt() {
         return true;
     }
 
+    // Waking raw pointers under the queue lock satisfies the wake pin
+    // contract: a sleeper re-takes this lock in cancel_sleep before it
+    // can exit, and it slept on this CPU, so wake never spins off-CPU
     while (!state.sleep_queue.empty()) {
         sched::task* t = state.sleep_queue.front();
         if (t->timer_deadline > now) break;

kernel/rc/strong_ref.h

@@ -79,8 +79,8 @@ class strong_ref {
     }
 
     /**
-     * Wrap a raw pointer whose refcount is already 1 (from allocation).
-     * Does NOT call add_ref.
+     * Wrap a raw pointer, taking ownership of one already-held reference,
+     * as when adopting a freshly allocated object. Does NOT call add_ref.
      */
     [[nodiscard]] static strong_ref adopt(T* raw) noexcept {
         return strong_ref(raw, ADOPT_REF);

kernel/resource/providers/proc_provider.cpp

@@ -1,12 +1,8 @@
 #include "resource/providers/proc_provider.h"
 #include "sched/sched.h"
 #include "sched/task.h"
-#include "sched/task_registry.h"
-#include "mm/mm.h"
 #include "mm/vma.h"
-#include "mm/vmm.h"
 #include "mm/heap.h"
-#include "fs/node.h"
 #include "common/logging.h"
 #include "sync/poll.h"
 #include "sync/wait_queue.h"
@@ -39,24 +35,29 @@ __PRIVILEGED_CODE static void proc_close(resource_object* obj) {
     auto* impl = static_cast<proc_resource_impl*>(obj->impl);
     auto* pr = impl->proc.ptr();
 
+    // The reference taken under pr->lock keeps the child reclaim-safe
+    // after the lock drops, even if it exits and is reaped concurrently
     sync::irq_state irq = sync::spin_lock_irqsave(pr->lock);
+    rc::strong_ref<sched::task> child;
+    bool unstarted = false;
 
-    if (pr->child && pr->child->state.load_relaxed() == sched::TASK_STATE_CREATED) {
-        auto* child = pr->child;
-        pr->child = nullptr;
-        sync::spin_unlock_irqrestore(pr->lock, irq);
+    if (pr->child) {
+        unstarted = pr->child->state.load_relaxed() == sched::TASK_STATE_CREATED;
+        if (unstarted || (!pr->exited && !pr->detached)) {
+            child = sched::task_ref(pr->child);
+        }
 
-        if (child->proc_res) {
-            (void)child->proc_res->release();
-            child->proc_res = nullptr;
+        if (unstarted) {
+            pr->child = nullptr;
         }
-        destroy_unstarted_task(child);
-    } else if (pr->child && !pr->exited && !pr->detached) {
-        sched::task* child = pr->child;
-        sync::spin_unlock_irqrestore(pr->lock, irq);
-        sched::force_wake_for_kill(child);
-    } else {
-        sync::spin_unlock_irqrestore(pr->lock, irq);
+    }
+
+    sync::spin_unlock_irqrestore(pr->lock, irq);
+
+    if (child && unstarted) {
+        destroy_unstarted_task(child.ptr());
+    } else if (child) {
+        sched::force_wake_for_kill(child.ptr());
     }
 
     heap::kfree_delete(impl);
@@ -164,44 +165,37 @@ __PRIVILEGED_CODE proc_resource* get_proc_resource(resource_object* obj) {
 
 __PRIVILEGED_CODE void destroy_unstarted_task(sched::task* t) {
     // Claim the task, a concurrent group teardown may have already
-    // moved it to dead and handed the memory to the reaper
+    // moved it to dead and handed it to the reaper
     uint32_t expected = sched::TASK_STATE_CREATED;
     if (!t->state.cmpxchg_strong_acq_rel(expected, sched::TASK_STATE_DEAD)) {
         return;
     }
 
-    // Leave the registry before the group teardown so registry walkers
-    // never see a task whose group is being freed (same order as reap_task)
-    sched::g_task_registry.remove(*t);
-
-    resource::release_task_handles(t);
-    if (t->cwd) {
-        if (t->cwd->release()) {
-            fs::node::ref_destroy(t->cwd);
+    // Winning the claim confers sole ownership of the task's proc
+    // resource reference, a losing teardown path must not touch it
+    if (t->proc_res) {
+        if (t->proc_res->release()) {
+            proc_resource::ref_destroy(t->proc_res);
         }
-        t->cwd = nullptr;
+        t->proc_res = nullptr;
     }
 
-    if (t->exec.mm_ctx) {
-        mm::mm_context_release(t->exec.mm_ctx);
-        t->exec.mm_ctx = nullptr;
+    // Unlink from the group list here, reap_task releases the group
+    // reference but never touches the thread list
+    if (t->group && t->group->leader != t && t->group_link.is_linked()) {
+        sync::irq_state irq = sync::spin_lock_irqsave(t->group->lock);
+        t->group->threads.remove(t);
+        t->group->thread_count--;
+        sync::spin_unlock_irqrestore(t->group->lock, irq);
     }
 
-    if (t->group) {
-        if (t->group->leader != t && t->group_link.is_linked()) {
-            sync::irq_state irq = sync::spin_lock_irqsave(t->group->lock);
-            t->group->threads.remove(t);
-            t->group->thread_count--;
-            sync::spin_unlock_irqrestore(t->group->lock, irq);
-        }
-        if (t->group->release()) {
-            sched::thread_group::ref_destroy(t->group);
-        }
-        t->group = nullptr;
-    }
+    // Never started means already off-CPU, so the task goes straight to
+    // the reaper for the same staged teardown as a normal exit
+    t->cleanup_stage.store_release(sched::TASK_CLEANUP_STAGE_SCHEDULER_DETACHED);
 
-    vmm::free(t->sys_stack_base);
-    heap::kfree_delete(t);
+    if (t->release()) {
+        sched::task::ref_destroy(t);
+    }
 }
 
 } // namespace resource::proc_provider

kernel/resource/providers/proc_provider.h

@@ -47,7 +47,8 @@ __PRIVILEGED_CODE int32_t create_proc_resource(
 
 /**
  * @brief Destroy a task that was created but never started (TASK_STATE_CREATED).
- * Frees mm_ctx, system stack, and the task struct. Does NOT release proc_res ref.
+ * Claims the task against concurrent group teardown, drops its proc
+ * resource reference, and defers reclamation to the reaper.
  * @note Privilege: **required**
  */
 __PRIVILEGED_CODE void destroy_unstarted_task(sched::task* t);