1 files changed, 100 insertions, 1 deletions
diff --git a/kernel/vm/brk.c b/kernel/vm/brk.c
index 46d6fc2..69a315f 100644
--- a/kernel/vm/brk.c
+++ b/kernel/vm/brk.c
@@ -53,6 +53,105 @@
  */
 long do_brk(void *addr, void **ret)
 {
-    NOT_YET_IMPLEMENTED("VM: do_brk");
+    // NOT_YET_IMPLEMENTED("VM: do_brk");
+
+    // If addr is NULL, return the current break
+    if (addr == NULL) 
+    {
+        *ret = curproc->p_brk;
+        return 0;
+    }
+
+    // Check if the address is within the valid range
+    if ((uintptr_t)addr > USER_MEM_HIGH)
+    {
+        return -ENOMEM;
+    }
+
+    // Check if the address is within the valid range
+    if (addr < curproc->p_start_brk)
+    {
+        return -ENOMEM;
+    }
+
+    // Check if the address is the same as the current break
+    // if (addr == curproc->p_brk)
+    // {
+    //     *ret = curproc->p_brk;
+    //     return 0;
+    // }
+
+    // Check if the address is page aligned
+    uintptr_t addr_page_aligned = ADDR_TO_PN(PAGE_ALIGN_UP(addr));
+    uintptr_t p_brk_page_aligned = ADDR_TO_PN(PAGE_ALIGN_UP(curproc->p_brk));
+    uintptr_t p_start_brk_page_aligned = ADDR_TO_PN(PAGE_ALIGN_UP(curproc->p_start_brk));
+
+    // Lookup the vmarea that represents the heap
+    vmarea_t *heap_vmarea = vmmap_lookup(curproc->p_vmmap, p_start_brk_page_aligned);
+
+    // Check if the address is the same as the current break
+    // If so, set rets and end here
+    if (addr_page_aligned == p_brk_page_aligned)
+    {
+        curproc->p_brk = addr;
+        *ret = addr;
+        return 0;
+    }
+
+    // Check the three cases, whether the heap needs to be created, modified or shrinked
+    if (heap_vmarea == NULL)
+    {
+        // Create the heap
+        long ret = vmmap_is_range_empty(curproc->p_vmmap, p_start_brk_page_aligned, addr_page_aligned - p_brk_page_aligned);
+        if (!ret)
+        {
+            // On fail, return -ENOMEM
+            return -ENOMEM;
+        }
+
+        // Map the heap
+        int flags = MAP_PRIVATE | MAP_ANON;
+        int prot = PROT_READ | PROT_WRITE;
+        ret = vmmap_map(
+            curproc->p_vmmap, NULL, 
+            p_start_brk_page_aligned, 
+            addr_page_aligned - p_start_brk_page_aligned, 
+            prot, flags, 
+            0, 
+            VMMAP_DIR_LOHI, &heap_vmarea
+        );
+        if (ret < 0)
+        {
+            // On fail, return ret
+            return ret;
+        }
+    }
+    else if (addr_page_aligned < p_brk_page_aligned)
+    {
+        // Shrink the heap
+        long ret = vmmap_remove(curproc->p_vmmap, addr_page_aligned, p_brk_page_aligned - addr_page_aligned);
+        if (ret < 0)
+        {
+            // On fail, return ret
+            return ret;
+        }
+    }
+    else
+    {
+        // Modify the heap
+        long ret = vmmap_is_range_empty(curproc->p_vmmap, p_brk_page_aligned, addr_page_aligned - p_brk_page_aligned);
+        if (!ret)
+        {
+            // On fail, return -ENOMEM
+            return -ENOMEM;
+        }
+        // Update the heap
+        heap_vmarea->vma_end = addr_page_aligned;
+    }
+
+    
+    // Update rets & return 0 on success
+    curproc->p_brk = addr;
+    *ret = addr;
     return 0;
 }