From 5befb8b4e67e1d06ba9729031dad7bda23a762c7 Mon Sep 17 00:00:00 2001
From: Freya Murphy <freya@freyacat.org>
Date: Fri, 4 Oct 2024 22:09:42 -0400
Subject: [PATCH] fix msim to load non page aligned segments

---
 msim/load.c | 156 +++++++++++++++++++++++++++++-----------------------
 msim/sim.h  |   1 -
 2 files changed, 88 insertions(+), 69 deletions(-)

diff --git a/msim/load.c b/msim/load.c
index a4029cf..e768740 100644
--- a/msim/load.c
+++ b/msim/load.c
@@ -9,6 +9,9 @@
 #include "sim.h"
 
 #define SEC_ALIGN 0x1000
+#define PAGE_SIZE 4096
+
+#define BITFILED_LEN (UINT32_MAX / PAGE_SIZE / 8)
 
 struct load_state {
 	FILE *file;
@@ -106,95 +109,112 @@ static int load_phdr(struct load_state *state)
 	return M_SUCCESS;
 }
 
-static int load_segment(struct simulator *sim,
-			struct load_state *state, Elf32_Phdr *hdr)
+static void set_page(uint8_t *bitfield, uint32_t addr)
 {
-	uint32_t off = B32(hdr->p_offset);
-	uint32_t len = B32(hdr->p_filesz);
-	uint32_t is_text = B32(hdr->p_flags) & PF_X;
+	int idx = (addr / PAGE_SIZE) / 8;
+	int off = (addr / PAGE_SIZE) % 8;
 
+	bitfield[idx] |= 1 << off;
+}
+
+static int get_page(const uint8_t *const bitfield, uint32_t addr)
+{
+	int idx = (addr / PAGE_SIZE) / 8;
+	int off = (addr / PAGE_SIZE) % 8;
+
+	return ((bitfield[idx] >> off) & 1);
+}
+
+static int load_segment(struct simulator *sim, struct load_state *state,
+			Elf32_Phdr *hdr, uint8_t* bitfield)
+{
+	uint32_t addr = B32(hdr->p_vaddr),
+	         off = B32(hdr->p_offset),
+	         len = B32(hdr->p_filesz),
+	         flags = B32(hdr->p_flags);
+
+	bool exec = flags & PF_X;
+
+	// ignore if empty
+	if (len < 1)
+		return M_SUCCESS;
+
+	// make sure segment is acutally inside
+	// the file
 	if (BOUND_CHK(state->file_sz, len, off)) {
 		ERROR("segment location invalid");
 		return M_ERROR;
 	}
 
-	uintptr_t addr = 0;
-	uint32_t add = 0;
-
-	if (len % SEC_ALIGN) {
-		add = SEC_ALIGN - (len % SEC_ALIGN);
-		len += add;
-	}
-
-	if (is_text) {
-		addr = sim->text_max;
-		sim->text_max += len;
-	} else {
-		addr = sim->data_max;
-		sim->data_max += len;
-	}
-
-	bool read = B32(hdr->p_flags) & PF_R;
-	bool write = B32(hdr->p_flags) & PF_W;
-
-	uint32_t prot = 0;
-	if (read)
-		prot |= PROT_READ;
-	if (write)
-		prot |= PROT_WRITE;
-
-	void *res = mmap((void*)addr, len, prot, MAP_PRIVATE | MAP_FIXED,
-		  state->fd, off);
-
-	if ((uintptr_t)res != addr) {
-		PERROR("failed to map executable");
+	// make sure the vitural address is also
+	// valid
+	if (BOUND_CHK(UINT32_MAX, len, addr)) {
+		ERROR("segment vitural addr invalid");
 		return M_ERROR;
 	}
 
+	// update text seg bounds
+	if (exec) {
+		if (addr < sim->text_min)
+			sim->text_min = addr;
+		if (addr + len > sim->text_max)
+			sim->text_max = addr + len;
+	}
+
+	// align the mapping ptr to
+	// the page size
+	uintptr_t ptr = (addr / PAGE_SIZE) * PAGE_SIZE;
+
+	// map each page that the segment
+	// requires
+	for (; ptr < addr + len; ptr += PAGE_SIZE) {
+		// dont remap if address is
+		// already mapped
+		if (get_page(bitfield, ptr))
+			continue;
+
+		// set page as mapped
+		set_page(bitfield, ptr);
+
+		void *res = mmap(
+			(void *) ptr, PAGE_SIZE,
+			PROT_READ | PROT_WRITE,
+			MAP_PRIVATE | MAP_FIXED | MAP_ANONYMOUS, -1, 0);
+
+		if ((uintptr_t) res != ptr) {
+			PERROR("failed to map executable");
+			return M_ERROR;
+		}
+
+	}
+
+	// load the segment into the mapped memory
+	fseek(state->file, off, SEEK_SET);
+	fread((void *) (uintptr_t) addr, 1, len, state->file);
+
 	return M_SUCCESS;
 }
 
 static int load_memory(struct simulator *sim, struct load_state *state)
 {
-	uint32_t base = 0;
+	// map each page in a 32bit address space to a single bit
+	// in the bitfield
+	uint8_t *bitfield = malloc(BITFILED_LEN);
+	if (bitfield == NULL) {
+		PERROR("cannot alloc");
+		return M_ERROR;
+	}
+	memset(bitfield, 0, BITFILED_LEN);
 
 	for (uint32_t i = 0; i < state->phdr_len; i++) {
-		Elf32_Phdr *hdr = NULL;
-		uint32_t min = UINT32_MAX;
+		Elf32_Phdr *hdr = &state->phdr[i];
 
-		if (B32(state->phdr[i].p_filesz) < 1)
-			continue;
-
-		// we need to load segments in order
-		for (uint32_t j = 0; j < state->phdr_len; j++) {
-			Elf32_Phdr *temp = &state->phdr[j];
-			uint32_t off = B32(temp->p_offset);
-			uint32_t len = B32(temp->p_filesz);
-
-			if (len < 1)
-				continue;
-
-			if (off <= base)
-				continue;
-
-			if (off >= min)
-				continue;
-
-			min = off;
-			hdr = temp;
-		}
-
-		base = min;
-
-		if (hdr == NULL) {
-			ERROR("invalid elf phdr");
-			return M_ERROR;
-		}
-
-		if (load_segment(sim, state, hdr))
+		if (load_segment(sim, state, hdr, bitfield))
 			return M_ERROR;
 	}
 
+	free(bitfield);
+
 	return M_SUCCESS;
 }
 
diff --git a/msim/sim.h b/msim/sim.h
index 6913792..e0ac4bb 100644
--- a/msim/sim.h
+++ b/msim/sim.h
@@ -34,7 +34,6 @@ struct simulator_args {
 struct simulator {
 	struct simulator_args *args;
 
-
 	/// the registers
 	uint32_t reg[32];