From 06e6429567679a705ddba0d13ee1486529091d33 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Thu, 27 Oct 2022 08:07:34 +0200
Subject: [PATCH] Kernel: Reorganize a bit of scheduler code

---
 kernel/src/thread/Scheduler.cpp | 53 +++++++++++++++++----------------
 1 file changed, 28 insertions(+), 25 deletions(-)

diff --git a/kernel/src/thread/Scheduler.cpp b/kernel/src/thread/Scheduler.cpp
index 5526ba4d..8b31ce4a 100644
--- a/kernel/src/thread/Scheduler.cpp
+++ b/kernel/src/thread/Scheduler.cpp
@@ -25,7 +25,7 @@ static uint64_t task_num = 0;
 
 static Task idle_task;
 
-static uint64_t free_tid = 0;
+static uint64_t free_pid = 0;
 
 static Task* sched_current_task;
 static Task* base_task;
@@ -91,7 +91,7 @@ void Scheduler::append_task(Task* task)
 void Scheduler::init()
 {
     memset(&idle_task, 0, sizeof(Task));
-    idle_task.id = free_tid++;
+    idle_task.id = free_pid++;
     idle_task.regs.rip = (uint64_t)idle_task_function;
     idle_task.regs.rsp = get_top_of_stack((uint64_t)MemoryManager::get_page(), 1);
     idle_task.regs.cs = 0x08;
@@ -113,7 +113,7 @@ void Scheduler::add_kernel_task(const char* taskname, void (*task)(void))
     Task* new_task = new Task;
     ASSERT(new_task);
     new_task->user_task = false;
-    new_task->id = free_tid++;
+    new_task->id = free_pid++;
     new_task->ppid = 0;
     new_task->regs.rip = (uint64_t)task;
     new_task->allocated_stack =
@@ -140,7 +140,7 @@ Task* Scheduler::create_user_task()
     if (!new_task) return nullptr;
     memset(&new_task->regs, 0, sizeof(Context));
     new_task->user_task = true;
-    new_task->id = free_tid++;
+    new_task->id = free_pid++;
     new_task->ppid = 0;
     new_task->task_sleep = 0;
     new_task->task_time = 0;
@@ -154,31 +154,38 @@ long Scheduler::load_user_task(const char* filename)
 {
     kinfoln("Loading user task: %s", filename);
     Interrupts::push_and_disable();
-    long result;
-    if ((result = ELFLoader::check_elf_image_from_filesystem(filename)) < 0)
-    {
-        kerrorln("Failed to load %s from initrd", filename);
-        Interrupts::pop();
-        return result;
-    }
-    if ((uint64_t)result > PMM::get_free()) { return -ENOMEM; }
     Task* new_task = new Task;
     ASSERT(new_task);
     memset(&new_task->regs, 0, sizeof(Context));
-    new_task->id = free_tid++;
+    new_task->id = free_pid++;
     new_task->ppid = 0;
     if (!new_task->allocator.init())
     {
         delete new_task;
-        free_tid--;
+        free_pid--;
         Interrupts::pop();
         return -ENOMEM;
     }
     new_task->address_space = AddressSpace::create();
     VMM::switch_to_user_address_space(new_task->address_space);
-    ELFImage* image = ELFLoader::load_elf_from_filesystem(
-        filename); // FIXME: TOCTOU? Right now, impossible, since interrupts are disabled and SMP is not a thing. But in
-                   // the future, it might be possible.
+    long result;
+    if ((result = ELFLoader::check_elf_image_from_filesystem(filename)) < 0)
+    {
+        delete new_task;
+        free_pid--;
+        kerrorln("Failed to load %s from initrd", filename);
+        Interrupts::pop();
+        return result;
+    }
+    if ((uint64_t)result > PMM::get_free())
+    {
+        delete new_task;
+        free_pid--;
+        kerrorln("Not enough memory for task %s", filename);
+        Interrupts::pop();
+        return -ENOMEM;
+    }
+    ELFImage* image = ELFLoader::load_elf_from_filesystem(filename);
     ASSERT(image);
     new_task->user_task = true;
     new_task->regs.rip = image->entry;
@@ -189,7 +196,7 @@ long Scheduler::load_user_task(const char* filename)
     {
         new_task->address_space.destroy();
         delete new_task;
-        free_tid--;
+        free_pid--;
         ELFLoader::release_elf_image(image);
         VMM::switch_back_to_kernel_address_space();
         Interrupts::pop();
@@ -245,13 +252,10 @@ void Scheduler::reap_task(Task* task)
     }
     kinfoln("reaping task %s, PID %ld, exited with code %ld", exiting_task->name, exiting_task->id,
             exiting_task->exit_status);
+    if (exiting_task->id == (free_pid - 1)) free_pid--; // If we are the last spawned thread, free our PID.
     if (exiting_task->allocated_stack && !exiting_task->is_user_task())
         MemoryManager::release_pages((void*)exiting_task->allocated_stack, TASK_PAGES_IN_STACK);
-    if (exiting_task->image) // FIXME: Also free pages the task has mmap-ed but not munmap-ed.
-    {
-        // ELFLoader::release_elf_image(exiting_task->image);
-        kfree(exiting_task->image);
-    }
+    if (exiting_task->image) kfree(exiting_task->image);
     if (exiting_task->is_user_task())
     {
         exiting_task->allocator.free();
@@ -262,7 +266,6 @@ void Scheduler::reap_task(Task* task)
         Interrupts::pop();
     }
     for (int i = 0; i < TASK_MAX_FDS; i++) { exiting_task->files[i].close(); }
-    if (exiting_task->id == (free_tid - 1)) free_tid--; // If we are the last spawned thread, free our PID.
     delete exiting_task;
 }
 
@@ -527,7 +530,7 @@ struct pstat
 void sys_pstat(Context* context, long pid, struct pstat* buf)
 {
     Task* task;
-    if (pid == -1) task = Scheduler::find_by_pid(free_tid - 1);
+    if (pid == -1) task = Scheduler::find_by_pid(free_pid - 1);
     else if (pid == 0)
         task = &idle_task;
     else