From 636c308993ea90040745d44d5b93ad10978279d6 Mon Sep 17 00:00:00 2001
From: Bananymous <bananymousosq@gmail.com>
Date: Sun, 28 May 2023 20:37:39 +0300
Subject: [PATCH] Kernel: fork() now copies allocation done through
 GeneralAllocator

---
 .../include/kernel/Memory/GeneralAllocator.h  |  2 +
 kernel/kernel/Memory/GeneralAllocator.cpp     | 42 +++++++++++++++++++
 kernel/kernel/Process.cpp                     |  6 +--
 userspace/test/test.cpp                       |  7 +++-
 4 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/kernel/include/kernel/Memory/GeneralAllocator.h b/kernel/include/kernel/Memory/GeneralAllocator.h
index 2c4dee55..19a5bb4f 100644
--- a/kernel/include/kernel/Memory/GeneralAllocator.h
+++ b/kernel/include/kernel/Memory/GeneralAllocator.h
@@ -19,6 +19,8 @@ namespace Kernel
 		vaddr_t allocate(size_t);
 		bool deallocate(vaddr_t);
 
+		BAN::ErrorOr<GeneralAllocator*> clone(MMU&);
+
 	private:
 		struct Allocation
 		{
diff --git a/kernel/kernel/Memory/GeneralAllocator.cpp b/kernel/kernel/Memory/GeneralAllocator.cpp
index 076a28d6..c2f24b75 100644
--- a/kernel/kernel/Memory/GeneralAllocator.cpp
+++ b/kernel/kernel/Memory/GeneralAllocator.cpp
@@ -1,4 +1,5 @@
 #include <kernel/Memory/GeneralAllocator.h>
+#include <kernel/Memory/MMUScope.h>
 #include <kernel/Process.h>
 
 namespace Kernel
@@ -61,4 +62,45 @@ namespace Kernel
 		return false;
 	}
 
+	BAN::ErrorOr<GeneralAllocator*> GeneralAllocator::clone(MMU& new_mmu)
+	{
+		GeneralAllocator* allocator = new GeneralAllocator(new_mmu);
+		if (allocator == nullptr)
+			return BAN::Error::from_errno(ENOMEM);
+
+		MMUScope _(m_mmu);
+		ASSERT(m_mmu.is_page_free(0));
+
+		for (auto& allocation : m_allocations)
+		{
+			Allocation new_allocation;
+			ASSERT(new_mmu.is_range_free(allocation.address, allocation.pages.size() * PAGE_SIZE));
+
+			new_allocation.address = allocation.address;
+			MUST(new_allocation.pages.reserve(allocation.pages.size()));
+
+			uint8_t flags = m_mmu.get_page_flags(allocation.address);
+			for (size_t i = 0; i < allocation.pages.size(); i++)
+			{
+				paddr_t paddr = Heap::get().take_free_page();
+				ASSERT(paddr);
+
+				vaddr_t vaddr = allocation.address + i * PAGE_SIZE;
+
+				MUST(new_allocation.pages.push_back(paddr));
+				new_mmu.map_page_at(paddr, vaddr, flags);
+
+				m_mmu.map_page_at(paddr, 0, MMU::Flags::ReadWrite | MMU::Flags::Present);
+				m_mmu.invalidate(0);
+				memcpy((void*)0, (void*)vaddr, PAGE_SIZE);
+			}
+
+			MUST(allocator->m_allocations.push_back(BAN::move(new_allocation)));
+		}
+		m_mmu.unmap_page(0);
+		m_mmu.invalidate(0);
+
+		return allocator;
+	}
+
 }
\ No newline at end of file
diff --git a/kernel/kernel/Process.cpp b/kernel/kernel/Process.cpp
index 425aa650..591cf4bc 100644
--- a/kernel/kernel/Process.cpp
+++ b/kernel/kernel/Process.cpp
@@ -218,10 +218,10 @@ namespace Kernel
 		ASSERT(m_threads.front() == &Thread::current());
 
 		//for (auto& allocator : m_fixed_width_allocators)
-		//	MUST(forked->m_fixed_width_allocators.push_back(allocator.clone()));
+		//	MUST(forked->m_fixed_width_allocators.push_back(MUST(allocator->clone(forked->mmu()))));
 
-		//if (m_general_allocator)
-		//	forked->m_general_allocator = m_general_allocator->clone();
+		if (m_general_allocator)
+			forked->m_general_allocator = MUST(m_general_allocator->clone(forked->mmu()));
 
 		Thread* thread = MUST(m_threads.front()->clone(forked, rsp, rip));
 		forked->add_thread(thread);
diff --git a/userspace/test/test.cpp b/userspace/test/test.cpp
index 8ba9b1e6..0f64e7ab 100644
--- a/userspace/test/test.cpp
+++ b/userspace/test/test.cpp
@@ -8,15 +8,18 @@
 
 int main()
 {
+	char* string = (char*)malloc(5000);
+	strcpy(string, "Hello");
+
 	printf("forking\n");
 
 	pid_t pid = fork();
 	if (pid == 0)
 	{
-		printf("child\n");
+		printf("child '%s'\n", string);
 		return 0;
 	}
-	printf("parent\n");
+	printf("parent '%s'\n", string);
 
 	return 0;
 }