Add mprotect()

Author: My-Bad-2

kernel/include/memory/pagemap.hpp

@@ -38,6 +38,8 @@ class PageMap {
 
     void unmap(uintptr_t virt_addr, uint16_t owner_pcid = 0, bool free_phys = false);
     uintptr_t translate(uintptr_t virt_addr);
+    size_t set_page_flags(uintptr_t virt_addr, uint8_t flags,
+                          CacheType cache = CacheType::WriteBack, uint16_t owner_pcid = 0);
 
     void load(uint16_t pcid = 0, bool flush = true);
 

kernel/include/task/process.hpp

@@ -56,7 +56,7 @@ struct Thread : public IntrusiveListNode<SchedulerTag>,
 
     Thread() = default;
     Thread(Process* parent, void (*callback)(void*), void* args);
-    ~Thread() {}
+    ~Thread();
 
    private:
     void arch_init(uintptr_t entry, uintptr_t arg);
@@ -88,6 +88,7 @@ struct Process : public IntrusiveListNode<ProcessTag> {
 
     void* mmap(void* addr, size_t len, int prot, int flags);
     void munmap(void* ptr, size_t len);
+    int mprotect(void* addr, size_t len, int prot);
 
     static void init();
 

kernel/src/arch/x86_64/memory/paging.cpp

@@ -2,6 +2,7 @@
 #include "memory/pagemap.hpp"
 #include "memory/memory.hpp"
 #include "memory/paging.hpp"
+#include <cstdint>
 #include "memory/pcid_manager.hpp"
 #include "memory/pmm.hpp"
 #include "cpu/registers.hpp"
@@ -525,6 +526,67 @@ uintptr_t PageMap::translate(uintptr_t virt_addr) {
     return 0;
 }
 
+size_t PageMap::set_page_flags(uintptr_t virt_addr, uint8_t flags, CacheType cache,
+                               uint16_t owner_pcid) {
+    uintptr_t curr_table_phys = this->phys_root_addr;
+    size_t new_flags          = convert_generic_flags(flags, cache, PageSize::Size4K);
+
+    for (int level = max_levels; level >= 1; --level) {
+        uintptr_t* table_virt = reinterpret_cast<uintptr_t*>(to_higher_half(curr_table_phys));
+
+        int shift = 12 + (level - 1) * 9;
+        int index = static_cast<int>((virt_addr >> shift) & 0x1FF);
+
+        uint64_t entry = table_virt[index];
+
+        if (!(entry & FlagPresent)) {
+            return false;
+        }
+
+        bool is_huge = (level > 1) && (entry & FlagHuge);
+        bool is_leaf = (level == 1);
+
+        if (is_huge || is_leaf) {
+            uintptr_t phys_base = entry & page_mask;
+            uint64_t new_entry  = phys_base | new_flags;
+
+            if (is_huge) {
+                new_entry |= FlagHuge;
+
+                // convert_generic_flags() doesn't handles Large PAT here.
+                // So, we do it manually.
+                switch (cache) {
+                    case CacheType::WriteProtected:
+                        new_entry |= FlagLPAT;
+                        break;
+                    case CacheType::WriteCombining:
+                        new_entry |= FlagLPAT;
+                        break;
+                    case CacheType::WriteBack:
+                    case CacheType::WriteThrough:
+                    case CacheType::Uncached:
+                    default:
+                        break;
+                }
+            }
+
+            table_virt[index] = new_entry;
+
+            if (this->is_active()) {
+                TLB::flush(virt_addr);
+            } else {
+                TLB::flush_specific(virt_addr, owner_pcid);
+            }
+
+            return 1ul << shift;
+        }
+
+        curr_table_phys = entry & page_mask;
+    }
+
+    return 0;
+}
+
 void PageMap::create_new(PageMap* map) {
     static bool kernel_initialized = false;
     uintptr_t* root_phys           = static_cast<uintptr_t*>(PhysicalManager::alloc());

kernel/src/libs/mutex.cpp

@@ -11,7 +11,7 @@ void Mutex::timeout_callback(void* data) {
 
     ctx->timed_out.store(true, std::memory_order_release);
 
-    task::Scheduler& sched = cpu::CpuCoreManager::get().get_current_core()->sched;
+    task::Scheduler& sched = task::Scheduler::get();
     sched.unblock(ctx->thread);
 }
 
@@ -130,9 +130,8 @@ void Mutex::wakeup_next() {
     LockGuard guard(this->queue_lock);
 
     if (this->wait_head) {
-        task::Thread* t      = this->wait_head->thread;
-        cpu::PerCpuData* cpu = cpu::CpuCoreManager::get().get_current_core();
-        cpu->sched.unblock(t);
+        task::Thread* t = this->wait_head->thread;
+        task::Scheduler::get().unblock(t);
 
         WaitNode* old   = this->wait_head;
         this->wait_head = this->wait_head->next;

kernel/src/task/process.cpp

@@ -36,6 +36,16 @@ Thread::Thread(Process* proc, void (*callback)(void*), void* args) {
     this->arch_init(reinterpret_cast<uintptr_t>(callback), reinterpret_cast<uintptr_t>(args));
 }
 
+Thread::~Thread() {
+    if (this->is_user_thread) {
+        this->owner->munmap(this->kernel_stack, USTACK_SIZE);
+    } else {
+        delete this->kernel_stack;
+    }
+
+    delete this->fpu_storage;
+}
+
 Process::Process(memory::PageMap* map) : map(map) {
     this->pid = next_pid.fetch_add(1, std::memory_order_relaxed);
     this->next_tid.store(1, std::memory_order_relaxed);
@@ -145,4 +155,43 @@ void Process::munmap(void* ptr, size_t) {
 void Process::init() {
     kernel_proc = new Process(memory::PageMap::get_kernel_map());
 }
+
+int Process::mprotect(void* addr, size_t len, int prot) {
+    uintptr_t virt_start = reinterpret_cast<uintptr_t>(addr);
+
+    if (virt_start & (memory::PAGE_SIZE_4K - 1)) {
+        return -1;
+    }
+
+    uint8_t flags           = memory::Read | memory::User;
+    memory::CacheType cache = memory::CacheType::WriteBack;
+
+    // Align length to 4KiB
+    size_t aligned_len = align_up(len, memory::PAGE_SIZE_4K);
+
+    if (prot & PROT_WRITE) {
+        flags |= memory::Write;
+    }
+
+    if (prot & PROT_EXEC) {
+        flags |= memory::Execute;
+    }
+
+    uintptr_t curr_virt  = virt_start;
+    uintptr_t virt_end   = virt_start + aligned_len;
+    uint16_t active_pcid = memory::PcidManager::get().get_pcid(this);
+
+    // TODO: Implement shatter page in PageMap
+    while (curr_virt < virt_end) {
+        size_t chunk_size = this->map->set_page_flags(curr_virt, flags, cache, active_pcid);
+
+        if (chunk_size == 0) {
+            return -1;
+        }
+
+        curr_virt += chunk_size;
+    }
+
+    return 0;
+}
 }  // namespace kernel::task

kernel/src/task/scheduler.cpp

@@ -17,6 +17,9 @@ constexpr uint16_t TIME_SLICE_QUANTA[64] = {
 }  // namespace
 
 void Scheduler::add_thread(Thread* t) {
+    cpu::CpuCoreManager& cpu_manager = cpu::CpuCoreManager::get();
+    cpu::PerCpuData* cpu             = cpu_manager.get_core_by_index(this->cpu_id);
+
     if (t->priority >= MLFQ_LEVELS) {
         t->priority = MLFQ_LEVELS - 1;
     }
@@ -25,9 +28,8 @@ void Scheduler::add_thread(Thread* t) {
         t->quantum = TIME_SLICE_QUANTA[t->priority];
     }
 
-    t->state             = Ready;
-    cpu::PerCpuData* cpu = cpu::CpuCoreManager::get().get_core_by_index(this->cpu_id);
-    t->cpu               = cpu;
+    t->state = Ready;
+    t->cpu   = cpu;
 
     {
         LockGuard guard(this->lock);
@@ -45,8 +47,8 @@ void Scheduler::add_thread(Thread* t) {
         cpu->reschedule_needed = true;
 
         // If this is running on a different core, send an IPI to wake it up
-        if (cpu::CpuCoreManager::get().get_current_core()->core_idx != this->cpu_id) {
-            cpu::CpuCoreManager::get().send_ipi(this->cpu_id, IPI_RESCHEDULE_VECTOR);
+        if (cpu_manager.get_current_core()->core_idx != this->cpu_id) {
+            cpu_manager.send_ipi(this->cpu_id, IPI_RESCHEDULE_VECTOR);
         }
     }
 }
@@ -77,8 +79,9 @@ Thread* Scheduler::get_next_thread() {
 }
 
 Thread* Scheduler::try_steal() {
+    cpu::CpuCoreManager& cpu_manager = cpu::CpuCoreManager::get();
     // Iterate over all other CPUs
-    size_t max_cpus = cpu::CpuCoreManager::get().get_total_cores();
+    size_t max_cpus = cpu_manager.get_total_cores();
 
     for (size_t i = 0; i < max_cpus; ++i) {
         size_t victim_id = (this->cpu_id + i) % max_cpus;
@@ -89,7 +92,7 @@ Thread* Scheduler::try_steal() {
         }
 
         cpu::PerCpuData* victim_cpu =
-            cpu::CpuCoreManager::get().get_core_by_index(static_cast<uint32_t>(victim_id));
+            cpu_manager.get_core_by_index(static_cast<uint32_t>(victim_id));
         Scheduler& victim_sched = victim_cpu->sched;
 
         // Use `try_lock()` to avoid deadlock
@@ -119,7 +122,7 @@ Thread* Scheduler::try_steal() {
 
             if (stolen) {
                 // Migrate the thread to this cpu
-                stolen->cpu = cpu::CpuCoreManager::get().get_core_by_index(this->cpu_id);
+                stolen->cpu = cpu_manager.get_core_by_index(this->cpu_id);
                 return stolen;
             }
         }