Track BRK in order to prevent lowering heap

Lowering the heap caused ghosts in the machine and a lot of confusion. Now fixed.

Author: fwsGonzo

lib/tinykvm/amd64/paging.hpp

@@ -28,4 +28,8 @@ static inline bool page_is_zeroed(const uint64_t* page) {
 	return true;
 }
 
+static constexpr inline uint64_t PageMask() {
+	return vMemory::PageSize() - 1UL;
+}
+
 } // tinykvm

lib/tinykvm/linux/system_calls.cpp

@@ -244,7 +244,7 @@ void Machine::setup_linux_system_calls()
 				} else {
 					regs.rax = dst;
 				}
-				PRINTMMAP("mmap(0x%llX, %llu, prot=%llX, flags=%llX) = 0x%llX\n",
+				PRINTMMAP("mmap(0x%lX, %lu, prot=%llX, flags=%llX) = 0x%llX\n",
 						  address, length, regs.rdx, regs.r10, regs.rax);
 				cpu.set_registers(regs);
 				return;
@@ -300,7 +300,7 @@ void Machine::setup_linux_system_calls()
 			{
 				cpu.machine().memzero(regs.rax, length);
 			}
-			PRINTMMAP("mmap(0x%llX, %llu, prot=%llX, flags=%llX) = 0x%llX\n",
+			PRINTMMAP("mmap(0x%lX, %lu, prot=%llX, flags=%llX) = 0x%llX\n",
 					  address, length, regs.rdx, regs.r10, regs.rax);
 			cpu.set_registers(regs);
 		});
@@ -331,17 +331,16 @@ void Machine::setup_linux_system_calls()
 	Machine::install_syscall_handler(
 		SYS_brk, [](vCPU& cpu) { // BRK
 			auto& regs = cpu.registers();
-			if (regs.rdi > cpu.machine().heap_address() + Machine::BRK_MAX)
-			{
-				regs.rax = cpu.machine().heap_address() + Machine::BRK_MAX;
-			}
-			else if (regs.rdi < cpu.machine().heap_address())
-			{
-				regs.rax = cpu.machine().heap_address();
-			}
-			else
-			{
-				regs.rax = regs.rdi;
+			const uint64_t old_brk = cpu.machine().brk_address();
+			uint64_t new_brk = regs.rdi;
+			if (new_brk < old_brk) {
+				// brk() to a lower address, keep the old one
+				// We can only grow the heap, not shrink it.
+				regs.rax = old_brk;
+			} else {
+				// clamp brk() outside to the heap range
+				new_brk = std::min(new_brk, cpu.machine().brk_end_address());
+				regs.rax = new_brk;
 			}
 			SYSPRINT("brk(0x%llX) = 0x%llX\n", regs.rdi, regs.rax);
 			cpu.set_registers(regs);
@@ -1049,7 +1048,7 @@ void Machine::setup_linux_system_calls()
 	Machine::install_syscall_handler(
 		SYS_getrlimit, [](vCPU& cpu) { // getrlimit
 			auto& regs = cpu.registers();
-			const auto g_rlim = regs.rsi;
+			[[maybe_unused]] const auto g_rlim = regs.rsi;
 			regs.rax = -ENOSYS;
 			SYSPRINT("getrlimit(0x%llX) = %lld\n", g_rlim, regs.rax);
 			cpu.set_registers(regs);

lib/tinykvm/machine.cpp

@@ -65,6 +65,7 @@ Machine::Machine(const Machine& other, const MachineOptions& options)
 	  m_image_base    {other.m_image_base},
 	  m_stack_address {other.m_stack_address},
 	  m_heap_address  {other.m_heap_address},
+	  m_brk_address   {other.m_brk_address},
 	  m_start_address {other.m_start_address},
 	  m_kernel_end    {other.m_kernel_end},
 	  m_mm            {other.m_mm},
@@ -158,6 +159,7 @@ void Machine::reset_to(const Machine& other, const MachineOptions& options)
 		this->m_image_base    = other.m_image_base;
 		this->m_stack_address = other.m_stack_address;
 		this->m_heap_address  = other.m_heap_address;
+		this->m_brk_address   = other.m_brk_address;
 		this->m_start_address = other.m_start_address;
 		this->m_kernel_end    = other.m_kernel_end;
 		memory.fork_reset(other.memory, options);

lib/tinykvm/machine.hpp

@@ -158,7 +158,10 @@ struct Machine
 	address_t kernel_end_address() const noexcept { return m_kernel_end; }
 	address_t max_address() const noexcept { return memory.physbase + memory.size; }
 
-	static constexpr uint64_t BRK_MAX = 0x20000;
+	static constexpr uint64_t BRK_MAX = 0x22000;
+	address_t brk_address() const noexcept { return this->m_brk_address; }
+	address_t brk_end_address() const noexcept { return this->m_heap_address + BRK_MAX; }
+	void set_brk_address(address_t addr) { this->m_brk_address = addr; }
 	address_t mmap_start() const noexcept { return this->m_heap_address + BRK_MAX; }
 	address_t mmap_allocate(size_t bytes);
 	bool      mmap_unmap(uint64_t addr, size_t size);
@@ -271,6 +274,7 @@ struct Machine
 	address_t m_image_base = 0x0;
 	address_t m_stack_address;
 	address_t m_heap_address;
+	address_t m_brk_address;
 	address_t m_start_address;
 	address_t m_kernel_end;
 

lib/tinykvm/machine_elf.cpp

@@ -5,6 +5,7 @@
 #include <stdexcept>
 #ifdef TINYKVM_ARCH_AMD64
 #include "amd64/idt.hpp" // interrupt_header()
+#include "amd64/paging.hpp"
 #endif
 #include "util/elf.hpp"
 
@@ -134,18 +135,19 @@ void Machine::elf_loader(std::string_view binary, const MachineOptions& options)
 				break;
 		}
 
-		uint64_t endm = hdr->p_vaddr + hdr->p_memsz;
-		endm += vMemory::PageSize()-1; endm &= ~(vMemory::PageSize()-1);
+		const uint64_t endm = hdr->p_vaddr + hdr->p_memsz;
 		if (this->m_heap_address < endm)
 			this->m_heap_address = endm;
 	}
 
 	/* Make sure mmap starts at a sane offset */
 	this->m_heap_address += this->m_image_base;
+	this->m_heap_address = (this->m_heap_address + PageMask()) & ~PageMask();
+	this->m_brk_address = this->m_heap_address;
 	this->m_mm = this->mmap_start();
 
 	/* If there is not enough room for stack, move it */
-	const uint32_t STACK_SIZE = options.stack_size & ~0xFFFLL;
+	const uint32_t STACK_SIZE = (options.stack_size + PageMask()) & ~PageMask();
 	if (this->m_stack_address < STACK_SIZE) {
 		this->m_stack_address = this->mmap_allocate(STACK_SIZE) + STACK_SIZE;
 	}
@@ -157,7 +159,8 @@ void Machine::elf_loader(std::string_view binary, const MachineOptions& options)
 
 	if (options.verbose_loader) {
 	printf("* Entry is at %p\n", (void*) m_start_address);
-	printf("* Stack is at %p\n", (void*) m_stack_address);
+	printf("* Stack is at %p -> %p\n", (void*) (m_stack_address - STACK_SIZE),
+		(void*) (m_stack_address));
 	}
 }
 

lib/tinykvm/machine_utils.cpp

@@ -4,9 +4,6 @@
 #include "amd64/paging.hpp"
 
 namespace tinykvm {
-static constexpr uint64_t PageMask() {
-	return vMemory::PageSize() - 1UL;
-}
 
 void Machine::memzero(address_t addr, size_t len)
 {

lib/tinykvm/memory.cpp

@@ -16,9 +16,6 @@
 
 namespace tinykvm {
 #define USERMODE_FLAGS (0x7 | 1UL << 63) /* USER, READ/WRITE, PRESENT, NX */
-static constexpr uint64_t PageMask() {
-	return vMemory::PageSize() - 1UL;
-}
 
 vMemory::vMemory(Machine& m, const MachineOptions& options,
 	uint64_t ph, uint64_t sf, char* p, size_t s, bool own)