Kernel, libc: Implement fork()

This time for real.

Also, add a new per-user-task virtual address allocator (UserHeap), so that mmap'ed pages are in user range and can be copied.

apps/src/init.c

@@ -154,19 +154,28 @@ int main()
 
     printf("Success!!\n");
 
-    const char* execpath = "/bin/sym";
+    printf("Forking...\n");
 
-    printf("Spawning %s\n", execpath);
+    pid_t child = fork();
 
-    pid_t child;
+    if (child < 0)
-
-    if ((child = spawn(execpath)) < 0)
     {
-        perror("spawn");
+        perror("fork");
         return 1;
     }
-
+    if (child == 0)
-    printf("Success!! Got PID %ld\n", child);
+    {
+        msleep(500);
+        printf("I am the child, who is my parent?\n");
+        execv("/bin/sym", NULL);
+        perror("execv");
+        return 1;
+    }
+    else
+    {
+        printf("Success!! Got PID %ld\n", child);
+        return 0;
+    }
 
     return 0;
 }

kernel/include/memory/AddressSpace.h

@@ -7,6 +7,8 @@ struct AddressSpace
 
     void destroy();
 
+    void clear();
+
     AddressSpace clone();
 
     PageTable* get_pml4()

kernel/include/memory/UserHeap.h

@@ -0,0 +1,26 @@
+#pragma once
+#include <stdint.h>
+
+struct UserHeap
+{
+    bool init();
+
+    uint64_t request_virtual_page();
+    uint64_t request_virtual_pages(uint64_t count);
+
+    void free_virtual_page(uint64_t address);
+    void free_virtual_pages(uint64_t address, uint64_t count);
+
+    void free();
+
+    bool inherit(UserHeap& other);
+
+  private:
+    uint8_t* bitmap = nullptr;
+    uint64_t bitmap_size = 0;
+    uint64_t start_index = 0;
+    bool bitmap_read(uint64_t index);
+    void bitmap_set(uint64_t index, bool value);
+
+    bool try_expand();
+};

kernel/include/sys/Syscall.h

@@ -20,6 +20,7 @@
 #define SYS_clock 15
 #define SYS_spawn 16
 #define SYS_mkdir 17
+#define SYS_fork 18
 
 namespace Syscall
 {
@@ -45,4 +46,5 @@ void sys_fcntl(Context* context, int fd, int command, uintptr_t arg);
 void sys_mprotect(Context* context, void* address, size_t size, int prot);
 void sys_clock(Context* context);
 void sys_spawn(Context* context, const char* pathname);
-void sys_mkdir(Context* context, const char* filename);
+void sys_mkdir(Context* context, const char* filename);
+void sys_fork(Context* context);

kernel/include/thread/Task.h

@@ -2,6 +2,7 @@
 #include "fs/FileDescriptor.h"
 #include "interrupts/Context.h"
 #include "memory/AddressSpace.h"
+#include "memory/UserHeap.h"
 #include "sys/elf/Image.h"
 
 #define TASK_MAX_FDS 32
@@ -47,6 +48,8 @@ struct Task
 
     AddressSpace address_space;
 
+    UserHeap allocator;
+
     int alloc_fd();
 
     int alloc_fd_greater_than_or_equal(int base_fd);

kernel/src/memory/AddressSpace.cpp

@@ -74,6 +74,62 @@ void AddressSpace::destroy()
     kdbgln("Reclaimed %ld pages from address space!", pages_freed);
 }
 
+void AddressSpace::clear()
+{
+    uint64_t pages_freed = 0;
+    for (int i = 0; i < 512; i++)
+    {
+        PageDirectoryEntry& pdp_pde = m_pml4->entries[i];
+        if (!pdp_pde.present) continue;
+        if (pdp_pde.larger_pages)
+        {
+            pages_freed++;
+            PMM::free_page((void*)pdp_pde.get_address());
+            continue;
+        }
+        PageTable* pdp = (PageTable*)pdp_pde.get_address();
+        for (int j = 0; j < 511; j++) // skip the last page directory, it's the kernel one
+        {
+            PageDirectoryEntry& pd_pde = pdp->entries[j];
+            if (!pd_pde.present) continue;
+            if (pd_pde.larger_pages)
+            {
+                pages_freed++;
+                PMM::free_page((void*)pd_pde.get_address());
+                continue;
+            }
+            PageTable* pd = (PageTable*)pd_pde.get_address();
+            for (int k = 0; k < 512; k++)
+            {
+                PageDirectoryEntry& pt_pde = pd->entries[k];
+                if (!pt_pde.present) continue;
+                if (pt_pde.larger_pages)
+                {
+                    pages_freed++;
+                    PMM::free_page((void*)pt_pde.get_address());
+                    continue;
+                }
+                PageTable* pt = (PageTable*)pt_pde.get_address();
+                for (int l = 0; l < 512; l++)
+                {
+                    PageDirectoryEntry& pde = pt->entries[l];
+                    if (!pde.present) continue;
+                    pages_freed++;
+                    PMM::free_page((void*)pde.get_address());
+                }
+                pages_freed++;
+                PMM::free_page(pt);
+            }
+            pages_freed++;
+            PMM::free_page(pd);
+        }
+        pages_freed++;
+        PMM::free_page(pdp);
+    }
+
+    kdbgln("Reclaimed %ld pages from address space!", pages_freed);
+}
+
 AddressSpace AddressSpace::clone() // FIXME: Add out-of-memory checks to this function.
 {
     AddressSpace result;

kernel/src/memory/UserHeap.cpp

@@ -0,0 +1,143 @@
+#define MODULE "mem"
+
+#include "memory/UserHeap.h"
+#include "log/Log.h"
+#include "std/stdlib.h"
+#include "std/string.h"
+
+#ifndef PAGE_SIZE
+#define PAGE_SIZE 4096
+#endif
+
+#define ALLOC_BASE 0xa00000
+
+#define INITIAL_SIZE 0x2000
+#define EXPAND_SIZE 0x1000
+
+bool UserHeap::init()
+{
+    bitmap = (uint8_t*)kmalloc(INITIAL_SIZE);
+    if (!bitmap) return false;
+    bitmap_size = INITIAL_SIZE;
+    memset(bitmap, 0, bitmap_size);
+    kdbgln("new user heap, bitmap at %p, size %ld", (void*)bitmap, bitmap_size);
+    return true;
+}
+
+bool UserHeap::inherit(UserHeap& other)
+{
+    bitmap = (uint8_t*)kmalloc(other.bitmap_size);
+    if (!bitmap) return false;
+    bitmap_size = other.bitmap_size;
+    memcpy(bitmap, other.bitmap, bitmap_size);
+    kdbgln("child user heap, bitmap at %p, size %ld", (void*)bitmap, bitmap_size);
+    return true;
+}
+
+void UserHeap::free()
+{
+    kdbgln("freeing user heap, bitmap at %p, size %ld", (void*)bitmap, bitmap_size);
+    kfree(bitmap);
+}
+
+bool UserHeap::try_expand()
+{
+    kdbgln("attempting to expand user heap");
+    void* new_bitmap = krealloc(bitmap, bitmap_size + EXPAND_SIZE);
+    if (!new_bitmap)
+    {
+        kdbgln("expansion failed");
+        return false;
+    }
+    bitmap = (uint8_t*)new_bitmap;
+    memset(bitmap + bitmap_size, 0, EXPAND_SIZE);
+    bitmap_size += EXPAND_SIZE;
+    kdbgln("expanded user heap, bitmap at %p, size %ld", (void*)bitmap, bitmap_size);
+    return true;
+}
+
+bool UserHeap::bitmap_read(uint64_t index)
+{
+    return (bitmap[index / 8] & (0b10000000 >> (index % 8))) > 0;
+}
+
+void UserHeap::bitmap_set(uint64_t index, bool value)
+{
+    uint64_t byteIndex = index / 8;
+    uint8_t bitIndexer = 0b10000000 >> (index % 8);
+    bitmap[byteIndex] &= ~bitIndexer;
+    if (value) { bitmap[byteIndex] |= bitIndexer; }
+}
+
+uint64_t UserHeap::request_virtual_page()
+{
+    uint64_t attempts = 0;
+allocate:
+    for (uint64_t index = start_index; index < bitmap_size * 8; index++)
+    {
+        if (bitmap_read(index)) continue;
+        bitmap_set(index, true);
+        start_index = index + 1;
+        return ALLOC_BASE + (index * PAGE_SIZE);
+    }
+
+    if (attempts < 5 && try_expand())
+    {
+        attempts++;
+        goto allocate;
+    }
+
+    return 0;
+}
+
+uint64_t UserHeap::request_virtual_pages(uint64_t count)
+{
+    uint64_t attempts = 0;
+allocate:
+    uint64_t contiguous = 0;
+    uint64_t contiguous_start = 0;
+    for (uint64_t index = start_index; index < bitmap_size * 8; index++)
+    {
+        if (bitmap_read(index))
+        {
+            contiguous = 0;
+            continue;
+        }
+        if (contiguous == 0)
+        {
+            contiguous_start = index;
+            contiguous++;
+        }
+        else
+            contiguous++;
+        if (contiguous == count)
+        {
+            for (uint64_t i = 0; i < count; i++) bitmap_set(contiguous_start + i, true);
+            return ALLOC_BASE + (contiguous_start * PAGE_SIZE);
+        }
+    }
+
+    if (attempts < 5 && try_expand())
+    {
+        attempts++;
+        goto allocate;
+    }
+
+    return 0;
+}
+
+void UserHeap::free_virtual_page(uint64_t address)
+{
+    if (address < ALLOC_BASE || address >= (ALLOC_BASE + bitmap_size * 8 * PAGE_SIZE)) return;
+    uint64_t index = (address - ALLOC_BASE) / PAGE_SIZE;
+    bitmap_set(index, false);
+    if (start_index > index) start_index = index;
+}
+
+void UserHeap::free_virtual_pages(uint64_t address, uint64_t count)
+{
+    if (address < ALLOC_BASE || address >= (ALLOC_BASE + bitmap_size * 8 * PAGE_SIZE)) return;
+    uint64_t index = (address - ALLOC_BASE) / PAGE_SIZE;
+    for (uint64_t i = 0; i < count; i++) { bitmap_set(index + i, false); }
+    if (start_index > index) start_index = index;
+}

kernel/src/sys/Syscall.cpp

@@ -29,6 +29,7 @@ void Syscall::entry(Context* context)
     case SYS_clock: sys_clock(context); break;
     case SYS_mkdir: sys_mkdir(context, (const char*)context->rdi); break;
     case SYS_spawn: sys_spawn(context, (const char*)context->rdi); break;
+    case SYS_fork: sys_fork(context); break;
     default: context->rax = -ENOSYS; break;
     }
     VMM::exit_syscall_context();

kernel/src/sys/exec.cpp

@@ -13,11 +13,46 @@
 #include "sys/elf/ELFLoader.h"
 #include "thread/Scheduler.h"
 
+void sys_fork(Context* context)
+{
+    kinfoln("fork(): attempting fork");
+
+    Task* parent = Scheduler::current_task();
+
+    Task* child = Scheduler::create_user_task();
+    if (!child)
+    {
+        context->rax = -ENOMEM;
+        return;
+    }
+
+    if (!child->allocator.inherit(parent->allocator))
+    {
+        child->state = child->Exited;
+        child->exit_status = -127; // so the reaper reaps it on next reaping
+        context->rax = -ENOMEM;
+        return;
+    }
+
+    child->save_context(context);
+    child->save_floating();
+
+    for (int i = 0; i < TASK_MAX_FDS; i++) { child->files[i] = parent->files[i]; }
+
+    child->address_space = parent->address_space.clone();
+
+    child->regs.rax = 0;
+    context->rax = child->id;
+
+    child->state = child->Running;
+
+    kinfoln("fork(): forked parent %ld into child %ld", parent->id, child->id);
+
+    return;
+}
+
 void sys_exec(Context* context, const char* pathname)
 {
-    /*context->rax = -ENOSYS; // FIXME: Make exec() work under separate address spaces.
-    return;*/
-
     char* kpathname = Syscall::strdup_from_user(pathname);
     if (!kpathname)
     {
@@ -65,6 +100,16 @@ void sys_exec(Context* context, const char* pathname)
 
     // At this point, pretty much nothing can fail.
 
+    task->allocator.free();
+    task->allocator
+        .init(); // If we had enough space for the old bitmap, we should have enough space for the new bitmap.
+
+    task->address_space.clear();
+    task->allocated_stack = (uint64_t)MemoryManager::get_pages_at(
+        0x100000, TASK_PAGES_IN_STACK,
+        MAP_USER | MAP_READ_WRITE); // If we had enough space for the old stack, there should be enough space for the
+                                    // new stack.
+
     ELFImage* image = ELFLoader::load_elf_from_vfs(program);
     ASSERT(image); // If check_elf_image succeeded, load_elf_from_vfs MUST succeed, unless something has gone terribly
                    // wrong.

kernel/src/sys/mem.cpp

@@ -6,6 +6,7 @@
 #include "memory/MemoryManager.h"
 #include "memory/VMM.h"
 #include "misc/utils.h"
+#include "thread/Scheduler.h"
 #include <stddef.h>
 
 #define MAP_READ 1
@@ -72,7 +73,9 @@ void sys_mmap(Context* context, void* address, size_t size, int prot)
         }
     }
     kdbgln("mmap(): %ld pages at any address, %s", Utilities::get_blocks_from_size(PAGE_SIZE, size), format_prot(prot));
-    void* result = MemoryManager::get_pages(Utilities::get_blocks_from_size(PAGE_SIZE, size), real_flags);
+    uint64_t ptr =
+        Scheduler::current_task()->allocator.request_virtual_pages(Utilities::get_blocks_from_size(PAGE_SIZE, size));
+    void* result = MemoryManager::get_pages_at(ptr, Utilities::get_blocks_from_size(PAGE_SIZE, size), real_flags);
     if (result)
     {
         kdbgln("mmap() succeeded: %p", result);
@@ -116,6 +119,8 @@ void sys_munmap(Context* context, void* address, size_t size)
         return;
     }
     uint64_t offset = (uint64_t)address % PAGE_SIZE;
+    Scheduler::current_task()->allocator.free_virtual_pages(((uint64_t)address - offset),
+                                                            Utilities::get_blocks_from_size(PAGE_SIZE, size));
     MemoryManager::release_pages((void*)((uint64_t)address - offset), Utilities::get_blocks_from_size(PAGE_SIZE, size));
     kdbgln("munmap() succeeded");
     context->rax = 0;

kernel/src/thread/Scheduler.cpp

@@ -108,13 +108,10 @@ void Scheduler::add_kernel_task(void (*task)(void))
 Task* Scheduler::create_user_task()
 {
     Task* new_task = new Task;
-    ASSERT(new_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->allocated_stack = (uint64_t)MemoryManager::get_pages_at(
-        0x100000, TASK_PAGES_IN_STACK, MAP_READ_WRITE | MAP_USER); // 16 KB is enough for everyone, right?
-    new_task->regs.rsp = get_top_of_stack(new_task->allocated_stack, TASK_PAGES_IN_STACK);
     new_task->task_sleep = 0;
     new_task->task_time = 0;
     new_task->cpu_time = 0;
@@ -139,6 +136,12 @@ long Scheduler::load_user_task(const char* filename)
     ASSERT(new_task);
     memset(&new_task->regs, 0, sizeof(Context));
     new_task->id = free_tid++;
+    if (!new_task->allocator.init())
+    {
+        delete new_task;
+        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(
@@ -152,6 +155,7 @@ long Scheduler::load_user_task(const char* filename)
         0x100000, TASK_PAGES_IN_STACK, MAP_READ_WRITE | MAP_USER); // 16 KB is enough for everyone, right?
     if (!new_task->allocated_stack)
     {
+        new_task->address_space.destroy();
         delete new_task;
         ELFLoader::release_elf_image(image);
         VMM::switch_back_to_kernel_address_space();
@@ -215,6 +219,7 @@ void Scheduler::reap_task(Task* task)
     }
     if (exiting_task->is_user_task())
     {
+        exiting_task->allocator.free();
         VMM::switch_back_to_kernel_address_space();
         VMM::apply_address_space();
         Interrupts::push_and_enable();

libs/libc/include/luna/syscall.h

@@ -19,6 +19,7 @@
 #define SYS_clock 15
 #define SYS_spawn 16
 #define SYS_mkdir 17
+#define SYS_fork 18
 
 #ifndef __want_syscalls
 #ifdef __cplusplus

libs/libc/src/unistd.cpp

@@ -19,9 +19,10 @@ extern "C"
     {
         NOT_IMPLEMENTED("execvp");
     }
+
     pid_t fork(void)
     {
-        NOT_IMPLEMENTED("fork");
+        return syscall(SYS_fork);
     }
 
     long syscall(long number, ...)
@@ -34,6 +35,7 @@ extern "C"
         {
         case SYS_clock:
         case SYS_yield:
+        case SYS_fork:
         case SYS_gettid: result = __luna_syscall0(number); break;
         case SYS_exit:
         case SYS_close: