user/elf: free memory not belonging to the elf when jumping to it

4 files changed, 70 insertions, 13 deletions

diff --git a/src/kernel/arch/amd64/pagedir.c b/src/kernel/arch/amd64/pagedir.c
index f043289..f4fbc09 100644
--- a/src/kernel/arch/amd64/pagedir.c
+++ b/src/kernel/arch/amd64/pagedir.c
@@ -21,6 +21,14 @@ static bool addr_canonical(const __user void *addr) {
 	return (n == 0) || ((~n) << addr_bits == 0);
 }
 
+/* the types here are idiotic because C is idiotic */
+static __user void *addr_canonize(const __user void *addr) {
+	union virt_addr v = {.full = (void __force*)addr};
+	v.sign = (((uintptr_t)addr >> 47) & 1) * 0xFFFF;
+	assert(addr_canonical(addr));
+	return v.full;
+}
+
 
 struct pagedir *pagedir_new(void) {
 	struct pagedir *dir = page_alloc(1);
@@ -81,11 +89,60 @@ get_entry(struct pagedir *dir, const void __user *virt) {
 	return pte;
 }
 
-void *pagedir_unmap(struct pagedir *dir, void __user *virt) {
-	pe_generic_t *page = get_entry(dir, virt);
-	if (!page) return NULL;
-	page->present = false;
-	return addr_extract(*page);
+void pagedir_unmap_user(struct pagedir *dir, void __user *virt, size_t len) {
+	// TODO rewrite this
+	const void __user *end = addr_canonize(virt + len);
+	union virt_addr v = {.full = virt};
+	v.off_4k = 0;
+
+	for (int pml4i = v.pml4; v.full < end && pml4i < 512; pml4i++) {
+		v.pml4 = pml4i;
+		v.full = addr_canonize(v.full);
+		if (v.full >= end) return;
+		if (!dir->e[v.pml4].present) {
+			v.pdpt = 0;
+			v.pd = 0;
+			v.pt = 0;
+			continue;
+		}
+		assert(!dir->e[v.pml4].large);
+		pe_generic_t *pdpt = addr_extract(dir->e[v.pml4]);
+
+		for (int pdpti = v.pdpt; v.full < end && pdpti < 512; pdpti++) {
+			v.pdpt = pdpti;
+			if (v.full >= end) return;
+			if (!pdpt[v.pdpt].present) {
+				v.pd = 0;
+				v.pt = 0;
+				continue;
+			}
+			assert(!pdpt[v.pdpt].large);
+			pe_generic_t *pd = addr_extract(pdpt[v.pdpt]);
+
+			for (int pdi = v.pd; v.full < end && pdi < 512; pdi++) {
+				v.pd = pdi;
+				if (v.full >= end) return;
+				if (!pd[v.pd].present) {
+					v.pt = 0;
+					continue;
+				}
+				assert(!pd[v.pd].large);
+				pe_generic_t *pt = addr_extract(pd[v.pd]);
+
+				for (int pti = v.pt; v.full < end && pti < 512; pti++) {
+					v.pt = pti;
+					if (v.full >= end) return;
+					if (!pt[v.pt].present) continue;
+					if (!pt[v.pt].user) continue;
+					page_free(addr_extract(pt[v.pt]), 1);
+					pt[v.pt].present = 0;
+				}
+				v.pt = 0;
+			}
+			v.pd = 0;
+		}
+		v.pdpt = 0;
+	}
 }
 
 void pagedir_map(struct pagedir *dir, void __user *virt, void *phys,
diff --git a/src/kernel/arch/generic.h b/src/kernel/arch/generic.h
index d1ff292..941bfa4 100644
--- a/src/kernel/arch/generic.h
+++ b/src/kernel/arch/generic.h
@@ -30,7 +30,7 @@ struct pagedir *pagedir_new(void);
 struct pagedir *pagedir_copy(const struct pagedir *orig);
 
 void pagedir_free(struct pagedir *);
-void *pagedir_unmap(struct pagedir *dir, void __user *virt);
+void pagedir_unmap_user(struct pagedir *dir, void __user *virt, size_t len);
 void pagedir_map(struct pagedir *dir, void __user *virt, void *phys,
                  bool user, bool writeable);
 bool pagedir_iskern(struct pagedir *, const void __user *virt);
diff --git a/src/kernel/mem/alloc.c b/src/kernel/mem/alloc.c
index e0cc5f4..8d6fa5b 100644
--- a/src/kernel/mem/alloc.c
+++ b/src/kernel/mem/alloc.c
@@ -101,6 +101,8 @@ void *page_zalloc(size_t pages) {
 
 // frees `pages` consecutive pages starting from *first
 void page_free(void *first_addr, size_t pages) {
+	if (first_addr < page_bitmap_start) // TODO unimplemented
+		return;
 	size_t first = (uintptr_t)(first_addr - page_bitmap_start) / PAGE_SIZE;
 	for (size_t i = 0; i < pages; i++) {
 		assert(bitmap_get(first + i));
diff --git a/src/kernel/syscalls.c b/src/kernel/syscalls.c
index a2ae5dd..f884fd3 100644
--- a/src/kernel/syscalls.c
+++ b/src/kernel/syscalls.c
@@ -317,6 +317,11 @@ void __user *_syscall_memflag(void __user *addr, size_t len, int flags) {
 			SYSCALL_RETURN((uintptr_t)addr);
 	}
 
+	if (!(flags & MEMFLAG_PRESENT)) {
+		pagedir_unmap_user(pages, addr, len);
+		SYSCALL_RETURN((uintptr_t)addr);
+	}
+
 
 	for (userptr_t iter = addr; iter < addr + len; iter += PAGE_SIZE) {
 		if (pagedir_iskern(pages, iter)) {
@@ -325,13 +330,6 @@ void __user *_syscall_memflag(void __user *addr, size_t len, int flags) {
 		}
 
 		phys = pagedir_virt2phys(pages, iter, false, false);
-
-		if (!(flags & MEMFLAG_PRESENT)) {
-			if (phys)
-				page_free(pagedir_unmap(pages, iter), 1);
-			continue;
-		}
-
 		if (!phys) {
 			phys = page_alloc(1);
 			memset(phys, 0, PAGE_SIZE); // TODO somehow test this