Fix issues with fork initial state and dirty bits

Author: fwsGonzo

lib/tinykvm/amd64/paging.cpp

@@ -305,7 +305,7 @@ uint64_t setup_amd64_paging(vMemory& memory,
 							const uint64_t addr2m = (base_giga_page << 30) | (pdidx << 21);
 							ptentry = PDE64_PRESENT | PDE64_USER | PDE64_NX | PDE64_PS | addr2m;
 							if (!read) ptentry &= ~PDE64_PRESENT; // A weird one, but... AMD64.
-							if (write) ptentry |= PDE64_RW;
+							if (write) ptentry |= PDE64_RW | PDE64_DIRTY;
 							else ptentry |= PDE64_G; // Global bit for read-only pages
 							if (exec) ptentry &= ~PDE64_NX;
 							// Increment whole 2MB page
@@ -349,6 +349,8 @@ uint64_t setup_amd64_paging(vMemory& memory,
 					if (!write) {
 						ptentry &= ~PDE64_RW;
 						ptentry |= PDE64_G; // Global bit for read-only pages
+					} else {
+						ptentry |= PDE64_DIRTY;
 					}
 					addr += 0x1000;
 				}

lib/tinykvm/common.hpp

@@ -106,7 +106,7 @@ namespace tinykvm
 		   to the given file. Depending on `snapshot_mode`,
 		   the file may be created if it does not exist,
 		   and must be of the correct size if it does exist. */
-		std::string snapshot_file;
+		std::string snapshot_file {};
 		enum SnapshotMode {
 			Disabled = 0,
 			Open = 1,

lib/tinykvm/machine_utils.cpp

@@ -41,34 +41,27 @@ void Machine::memzero(address_t addr, size_t len)
 
 void Machine::copy_to_guest(address_t addr, const void* vsrc, size_t len, bool zeroes)
 {
-	if (uses_cow_memory() || !memory.safely_within(addr, len))
+	auto* src = (const uint8_t *)vsrc;
+	while (len != 0)
 	{
-		auto* src = (const uint8_t *)vsrc;
-		while (len != 0)
-		{
-			const size_t offset = addr & PageMask();
-			const size_t size = std::min(vMemory::PageSize() - offset, len);
-			const bool full_page = (size == vMemory::PageSize());
-			WritablePageOptions opts;
-			opts.allow_dirty = full_page;
-			opts.zeroes = zeroes;
-			// Get a writable page, possibly allocating a new one
-			WritablePage page = writable_page_at(memory, addr & ~PageMask(), memory.expectedUsermodeFlags(), opts);
-			// Page is always dirty
-			page.set_dirty();
-			// Copy data to the page
-			char* page_data = page.page;
-			std::memcpy(&page_data[offset], src, size);
+		const size_t offset = addr & PageMask();
+		const size_t size = std::min(vMemory::PageSize() - offset, len);
+		const bool full_page = (size == vMemory::PageSize());
+		WritablePageOptions opts;
+		opts.allow_dirty = full_page;
+		opts.zeroes = zeroes;
+		// Get a writable page, possibly allocating a new one
+		WritablePage page = writable_page_at(memory, addr & ~PageMask(), memory.expectedUsermodeFlags(), opts);
+		// Page is always dirty
+		page.set_dirty();
+		// Copy data to the page
+		char* page_data = page.page;
+		std::memcpy(&page_data[offset], src, size);
 
-			addr += size;
-			src += size;
-			len -= size;
-		}
-		return;
+		addr += size;
+		src += size;
+		len -= size;
 	}
-	/* Original VM uses identity-mapped memory */
-	auto* dst = memory.safely_at(addr, len);
-	std::memcpy(dst, vsrc, len);
 }
 
 void Machine::copy_from_guest(void* vdst, address_t addr, size_t len) const

lib/tinykvm/memory_bank.cpp

@@ -156,7 +156,7 @@ void MemoryBanks::reset(const MachineOptions& options)
 	//size_t limit_pages = options.reset_free_work_mem / vMemory::PageSize();
 
 	/* Instead of removing the banks, give memory back to kernel */
-	for (size_t i = 1u; i < m_mem.size(); i++) {
+	for (size_t i = 0u; i < m_mem.size(); i++) {
 		/* WARNING: MADV_FREE *does not* immediately free, so use MADV_DONTNEED instead. */
 		if (m_mem[i].dirty_size() > 0)
 			madvise(m_mem[i].mem, m_mem[i].dirty_size(), MADV_DONTNEED);

tests/unit/fork.cpp

@@ -2,9 +2,10 @@
 #include <catch2/matchers/catch_matchers_string.hpp>
 
 #include <tinykvm/machine.hpp>
+#include <linux/kvm.h>
 extern std::vector<uint8_t> build_and_load(const std::string& code);
 static const uint64_t MAX_MEMORY = 8ul << 20; /* 8MB */
-static const uint64_t MAX_COWMEM = 3ul << 20; /* 1MB */
+static const uint64_t MAX_COWMEM = 3ul << 20; /* 3MB */
 static const std::vector<std::string> env {
 	"LC_TYPE=C", "LC_ALL=C", "USER=root"
 };
@@ -566,3 +567,74 @@ int func2() {
 		REQUIRE(fork2.return_value() == 22222);
 	}
 }
+
+TEST_CASE("Fork before main()", "[Fork]")
+{
+	const auto binary = build_and_load(R"M(
+#include <stdio.h>
+extern void _exit(int);
+int main() {
+	printf("Hello World!\n");
+	_exit(666);
+	return 666;
+}
+static unsigned value = 12345;
+void set_value(int v) {
+	value = v;
+}
+int func1() {
+	return value;
+}
+int func2() {
+	return 54321;
+}
+)M");
+
+	tinykvm::Machine machine1 { binary, {
+		.max_mem = 10ull << 20, // We need 10mb because of fragmentation
+	} };
+	machine1.setup_linux({"fork"}, env);
+	machine1.prepare_copy_on_write();
+	REQUIRE(machine1.is_forkable());
+	REQUIRE(!machine1.is_forked());
+
+	tinykvm::Machine machine2 { binary, {
+		.max_mem = 10ull << 20,
+	} };
+	machine2.prepare_copy_on_write(); // No Linux setup
+	REQUIRE(machine2.is_forkable());
+	REQUIRE(!machine2.is_forked());
+
+	auto fork1 = tinykvm::Machine { machine1, {
+		.max_cow_mem = MAX_COWMEM,
+		.split_hugepages = true
+	} };
+
+	for (int i = 0; i < 100; i++)
+	{
+		fork1.run(4.0f);
+		REQUIRE(fork1.return_value() == 666);
+
+		fork1.reset_to(machine1, {
+			.max_cow_mem = 4ul << 20,
+			.split_hugepages = true
+		});
+	}
+
+	auto fork2 = tinykvm::Machine { machine2, {
+		.max_cow_mem = MAX_COWMEM,
+		.split_hugepages = true
+	} };
+
+	for (int i = 0; i < 100; i++)
+	{
+		fork2.setup_linux({"fork"}, env);
+		fork2.run(4.0f);
+		REQUIRE(fork2.return_value() == 666);
+
+		fork2.reset_to(machine2, {
+			.max_cow_mem = 4ul << 20,
+			.split_hugepages = true
+		});
+	}
+}