From 0252f50c70010193b97cba2f042fee5711084057 Mon Sep 17 00:00:00 2001
From: Qwinci <32550582+Qwinci@users.noreply.github.com>
Date: Fri, 20 Oct 2023 20:29:09 +0300
Subject: [PATCH] mem: Implement CoW

---
 src/mem/vm.c        | 66 +++++++++++++++++++++++++++++++++++++++++++++
 src/mem/vm.h        |  3 +++
 src/sched/process.c | 32 +++++++++++++---------
 3 files changed, 88 insertions(+), 13 deletions(-)

diff --git a/src/mem/vm.c b/src/mem/vm.c
index 1ecb9fa..d652556 100644
--- a/src/mem/vm.c
+++ b/src/mem/vm.c
@@ -6,6 +6,7 @@
 #include "sched/process.h"
 #include "vmem.h"
 #include "string.h"
+#include "utils.h"
 
 static VMem kernel_vmem = {};
 static Mutex KERNEL_VM_LOCK = {};
@@ -164,6 +165,47 @@ void* vm_user_alloc_backed(Process* process, void* at, usize count, PageFlags fl
 	return vm;
 }
 
+void* vm_user_create_cow(Process* process, void* original_map, Mapping* original, void* at) {
+	usize pages = ALIGNUP(original->size, PAGE_SIZE) / PAGE_SIZE;
+	void* vm = vm_user_alloc(process, at, pages);
+	if (!vm) {
+		return NULL;
+	}
+
+	PageFlags page_flags = PF_USER;
+	if ((original->flags & MAPPING_FLAG_R) || (original->flags & MAPPING_FLAG_W)) {
+		page_flags |= PF_READ;
+	}
+	if (original->flags & MAPPING_FLAG_X) {
+		page_flags |= PF_EXEC;
+	}
+
+	for (usize i = 0; i < pages; ++i) {
+		usize phys = arch_virt_to_phys(original_map, (usize) original->base + i * PAGE_SIZE);
+		if (phys) {
+			Page* page = page_from_addr(phys);
+			page->refs += 1;
+			arch_map_page(process->map, (usize) vm + i * PAGE_SIZE, phys, page_flags);
+		}
+	}
+
+	MappingFlags flags = original->flags | MAPPING_FLAG_COW;
+
+	mutex_lock(&process->mapping_lock);
+	if (!process_add_mapping(process, (usize) vm, original->size, flags)) {
+		mutex_unlock(&process->mapping_lock);
+
+		for (usize i = 0; i < pages; ++i) {
+			arch_unmap_page(process->map, (usize) vm + i * PAGE_SIZE, true);
+		}
+
+		vm_user_dealloc(process, vm, pages);
+		return NULL;
+	}
+	mutex_unlock(&process->mapping_lock);
+	return vm;
+}
+
 void* vm_user_alloc_on_demand(Process* process, void* at, usize count, MappingFlags flags, void** kernel_mapping) {
 	void* vm = vm_user_alloc(process, at, count);
 	if (!vm) {
@@ -247,3 +289,27 @@ bool vm_user_dealloc_on_demand(Process* process, void* ptr, usize count, void* k
 	}
 	return true;
 }
+
+bool vm_user_dealloc_cow(Process* process, void* ptr, usize count) {
+	if (!process_remove_mapping(process, (usize) ptr)) {
+		return false;
+	}
+	for (usize i = 0; i < count; ++i) {
+		usize virt = (usize) ptr + i * PAGE_SIZE;
+		usize phys = arch_virt_to_phys(process->map, virt);
+		if (!phys) {
+			continue;
+		}
+		arch_user_unmap_page(process, virt, true);
+
+		Page* page = page_from_addr(phys);
+		if (page->refs == 0) {
+			pfree(page, 1);
+		}
+		else {
+			page->refs -= 1;
+		}
+	}
+	vm_user_dealloc(process, ptr, count);
+	return true;
+}
diff --git a/src/mem/vm.h b/src/mem/vm.h
index 5e8641b..3d9d80a 100644
--- a/src/mem/vm.h
+++ b/src/mem/vm.h
@@ -22,3 +22,6 @@ void* vm_user_alloc_on_demand(Process* process, void* at, usize count, MappingFl
 void vm_user_dealloc_kernel(void* kernel_mapping, usize count);
 bool vm_user_dealloc_backed(Process* process, void* ptr, usize count, void* kernel_mapping);
 bool vm_user_dealloc_on_demand(Process* process, void* ptr, usize count, void* kernel_mapping);
+
+void* vm_user_create_cow(Process* process, void* original_map, Mapping* original, void* at);
+bool vm_user_dealloc_cow(Process* process, void* ptr, usize count);
diff --git a/src/sched/process.c b/src/sched/process.c
index 37d79cb..a1fda31 100644
--- a/src/sched/process.c
+++ b/src/sched/process.c
@@ -4,6 +4,7 @@
 #include "mem/vm.h"
 #include "task.h"
 #include "mem/pmalloc.h"
+#include "mem/utils.h"
 
 Process* ACTIVE_INPUT_PROCESS = NULL;
 
@@ -172,23 +173,28 @@ bool process_handle_fault(Process* process, usize addr) {
 	if (!mapping || !((mapping->flags & MAPPING_FLAG_COW) | (mapping->flags & MAPPING_FLAG_ON_DEMAND))) {
 		return false;
 	}
+	if (!(mapping->flags & MAPPING_FLAG_R)) {
+		return false;
+	}
 
 	assert(addr >= mapping->base);
-	if (mapping->flags & MAPPING_FLAG_ON_DEMAND) {
-		Page* page = pmalloc(1);
-		if (!page) {
-			// todo oom
-			return false;
-		}
-
-		PageFlags flags = flags_to_pf(mapping->flags);
-
-		arch_user_map_page(process, addr & ~(PAGE_SIZE - 1), page->phys, flags);
-		arch_invalidate_mapping(process);
+	Page* page = pmalloc(1);
+	if (!page) {
+		// todo oom
+		return false;
 	}
-	else {
-		assert(!"cow is not implemented");
+
+	if (mapping->flags & MAPPING_FLAG_COW) {
+		memcpy(to_virt(page->phys), (void*) ALIGNDOWN(addr, PAGE_SIZE), PAGE_SIZE);
+		usize orig_addr = arch_virt_to_phys(process->map, ALIGNDOWN(addr, PAGE_SIZE));
+		Page* orig_page = page_from_addr(orig_addr);
+		assert(orig_page->refs > 0);
+		orig_page->refs -= 1;
 	}
 
+	PageFlags flags = flags_to_pf(mapping->flags);
+	arch_user_map_page(process, ALIGNDOWN(addr, PAGE_SIZE), page->phys, flags);
+	arch_invalidate_mapping(process);
+
 	return true;
 }