From e8a8e23a8543224be42798b4c357df67ef30b6f5 Mon Sep 17 00:00:00 2001
From: Danny Holman <dholman@gymli.org>
Date: Tue, 25 Jun 2024 20:44:14 -0500
Subject: arch: i386: kmalloc: don't ignore the last block

Fix a bug in which the kfree function would ignore the last mem_block
struct. This allows the kernel to free a mem_block at the end of the
linked list.

Signed-off-by: Danny Holman <dholman@gymli.org>
---
 arch/i386/kernel/kmalloc.c | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/arch/i386/kernel/kmalloc.c b/arch/i386/kernel/kmalloc.c
index 8118ecd..9c82e39 100644
--- a/arch/i386/kernel/kmalloc.c
+++ b/arch/i386/kernel/kmalloc.c
@@ -2,27 +2,25 @@
 #include <kernel/paging.h>
 #include <kernel/pmem.h>
 #include <kernel/asm.h>
+#include <kernel/vmem.h>
 #include <kernel/panic.h>
 #include <stdint.h>
 
 static uintptr_t heap_start = KEND + 0xC0001000;
-static uintptr_t heap_end = 0xFFFFF000;
+static uintptr_t heap_end = 0xFFBFF000;
 
 static void *kbrk = NULL;
 static struct mem_block *first = NULL;
 
 void _request_page(void) {
-        uintptr_t paddr = pfa_alloc();
-        if (paddr == PFA_ALLOC_ERR)
-                panic("Out of physical memory");
-
+        uintptr_t paddr = pfa_alloc(1);
         uintptr_t vaddr = PGROUNDDN((uintptr_t)kbrk + PAGE_SIZE);
         kbrk += PAGE_SIZE;
-        map_page(NULL, paddr, vaddr, PD_PRES | PD_RW);
+        map_page(paddr, vaddr, PD_PRES | PD_RW);
 }
 
 void kmalloc_init(void) {
-        kbrk = (void*)(heap_start & 0xFFFFF000);
+        kbrk = PGROUNDDN(heap_start);
         _request_page();
         first = (struct mem_block*)kbrk;
         first->start = (uintptr_t)kbrk + sizeof(struct mem_block);
@@ -39,12 +37,17 @@ void* kmalloc(size_t sz) {
 
         struct mem_block *temp = first;
         while (temp->next != NULL) {
-                if (temp->size <= sz && temp->alloc == 0)
+                if (temp->size <= sz && temp->alloc == 0) {
+                        temp->alloc = 1;
                         return (void*)temp->start;
+                }
                 temp = temp->next;
         }
 
         temp->next = (struct mem_block*)((uintptr_t)temp + sizeof(struct mem_block) + temp->size);
+        if (temp->next > kbrk)
+                _request_page();
+        temp->next->prev = temp;
         temp = temp->next;
         temp->start = (uintptr_t)(temp + sizeof(struct mem_block));
         temp->size = sz;
@@ -55,11 +58,12 @@ void* kmalloc(size_t sz) {
 
 void kfree(void *ptr) {
         struct mem_block *temp = first;
-        while (temp->next != NULL) {
+        while (temp != NULL) {
                 if (temp->start == (uintptr_t)ptr) {
                         temp->alloc = 0;
                         return;
                 }
+                temp = temp->next;
         }
         panic("Attempted to free memory not alloc'd");
 }
-- 
cgit v1.2.3