heap (brk / sbrk)

5 files changed, 135 insertions, 8 deletions

diff --git a/kernel/include/comus/procs.h b/kernel/include/comus/procs.h
index eb6c54f..cf3f90d 100644
--- a/kernel/include/comus/procs.h
+++ b/kernel/include/comus/procs.h
@@ -57,6 +57,10 @@ struct pcb {
 	size_t priority;
 	size_t ticks;
 
+	// heap
+	char *heap_start;
+	size_t heap_len;
+
 	// elf metadata
 	Elf64_Ehdr elf_header;
 	Elf64_Phdr elf_segments[N_ELF_SEGMENTS];
diff --git a/kernel/memory/virtalloc.c b/kernel/memory/virtalloc.c
index 863959e..0cbba33 100644
--- a/kernel/memory/virtalloc.c
+++ b/kernel/memory/virtalloc.c
@@ -182,7 +182,7 @@ void *virtaddr_alloc(struct virt_ctx *ctx, int n_pages)
 		if (length < n_length)
 			continue;
 
-		return (void*)node->start;
+		return (void *)node->start;
 	}
 
 	return NULL;
@@ -205,12 +205,12 @@ int virtaddr_take(struct virt_ctx *ctx, const void *virt, int n_pages)
 			continue;
 
 		// create new node on left
-		if (node->start < (uintptr_t) virt) {
+		if (node->start < (uintptr_t)virt) {
 			struct virt_addr_node *left = get_node(ctx);
 			left->next = node;
 			left->prev = node->prev;
 			left->start = node->start;
-			left->end = (uintptr_t) virt;
+			left->end = (uintptr_t)virt;
 			left->is_used = true;
 			left->is_alloc = false;
 			node->prev->next = left;
@@ -218,11 +218,11 @@ int virtaddr_take(struct virt_ctx *ctx, const void *virt, int n_pages)
 		}
 
 		// create new node on right
-		if (node->end > (uintptr_t) virt + n_length) {
+		if (node->end > (uintptr_t)virt + n_length) {
 			struct virt_addr_node *right = get_node(ctx);
 			right->prev = node;
 			right->next = node->next;
-			right->start = (uintptr_t) virt + n_length;
+			right->start = (uintptr_t)virt + n_length;
 			right->end = node->end;
 			right->is_used = true;
 			right->is_alloc = false;
@@ -230,7 +230,7 @@ int virtaddr_take(struct virt_ctx *ctx, const void *virt, int n_pages)
 			node->next = right;
 		}
 
-		node->start = (uintptr_t) virt;
+		node->start = (uintptr_t)virt;
 		node->end = node->start + n_length;
 		node->is_alloc = true;
 		node->is_used = true;
diff --git a/kernel/syscall.c b/kernel/syscall.c
index 16eb243..dd9bc16 100644
--- a/kernel/syscall.c
+++ b/kernel/syscall.c
@@ -1,3 +1,4 @@
+#include "lib/kio.h"
 #include <comus/cpu.h>
 #include <comus/syscalls.h>
 #include <comus/drivers/acpi.h>
@@ -7,6 +8,7 @@
 #include <comus/procs.h>
 #include <comus/time.h>
 #include <lib.h>
+#include <stddef.h>
 
 static struct pcb *pcb;
 
@@ -201,6 +203,68 @@ __attribute__((noreturn)) static int sys_poweroff(void)
 	acpi_shutdown();
 }
 
+static void *pcb_update_heap(intptr_t increment)
+{
+	char *curr_brk;
+	size_t curr_pages, new_pages, new_len;
+
+	new_len = pcb->heap_len + increment;
+	new_pages = (new_len + PAGE_SIZE - 1) / PAGE_SIZE;
+
+	curr_brk = pcb->heap_start + pcb->heap_len;
+	curr_pages = (pcb->heap_len + PAGE_SIZE - 1) / PAGE_SIZE;
+	if (pcb->heap_len == 0)
+		curr_pages = 0;
+
+	// do nothing i guess
+	if (increment == 0 || new_pages == curr_pages)
+		return curr_brk;
+
+	// unmap pages if decreasing
+	if (new_pages < curr_pages) {
+		void *new_end = pcb->heap_start + new_pages * PAGE_SIZE;
+		mem_free_pages(pcb->memctx, new_end);
+		pcb->heap_len = new_pages * PAGE_SIZE;
+	}
+
+	// map pages if increasing
+	else {
+		void *curr_end = pcb->heap_start + curr_pages * PAGE_SIZE;
+		if (mem_alloc_pages_at(pcb->memctx, new_pages - curr_pages, curr_end,
+							   F_PRESENT | F_WRITEABLE | F_UNPRIVILEGED) ==
+			NULL)
+			return NULL;
+		pcb->heap_len = new_pages * PAGE_SIZE;
+	}
+
+	return curr_brk;
+}
+
+static int sys_brk(void)
+{
+	RET(void *, brk);
+	ARG1(const void *, addr);
+
+	// cannot make heap smaller than start
+	if ((const char *)addr < pcb->heap_start) {
+		*brk = NULL;
+		return 0;
+	}
+
+	*brk = pcb_update_heap((intptr_t)addr -
+						   ((intptr_t)pcb->heap_start + pcb->heap_len));
+	return 0;
+}
+
+static int sys_sbrk(void)
+{
+	RET(void *, brk);
+	ARG1(intptr_t, increment);
+
+	*brk = pcb_update_heap(increment);
+	return 0;
+}
+
 static int sys_drm(void)
 {
 	ARG1(void **, res_fb);
@@ -255,8 +319,8 @@ static int (*syscall_tbl[N_SYSCALLS])(void) = {
 	[SYS_getpid] = sys_getpid,	 [SYS_getppid] = sys_getppid,
 	[SYS_gettime] = sys_gettime, [SYS_getprio] = sys_getprio,
 	[SYS_setprio] = sys_setprio, [SYS_kill] = sys_kill,
-	[SYS_sleep] = sys_sleep,	 [SYS_brk] = NULL,
-	[SYS_sbrk] = NULL,			 [SYS_poweroff] = sys_poweroff,
+	[SYS_sleep] = sys_sleep,	 [SYS_brk] = sys_brk,
+	[SYS_sbrk] = sys_sbrk,		 [SYS_poweroff] = sys_poweroff,
 	[SYS_drm] = sys_drm,		 [SYS_ticks] = sys_ticks,
 };
 
diff --git a/kernel/user.c b/kernel/user.c
index 8f626a7..fc1167d 100644
--- a/kernel/user.c
+++ b/kernel/user.c
@@ -71,6 +71,10 @@ static int user_load_segment(struct pcb *pcb, struct disk *disk, int idx)
 		total_read += read;
 	}
 
+	// update heap end
+	if (hdr.p_vaddr + mem_pages * PAGE_SIZE > (uint64_t)pcb->heap_start)
+		pcb->heap_start = (void *)(hdr.p_vaddr + mem_pages * PAGE_SIZE);
+
 	kunmapaddr(mapADDR);
 	return 0;
 }
@@ -79,6 +83,9 @@ static int user_load_segments(struct pcb *pcb, struct disk *disk)
 {
 	int ret = 0;
 
+	pcb->heap_start = NULL;
+	pcb->heap_len = 0;
+
 	if (load_buffer == NULL)
 		if ((load_buffer = kalloc(BLOCK_SIZE)) == NULL)
 			return 1;
@@ -86,6 +93,12 @@ static int user_load_segments(struct pcb *pcb, struct disk *disk)
 	for (int i = 0; i < pcb->n_elf_segments; i++)
 		if ((ret = user_load_segment(pcb, disk, i)))
 			return ret;
+
+	if (pcb->heap_start == NULL) {
+		WARN("No loadable ELF segments found.");
+		return 1;
+	};
+
 	return 0;
 }
 
diff --git a/user/heap.c b/user/heap.c
new file mode 100644
index 0000000..ffa107b
--- /dev/null
+++ b/user/heap.c
@@ -0,0 +1,46 @@
+
+#include <unistd.h>
+#include <string.h>
+#include <stdio.h>
+
+#define PAGE_SIZE 4096
+
+int main(void)
+{
+	void *start, *brk;
+
+	start = sbrk(0);
+	printf("heap start: %p\n", start);
+
+	// test extending
+	if (sbrk(PAGE_SIZE) == NULL) {
+		fprintf(stderr, "failed to extend break\n");
+		return 1;
+	}
+
+	brk = sbrk(0);
+	printf("new break: %p\n", brk);
+
+	// test reextending
+	if (sbrk(PAGE_SIZE) == NULL) {
+		fprintf(stderr, "failed to extend break\n");
+		return 1;
+	}
+
+	brk = sbrk(0);
+	printf("new new break: %p\n", brk);
+
+	// test shrinking
+	if (sbrk(-PAGE_SIZE) == NULL) {
+		fprintf(stderr, "failed to shrink break\n");
+		return 1;
+	}
+
+	brk = sbrk(0);
+	printf("new new new break: %p\n", brk);
+
+	// test write
+	memset(start, 1, PAGE_SIZE);
+
+	return 0;
+}