Merge pull request #25 from varnish/constrained_direct_memory_writes

Implement constrained direct memory writes

Author: fwsGonzo

lib/tinykvm/amd64/paging.cpp

@@ -572,8 +572,12 @@ char * writable_page_at(vMemory& memory, uint64_t addr, uint64_t verify_flags, b
 		auto* pdpt = memory.page_at(pdpt_mem);
 		/* Make copy of page if needed */
 		if (is_copy_on_write(pml4[i])) {
-			clone_and_update_entry(memory, pml4[i], pdpt, PDE64_RW);
-			CLPRINT("-> Cloning a PML4 entry %lu: 0x%lX at %p\n", i, pml4[i], pdpt);
+			if (memory.main_memory_writes) {
+				unlock_identity_mapped_entry(pml4[i]);
+			} else {
+				clone_and_update_entry(memory, pml4[i], pdpt, PDE64_RW);
+				CLPRINT("-> Cloning a PML4 entry %lu: 0x%lX at %p\n", i, pml4[i], pdpt);
+			}
 			assert(!is_copy_on_write(pml4[i]) && (pml4[i] & PDE64_PRESENT));
 		}
 		const uint64_t j = index_from_pdpt_entry(addr);
@@ -582,8 +586,12 @@ char * writable_page_at(vMemory& memory, uint64_t addr, uint64_t verify_flags, b
 			auto* pd = memory.page_at(pd_mem);
 			/* Make copy of page if needed */
 			if (is_copy_on_write(pdpt[j])) {
-				clone_and_update_entry(memory, pdpt[j], pd, PDE64_RW);
-				CLPRINT("-> Cloning a PDPT entry: 0x%lX\n", pdpt[j]);
+				if (memory.main_memory_writes) {
+					unlock_identity_mapped_entry(pdpt[j]);
+				} else {
+					clone_and_update_entry(memory, pdpt[j], pd, PDE64_RW);
+					CLPRINT("-> Cloning a PDPT entry: 0x%lX\n", pdpt[j]);
+				}
 			}
 			const uint64_t k = index_from_pd_entry(addr);
 			if (pd[k] & PDE64_PRESENT) {
@@ -594,7 +602,13 @@ char * writable_page_at(vMemory& memory, uint64_t addr, uint64_t verify_flags, b
 					/* Copy-on-write 2MB page */
 
 					/* NOTE: Make sure we are re-reading pd[k] */
-					if (memory.split_hugepages && (pd[k] & PDE64_PS)) { // 2MB page
+					if (memory.main_memory_writes) {
+						unlock_identity_mapped_entry(pd[k]);
+						if (pd[k] & PDE64_PS) {
+							memory.increment_unlocked_pages(512);
+						}
+						goto entry_is_no_longer_copy_on_write;
+					} else if (memory.split_hugepages && (pd[k] & PDE64_PS)) { // 2MB page
 						CLPRINT("-> Splitting a 2MB page, addr=0x%lX rw=%lu cloneable=%lu\n",
 							addr, pd[k] & PDE64_RW, pd[k] & PDE64_CLONEABLE);
 						/* Remove PS flag */
@@ -656,6 +670,7 @@ char * writable_page_at(vMemory& memory, uint64_t addr, uint64_t verify_flags, b
 					CLPRINT("-> Cloning a PD entry: 0x%lX\n", pd[k]);
 				}
 
+entry_is_no_longer_copy_on_write:
 				if (pd[k] & PDE64_PS) { // 2MB page
 					if (UNLIKELY((pd[k] & verify_flags) != verify_flags)) {
 						memory_exception("page_at: pt entry not user writable", addr, pd[k]);
@@ -673,6 +688,7 @@ char * writable_page_at(vMemory& memory, uint64_t addr, uint64_t verify_flags, b
 					if (is_copy_on_write(pt[e])) {
 						if (memory.is_forkable_master() && memory.main_memory_writes) {
 							unlock_identity_mapped_entry(pt[e]);
+							memory.increment_unlocked_pages(1);
 						} else if (write_zeroes || (pt[e] & PDE64_DIRTY) == 0x0) {
 							zero_and_update_entry(memory, pt[e], data, PDE64_RW);
 						} else {

lib/tinykvm/common.hpp

@@ -37,7 +37,7 @@ namespace tinykvm
 		uint32_t reset_free_work_mem = 0; /* reset_to() */
 		uint64_t vmem_base_address = 0;
 		std::string_view binary = {};
-		std::vector<VirtualRemapping> remappings;
+		std::vector<VirtualRemapping> remappings {};
 
 		bool verbose_loader = false;
 		bool short_lived = false;
@@ -46,7 +46,7 @@ namespace tinykvm
 		/* When enabled, master VMs will write directly
 		   to their own main memory instead of memory banks,
 		   allowing forks to immediately see changes. */
-		bool master_direct_memory_writes = true;
+		bool master_direct_memory_writes = false;
 		/* When enabled, split hugepages during page faults. */
 		bool split_hugepages = false;
 		/* When enabled, reset_to() will accept a different

lib/tinykvm/memory.cpp

@@ -1,6 +1,7 @@
 #include "machine.hpp"
 #include <cstring>
 #include <sys/mman.h>
+#include <stdexcept>
 #include <string>
 #include <unistd.h>
 #include <unordered_set>
@@ -285,11 +286,7 @@ size_t Machine::banked_memory_pages() const noexcept
 }
 size_t Machine::banked_memory_capacity_pages() const noexcept
 {
-	size_t count = 0;
-	for (const auto& bank : memory.banks) {
-		count += bank.n_pages;
-	}
-	return count;
+	return memory.banks.max_pages();
 }
 
 __attribute__((cold, noreturn))
@@ -298,6 +295,18 @@ void vMemory::memory_exception(const char* msg, uint64_t addr, uint64_t size)
 	throw MemoryException(msg, addr, size);
 }
 
+void vMemory::increment_unlocked_pages(size_t pages)
+{
+	if (this->main_memory_writes) {
+		this->unlocked_pages += pages;
+		if (this->unlocked_pages > this->banks.max_pages()) {
+			memory_exception("Out of working memory",
+				this->unlocked_pages * PAGE_SIZE, this->banks.max_pages() * PAGE_SIZE);
+		}
+	} else {
+		memory_exception("Memory::increment_unlocked_pages() without direct main memory writes enabled", 0, pages);
+	}
+}
 
 uint64_t vMemory::expectedUsermodeFlags() const noexcept
 {

lib/tinykvm/memory.hpp

@@ -22,6 +22,9 @@ struct vMemory {
 	/* Optional executable memory range */
 	uint64_t vmem_exec_begin = 0;
 	uint64_t vmem_exec_end   = 0;
+	/* Counter for the number of pages that have been unlocked
+	   in the main memory. */
+	size_t unlocked_pages = 0;
 	/* Linear memory */
 	char*  ptr;
 	size_t size;
@@ -61,6 +64,14 @@ struct vMemory {
 	MemoryBank::Page new_hugepage();
 
 	bool compare(const vMemory& other);
+	/* When a main VM has direct memory writes enabled, it can
+	   write directly to its own memory, but in order to constrain
+	   the memory usage, we need to keep track of the number of
+	   pages that have been unlocked. */
+	void increment_unlocked_pages(size_t pages);
+	size_t unlocked_memory_pages() const noexcept {
+		return unlocked_pages;
+	}
 
 	VirtualMem vmem() const;
 

lib/tinykvm/memory_bank.cpp

@@ -14,7 +14,7 @@ static constexpr bool MADVISE_NOT_DELETE  = true;
 
 MemoryBanks::MemoryBanks(Machine& machine, const MachineOptions& options)
 	: m_machine { machine },
-	  m_arena_begin { 0x7000000000 },
+	  m_arena_begin { ARENA_BASE_ADDRESS },
 	  m_arena_next { m_arena_begin },
 	  m_idx_begin { FIRST_BANK_IDX },
 	  m_idx { m_idx_begin },
@@ -27,10 +27,8 @@ void MemoryBanks::set_max_pages(size_t new_max)
 	this->m_max_pages = new_max;
 	//fprintf(stderr, "max_pages: %zu/%zu\n", m_mem.size(), new_max);
 	/* Reserve the maximum number of banks possible.
-	   We have to + 1 to make sure it's rounded up, avoiding
-	   any possible reallocations close to being out of memory.
 	   NOTE: DO NOT modify this! Needs deque behavior. */
-	m_mem.reserve(m_max_pages / MemoryBank::N_PAGES + 1);
+	m_mem.reserve((m_max_pages + MemoryBank::N_PAGES-1) / MemoryBank::N_PAGES);
 }
 
 char* MemoryBanks::try_alloc(size_t N)
@@ -81,8 +79,8 @@ MemoryBank& MemoryBanks::get_available_bank(size_t pages)
 	/* Allocate new memory bank if we are not maxing out memory */
 	if (m_num_pages < m_max_pages) {
 		if constexpr (VERBOSE_MEMORY_BANK) {
-			printf("Allocating new bank at 0x%lX with total pages %u\n",
-				m_arena_next, m_num_pages + MemoryBank::N_PAGES);
+			printf("Allocating new bank at 0x%lX with total pages %u/%u\n",
+				m_arena_next, m_num_pages + MemoryBank::N_PAGES, m_max_pages);
 		}
 		auto& bank = this->allocate_new_bank(m_arena_next);
 		m_num_pages += bank.n_pages;

lib/tinykvm/memory_bank.hpp

@@ -50,12 +50,14 @@ struct MemoryBank {
 
 struct MemoryBanks {
 	static constexpr unsigned FIRST_BANK_IDX = 2;
+	static constexpr uint64_t ARENA_BASE_ADDRESS = 0x7000000000;
 
 	MemoryBanks(Machine&, const MachineOptions&);
 
 	MemoryBank& get_available_bank(size_t n_pages);
 	void reset(const MachineOptions&);
 	void set_max_pages(size_t new_max);
+	size_t max_pages() const noexcept { return m_max_pages; }
 
 	bool using_hugepages() const noexcept { return m_using_hugepages; }
 

lib/tinykvm/vcpu.cpp

@@ -377,14 +377,12 @@ 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)
 {
-	assert(this->m_prepped == false);
 	this->m_prepped = true;
 	/* Make each writable page read-only, causing page fault.
 	   any page after the @shared_memory_boundary is untouched,
 	   effectively turning it into a shared memory area for all. */
 	if (shared_memory_boundary == 0)
 		shared_memory_boundary = UINT64_MAX;
-	foreach_page_makecow(this->memory, kernel_end_address(), shared_memory_boundary);
 
 	// Visualizing the page tables after makecow should show that all
 	// relevant user-writable pages have been made read-only and cloneable
@@ -395,10 +393,27 @@ void Machine::prepare_copy_on_write(size_t max_work_mem, uint64_t shared_memory_
 	memory.banks.set_max_pages(max_work_mem / PAGE_SIZE);
 	/* Without working memory we will not be able to make
 	   this master VM usable after prepare_copy_on_write. */
-	if (max_work_mem == 0)
+	if (max_work_mem == 0) {
+		/* If there are previously banked pages, we need to
+		   flatten them into the main memory. */
+		/// XXX: Implement memory flattening
+		memory.page_tables = memory.physbase + PT_ADDR;
+		struct kvm_sregs sregs = this->get_special_registers();
+
+		/* Page table entry will be cloned at the start */
+		sregs.cr3 = memory.page_tables;
+		sregs.cr0 |= CR0_WP;
+
+		vcpu.set_special_registers(sregs);
+		this->enter_usermode();
+
+		foreach_page_makecow(this->memory, kernel_end_address(), shared_memory_boundary);
 		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);
 	this->setup_cow_mode(this);
 }
 void Machine::setup_cow_mode(const Machine* other)

tests/unit/fork.cpp

@@ -42,7 +42,7 @@ extern void prints_hello_world() {
 
 	// Make machine forkable (no working memory)
 	machine.prepare_copy_on_write(65536);
-	REQUIRE(machine.banked_memory_pages() == 4);
+	REQUIRE(machine.banked_memory_pages() == 5);
 	REQUIRE(machine.is_forkable());
 	REQUIRE(!machine.is_forked());
 
@@ -249,7 +249,7 @@ extern int get_value() {
 
 	// Make machine forkable (with working memory)
 	machine.prepare_copy_on_write(65536);
-	REQUIRE(machine.banked_memory_pages() == 4);
+	REQUIRE(machine.banked_memory_pages() == 5);
 	REQUIRE(machine.is_forkable());
 	REQUIRE(!machine.is_forked());
 
@@ -293,10 +293,10 @@ extern int get_value() {
 
 	// Value now starts at 1 due to the change in main VM
 	fork1.timed_vmcall(funcaddr, 4.0f);
-	REQUIRE(fork1.return_value() == 2);
+	REQUIRE(fork1.return_value() == 1);
 
 	fork2.timed_vmcall(funcaddr, 4.0f);
-	REQUIRE(fork2.return_value() == 2);
+	REQUIRE(fork2.return_value() == 1);
 }
 
 TEST_CASE("Fork sanity checks w/crashes", "[Fork]")
@@ -363,7 +363,9 @@ extern void crash() {
 TEST_CASE("Fork and run main()", "[Fork]")
 {
 	const auto binary = build_and_load(R"M(
+#include <stdio.h>
 int main() {
+	printf("Hello World!\n");
 	return 666;
 }
 static unsigned value = 12345;
@@ -378,17 +380,20 @@ int func2() {
 }
 )M");
 
-	tinykvm::Machine machine { binary, { .max_mem = MAX_MEMORY } };
+	tinykvm::Machine machine { binary, {
+		.max_mem = MAX_MEMORY,
+		.master_direct_memory_writes = true
+	} };
 
 	// We need to create a Linux environment for runtimes to work well
 	machine.setup_linux({"fork"}, env);
 	REQUIRE(machine.banked_memory_pages() == 0);
 
-	// Make machine forkable (with *NO* working memory)
+	// Make machine forkable (with 4MB working memory)
 	machine.prepare_copy_on_write(4ULL << 20);
+	REQUIRE(machine.banked_memory_capacity_bytes() == 4ULL << 20);
 	REQUIRE(machine.is_forkable());
 	REQUIRE(!machine.is_forked());
-	REQUIRE(machine.return_value() == 0); // Initial register value
 
 	// Run for at most 4 seconds before giving up
 	machine.run(4.0f);
@@ -397,11 +402,10 @@ int func2() {
 	// We only gave it 4MB working memory, so lets mmap allocate that and verify
 	// that if we write more than that, we get an exception thrown
 	REQUIRE_THROWS([&] () {
-		const size_t size = 5ULL << 20;
-		uint64_t addr = machine.mmap_allocate(5ULL << 20);
+		const size_t size = 8ULL << 20;
+		uint64_t addr = machine.mmap_allocate(size);
 		char buffer[4096];
-		for (int i = 0; i < 4096; i++)
-			buffer[i] = 'a';
+		memset(buffer, 'a', sizeof(buffer));
 		for (size_t i = 0; i < size; i += 4096)
 		{
 			machine.copy_to_guest(addr + i, buffer, 4096);
@@ -411,11 +415,33 @@ int func2() {
 	}());
 
 	// There are banked pages now
-	const auto banked_pages_before = machine.banked_memory_pages();
+	const auto banked_pages_before = machine.main_memory().unlocked_memory_pages();
 	REQUIRE(banked_pages_before > 500);
 
+	// We have no free memory now, so make another VM
+	tinykvm::Machine machine2 { binary, {
+		.max_mem = MAX_MEMORY,
+		.master_direct_memory_writes = true
+	} };
+
+	// We need to create a Linux environment for runtimes to work well
+	machine2.setup_linux({"fork"}, env);
+	REQUIRE(machine2.banked_memory_pages() == 0);
+
+	// Make machine forkable (with 4MB working memory)
+	machine2.prepare_copy_on_write(4ULL << 20);
+	REQUIRE(machine2.banked_memory_capacity_bytes() == 4ULL << 20);
+	REQUIRE(machine2.is_forkable());
+	REQUIRE(!machine2.is_forked());
+
+	// Run for at most 4 seconds before giving up
+	machine2.run(4.0f);
+	REQUIRE(machine2.return_value() == 666); // Main() return value
+
+	machine2.prepare_copy_on_write(0);
+
 	// Create fork
-	auto fork1 = tinykvm::Machine { machine, {
+	auto fork1 = tinykvm::Machine { machine2, {
 		.max_mem = MAX_MEMORY, .max_cow_mem = MAX_COWMEM
 	} };
 	REQUIRE(fork1.return_value() == 666); // Main() return value
@@ -425,25 +451,17 @@ int func2() {
 
 	fork1.vmcall("func2");
 	REQUIRE(fork1.return_value() == 54321);
-
-	// This is problematic, but we will try to fix this later
-	// Forked VM is supposed to diverge from the main VM, regardless of mode
-	machine.vmcall("set_value", 99999);
-	fork1.vmcall("func1");
-	REQUIRE(fork1.return_value() == 99999);
-
-	REQUIRE(machine.banked_memory_pages() == banked_pages_before);
 	REQUIRE(fork1.banked_memory_pages() > 0);
 
 	for (int i = 0; i < 20; i++)
 	{
-		fork1.reset_to(machine, {
+		fork1.reset_to(machine2, {
 			.max_mem = MAX_MEMORY,
 			.max_cow_mem = MAX_COWMEM,
 		});
 
 		fork1.vmcall("func1");
-		REQUIRE(fork1.return_value() == 99999);
+		REQUIRE(fork1.return_value() == 12345);
 
 		fork1.vmcall("func2");
 		REQUIRE(fork1.return_value() == 54321);