Implement an ELFImage struct

This struct allows us to keep track of what memory is used by the loaded executable. For some reason, freeing this memory when the task exits triggers a kernel page fault, so I'm not doing that right now.

kernel/include/sys/elf/ELFLoader.h

@@ -1,8 +1,9 @@
 #pragma once
+#include "sys/elf/Image.h"
 #include <stdint.h>
 
 namespace ELFLoader
 {
-    void* load_elf_from_address(uintptr_t addr);
-    void* load_elf_from_initrd(const char* filename);
+    ELFImage* load_elf_from_address(uintptr_t addr);
+    ELFImage* load_elf_from_initrd(const char* filename);
 }

kernel/include/sys/elf/Image.h

@@ -0,0 +1,15 @@
+#pragma once
+#include <stdint.h>
+
+struct ELFSection
+{
+    uintptr_t base;
+    uint64_t pages;
+};
+
+struct ELFImage
+{
+    uintptr_t entry;
+    uint64_t section_count;
+    ELFSection sections[1];
+};

kernel/include/thread/Task.h

@@ -1,5 +1,6 @@
 #pragma once
 #include "interrupts/Context.h"
+#include "sys/elf/Image.h"
 
 struct Task
 {
@@ -30,6 +31,8 @@ struct Task
     bool floating_saved = false;
 
     bool is_user_task();
+
+    ELFImage* image = nullptr;
 };
 
 void set_context_from_task(Task& task, Context* ctx);

kernel/src/sys/elf/ELFLoader.cpp

@@ -5,6 +5,7 @@
 #include "log/Log.h"
 #include "memory/MemoryManager.h"
 #include "misc/utils.h"
+#include "std/stdlib.h"
 #include "std/string.h"
 #include "sys/elf/ELF.h"
 
@@ -18,7 +19,7 @@ static const char* format_permissions(uint32_t flags)
     return perms;
 }
 
-void* ELFLoader::load_elf_from_initrd(const char* filename)
+ELFImage* ELFLoader::load_elf_from_initrd(const char* filename)
 {
     InitRD::File elf_file = InitRD::open(filename);
     if (!elf_file.addr)
@@ -30,7 +31,7 @@ void* ELFLoader::load_elf_from_initrd(const char* filename)
     return load_elf_from_address((uintptr_t)elf_file.addr);
 }
 
-void* ELFLoader::load_elf_from_address(uintptr_t addr)
+ELFImage* ELFLoader::load_elf_from_address(uintptr_t addr)
 {
     Elf64_Ehdr* elf_ehdr = (Elf64_Ehdr*)addr;
     if (strncmp((const char*)elf_ehdr->e_ident, ELFMAG, SELFMAG) != 0)
@@ -63,6 +64,9 @@ void* ELFLoader::load_elf_from_address(uintptr_t addr)
         kwarnln("ELF file has no PHDRS");
         return 0;
     }
+    ELFImage* image = (ELFImage*)kmalloc(sizeof(ELFImage) - sizeof(ELFSection));
+    memset(image, 0, sizeof(ELFImage) - sizeof(ELFSection));
+    image->entry = elf_ehdr->e_entry;
     int i;
     Elf64_Phdr* phdr;
     for (phdr = (Elf64_Phdr*)((uint64_t)addr + elf_ehdr->e_phoff), i = 0; i < elf_ehdr->e_phnum;
@@ -82,8 +86,19 @@ void* ELFLoader::load_elf_from_address(uintptr_t addr)
                                                        phdr->p_flags & 2 ? MAP_READ_WRITE | MAP_USER : MAP_USER);
             memcpy(buffer, (void*)(addr + phdr->p_offset), phdr->p_filesz);
             memset((void*)((uint64_t)buffer + phdr->p_filesz), 0, phdr->p_memsz - phdr->p_filesz);
+            image = (ELFImage*)krealloc(image, (sizeof(ELFImage) - sizeof(ELFSection)) +
+                                                   (image->section_count + 1) * sizeof(ELFSection));
+            ELFSection& section = image->sections[image->section_count];
+            section.base = (uintptr_t)buffer;
+            section.pages = pages;
+            image->section_count++;
         }
         else { kdbgln("skipping non-loadable segment"); }
     }
-    return (void*)elf_ehdr->e_entry;
+    if (!image->section_count)
+    {
+        kfree(image);
+        return 0;
+    }
+    return image;
 }

kernel/src/thread/Scheduler.cpp

@@ -117,7 +117,14 @@ void Scheduler::load_user_task(const char* filename)
     Task* new_task = new Task;
     ASSERT(new_task);
     new_task->id = free_tid++;
-    new_task->regs.rip = (uint64_t)ELFLoader::load_elf_from_initrd(filename);
+    ELFImage* image = ELFLoader::load_elf_from_initrd(filename);
+    if (!image)
+    {
+        kerrorln("Failed to load %s from initrd", filename);
+        return;
+    }
+    new_task->regs.rip = image->entry;
+    new_task->image = image;
     if (!new_task->regs.rip)
     {
         kwarnln("Failed to load user task %s", filename);
@@ -154,6 +161,15 @@ void Scheduler::reap_task(Task* task)
     kinfoln("reaping task %ld", exiting_task->id);
     if (exiting_task->allocated_stack)
         MemoryManager::release_pages((void*)exiting_task->allocated_stack, TASK_PAGES_IN_STACK);
+    if (exiting_task->image)
+    {
+        for (uint64_t i = 0; i < exiting_task->image->section_count; i++)
+        {
+            ELFSection& section = exiting_task->image->sections[i];
+            kdbgln("Task was using region %lx, which used %ld pages", section.base, section.pages);
+        }
+        kfree(exiting_task->image);
+    }
     delete exiting_task;
 }