Add split_accessed_hugepages argument to Machine::prepare_cow

Author: fwsGonzo

lib/tinykvm/amd64/paging.cpp

@@ -164,9 +164,9 @@ uint64_t setup_amd64_paging(vMemory& memory,
 	const uint64_t vdso_pdpt_addr = pml4_addr + PD_END + 0x0;
 	const uint64_t vsyscall_pd_addr = pml4_addr + PD_END + 0x1000;
 	const uint64_t vsyscall_pt_addr = pml4_addr + PD_END + 0x2000;
-	auto *vdso_pdpt = (uint64_t *)(pagetable + PD_END + 0x0000);
-	auto *vsyscall_pd = (uint64_t *)(pagetable + PD_END + 0x1000);
-	auto *vsyscall_pt = (uint64_t *)(pagetable + PD_END + 0x2000);
+	auto* vdso_pdpt = (uint64_t *)(pagetable + PD_END + 0x0000);
+	auto* vsyscall_pd = (uint64_t *)(pagetable + PD_END + 0x1000);
+	auto* vsyscall_pt = (uint64_t *)(pagetable + PD_END + 0x2000);
 
 	// next free page for ELF loader
 	uint64_t free_page = pml4_addr + PD_END + 0x3000;
@@ -390,6 +390,10 @@ uint64_t setup_amd64_paging(vMemory& memory,
 			| (base_giga_page << 30) | (base_2mb_page << 21) | (i << 12);
 	}
 
+	if (free_page >= memory.physbase + (2ULL << 20)) {
+		throw MachineException("Pagetable setup exceeded 2MB limit");
+	}
+
 	/* Verify a kernel page */
 	page_at(memory, memory.physbase + 0x1000,
 		[&memory] (uint64_t addr, uint64_t& entry, size_t size) {
@@ -532,8 +536,9 @@ void foreach_page(vMemory& memory, foreach_page_t callback, bool skip_oob_addres
 					{
 						if (pd[k] & PDE64_PRESENT) {
 							const auto [pt_base, pt_mem, pt_size] = pt_from_index(k, pd_base, pd);
+							const bool is_2mb_page = (pd[k] & PDE64_PS) != 0;
 							callback(pt_base, pd[k], pt_size);
-							if (!(pd[k] & PDE64_PS)) { // not 2MB page
+							if (!is_2mb_page) {
 								auto* pt = memory.page_at(pt_mem);
 								for (uint64_t e = 0; e < 512; e++) {
 									const auto [pte_base, pte_mem, pte_size] = pte_from_index(e, pt_base, pt);
@@ -554,19 +559,49 @@ void foreach_page(const vMemory& mem, foreach_page_t callback, bool skip_oob_add
 	foreach_page(const_cast<vMemory&>(mem), std::move(callback), skip_oob_addresses);
 }
 
-void foreach_page_makecow(vMemory& mem, uint64_t kernel_end, uint64_t shared_memory_boundary)
+void foreach_page_makecow(vMemory& mem, uint64_t kernel_end,
+	uint64_t shared_memory_boundary, bool split_accessed_hugepages)
 {
 	if (UNLIKELY(shared_memory_boundary < kernel_end)) {
 		memory_exception("Shared memory boundary was illegal (zero)", shared_memory_boundary, 0u);
 	}
 	foreach_page(mem,
-	[=] (uint64_t addr, uint64_t& entry, size_t /*size*/) {
+	[=, m = &mem] (uint64_t addr, uint64_t& entry, size_t size) {
 		if (addr < shared_memory_boundary) {
 			const uint64_t flags = (PDE64_PRESENT | PDE64_RW);
 			if ((entry & flags) == flags) {
 				entry &= ~PDE64_RW;
 				entry |= PDE64_CLONEABLE | PDE64_G; // Global bit for read-only pages
 			}
+			if ((entry & PDE64_ACCESSED) != 0)
+			{
+				// Since this page has been accessed, check if it's a 2MB leaf page
+				// and if it is, split it into 4k pages with the ACCESS bit removed.
+				if (size == (1ULL << 21) && (entry & PDE64_PS) != 0 && split_accessed_hugepages)
+				{
+					// Split 2MB page into 4k pages
+					const uint64_t pd_base = entry & PDE64_ADDR_MASK;
+					// Allocate new page for page table
+					const auto new_page = m->allocate_unmapped_kernelpage();
+					if (new_page.addr != 0) {
+						// Set default attributes + new page table
+						const uint64_t pd_entry_flags = entry & ~(PDE64_ADDR_MASK | PDE64_PS | PDE64_ACCESSED);
+						entry = pd_entry_flags | new_page.addr;
+						// Fill new page with default attributes
+						uint64_t* pagetable = new_page.pmem;
+						for (uint64_t i = 0; i < 512; i++) {
+							// Set writable 4k attributes
+							uint64_t addr4k = pd_base | (i << 12);
+							pagetable[i] = pd_entry_flags | addr4k;
+						}
+						return;
+					} // new_page.addr
+					// Failing to allocate, so we simply won't
+					// split the page, but this may lead to extra
+					// memory usage if forks don't need the whole
+					// 2MB page.
+				}
+			}
 		}
 		// Clear accessed bit for *all* pages
 		// Doing this makes it possible to send a dummy request to estimate which

lib/tinykvm/amd64/paging.hpp

@@ -14,7 +14,7 @@ extern void print_pagetables(const vMemory&);
 using foreach_page_t = std::function<void(uint64_t, uint64_t&, size_t)>;
 extern void foreach_page(vMemory&, foreach_page_t callback, bool skip_oob_addresses = true);
 extern void foreach_page(const vMemory&, foreach_page_t callback, bool skip_oob_addresses = true);
-extern void foreach_page_makecow(vMemory&, uint64_t kernel_end, uint64_t shared_memory_boundary);
+extern void foreach_page_makecow(vMemory&, uint64_t kernel_end, uint64_t shared_memory_boundary, bool split_accessed_hugepages = false);
 extern std::vector<std::pair<uint64_t, uint64_t>> get_accessed_pages(const vMemory& memory);
 
 extern void page_at(vMemory&, uint64_t addr, foreach_page_t, bool ignore_missing = false);

lib/tinykvm/machine.hpp

@@ -163,6 +163,7 @@ struct Machine
 	address_t exit_address() const noexcept;
 	void set_stack_address(address_t addr) { this->m_stack_address = addr; }
 	address_t kernel_end_address() const noexcept { return m_kernel_end; }
+	void set_kernel_end_address(address_t addr) { m_kernel_end = addr; }
 	address_t max_address() const noexcept { return memory.physbase + memory.size; }
 
 	static constexpr uint64_t BRK_MAX = 0x22000;
@@ -241,7 +242,8 @@ struct Machine
 	/* Make VM copy-on-write in order to support fast forking.
 	   When @max_work_mem is non-zero, the master VM can still
 	   be used after preparation. */
-	void prepare_copy_on_write(size_t max_work_mem = 0, uint64_t shared_memory_boundary = UINT64_MAX);
+	void prepare_copy_on_write(size_t max_work_mem = 0, uint64_t shared_memory_boundary = UINT64_MAX,
+		bool split_accessed_hugepages = false);
 	void set_main_memory_writable(bool v) { memory.main_memory_writes = v; }
 	bool is_forked() const noexcept { return m_forked; }
 	bool uses_cow_memory() const noexcept { return m_forked || m_prepped; }

lib/tinykvm/memory.cpp

@@ -573,6 +573,48 @@ void vMemory::increment_unlocked_pages(size_t pages)
 	}
 }
 
+MemoryBank::Page vMemory::allocate_unmapped_kernelpage()
+{
+	if (this->machine.is_forked()) {
+		// It's too dangerous as forked VMs may have active
+		// TLB entries that user memory mappings that will
+		// suddenly become kernel mappings.
+		memory_exception("Cannot allocate unmapped kernel page in forked VM", 0, 0);
+	}
+
+	MemoryBank::Page result{};
+	result.pmem = nullptr;
+	result.addr = 0;
+	// 1. Check if kernel end is < 2MB userspace area
+	const uint64_t kernel_end = machine.kernel_end_address();
+	if (kernel_end < physbase + (1ULL << 21)) {
+		// Find the page directory entry for the kernel end
+		tinykvm::page_at(*this, kernel_end,
+		[&](uint64_t addr, uint64_t& entry, uint64_t page_size) mutable {
+			static constexpr uint64_t PDE64_ADDR_MASK = ~0x8000000000000FFF;
+			// Grab physical from the page directory entry
+			const uint64_t phys_addr = entry & PDE64_ADDR_MASK;
+			uint64_t* phys_mem = this->page_at(phys_addr);
+			if (phys_addr < PT_ADDR || phys_mem < (uint64_t*)(this->ptr + PT_ADDR)) {
+				memory_exception("Invalid page directory entry for kernel end", addr, page_size);
+			}
+			// Set up the page directory entry with:
+			// Kernel + unmapped + no-execute
+			entry = phys_addr | PDE64_NX;
+
+			result = MemoryBank::Page{
+				phys_mem,
+				kernel_end,
+				PageSize(),
+				false
+			};
+			this->machine.set_kernel_end_address(kernel_end + PageSize());
+		}, false);
+	}
+	//printf("*** Allocated unmapped kernel page at 0x%lX\n", result.addr);
+	return result;
+}
+
 uint64_t vMemory::expectedUsermodeFlags() const noexcept
 {
 	uint64_t flags = PDE64_PRESENT | PDE64_USER | PDE64_RW;

lib/tinykvm/memory.hpp

@@ -72,6 +72,7 @@ struct vMemory {
 	char *get_writable_page(uint64_t addr, uint64_t flags, bool zeroes, bool dirty);
 	MemoryBank::Page new_page();
 	MemoryBank::Page new_hugepage();
+	MemoryBank::Page allocate_unmapped_kernelpage();
 
 	bool compare(const vMemory& other);
 	/* When a main VM has direct memory writes enabled, it can

lib/tinykvm/vcpu.cpp

@@ -416,11 +416,10 @@ void vCPU::set_vcpu_table_at(unsigned index, int value)
 	}
 }
 
-void Machine::prepare_copy_on_write(size_t max_work_mem, uint64_t shared_memory_boundary)
+void Machine::prepare_copy_on_write(size_t max_work_mem,
+	uint64_t shared_memory_boundary, bool split_accessed_hugepages)
 {
 	this->m_prepped = true;
-	if (max_work_mem == 0) {
-	}
 
 	/* Make each writable page read-only, causing page fault.
 	   any page after the @shared_memory_boundary is untouched,
@@ -457,13 +456,13 @@ void Machine::prepare_copy_on_write(size_t max_work_mem, uint64_t shared_memory_
 		vcpu.set_special_registers(sregs);
 		this->enter_usermode();
 
-		foreach_page_makecow(this->memory, kernel_end_address(), shared_memory_boundary);
+		foreach_page_makecow(this->memory, kernel_end_address(), shared_memory_boundary, split_accessed_hugepages);
 		return;
 	}
 
 	/* This call makes this VM usable after making every page in the
 	   page tables read-only, enabling memory through page faults. */
-	foreach_page_makecow(this->memory, kernel_end_address(), shared_memory_boundary);
+	foreach_page_makecow(this->memory, kernel_end_address(), shared_memory_boundary, split_accessed_hugepages);
 	this->setup_cow_mode(this);
 }
 void Machine::setup_cow_mode(const Machine* other)