Fix the physical memory map overwriting modules

It turns out GRUB places modules after the kernel in memory. I did not
anticipate this, and thus I was using the memory right after the kernel
to store the bitmap used in the physical memory manager.

A strange bug led me to this fix:
 When using the new memory allocator and memory mapping the first 16 MiB
 of memory, the letter 'w' in the "loop.c" program got scrambled. Just
 that one letter. And no bug whatsoever if I identity mapped only 4 MiB.

Mapping pages allocates physical memory for new pages tables, so my
guess is that that was it.

kernel/include/kernel/paging.h

@@ -37,9 +37,6 @@ uintptr_t paging_virt_to_phys(uintptr_t virt);
 #define KERNEL_HEAP_BEGIN (KERNEL_BASE_VIRT + 0x1000*1024)
 #define KERNEL_HEAP_SIZE 0x1E00000
 
-#define PHYS_TO_VIRT(addr) ((addr) + KERNEL_BASE_VIRT)
-#define VIRT_TO_PHYS(addr) ((addr) - KERNEL_BASE_VIRT)
-
 #define PAGE_PRESENT 1
 #define PAGE_RW      2
 #define PAGE_USER    4

kernel/include/kernel/pmm.h

@@ -14,6 +14,7 @@ uintptr_t pmm_alloc_aligned_large_page();
 uintptr_t pmm_alloc_pages(uint32_t num);
 void pmm_free_page(uintptr_t addr);
 void pmm_free_pages(uintptr_t addr, uint32_t num);
+uintptr_t pmm_get_kernel_end();
 
 extern uint32_t* mem_map;
 

kernel/include/kernel/sys.h

@@ -10,6 +10,12 @@
 /* Returns the next multiple of `s` greater than `a`, or `a` if it is a
  * multiple of `s`.
  */
-#define ALIGN(a, s) (((a)/(s) + ((a) % (s) ? 1 : 0)) * (s))
+static uint32_t align_to(uint32_t n, uint32_t align) {
+    if (n % align == 0) {
+        return n;
+    }
+
+    return n + (align - n % align);
+}
 
 #endif

kernel/src/kernel.c

@@ -50,6 +50,7 @@ void kernel_main(multiboot_t* boot, uint32_t magic) {
 	init_timer();
 	init_wm();
 
+	// Load GRUB modules as programs
 	mod_t* modules = (mod_t*) boot->mods_addr;
 
 	for (uint32_t i = 0; i < boot->mods_count; i++) {
@@ -62,7 +63,5 @@ void kernel_main(multiboot_t* boot, uint32_t magic) {
 
 	proc_print_processes();
 
-
-
 	init_proc();
 }

kernel/src/mem/mem.c

@@ -14,6 +14,8 @@ void mem_print_blocks();
 static mem_block_t* bottom = NULL;
 static mem_block_t* top = NULL;
 
+/* Returns the size of a block, including the header.
+ */
 uint32_t mem_block_size(mem_block_t* block) {
 	return sizeof(mem_block_t) + (block->size & ~1);
 }
@@ -34,7 +36,7 @@ mem_block_t* mem_new_block(uint32_t size, uint32_t align) {
 
 	// We start the heap right where the first allocation works
 	if (!top) {
-		uintptr_t addr = ALIGN(KERNEL_HEAP_BEGIN+header_size, align) - header_size;
+		uintptr_t addr = align_to(KERNEL_HEAP_BEGIN+header_size, align) - header_size;
 		bottom = (mem_block_t*) addr;
 		top = bottom;
 		top->size = size | 1;
@@ -45,15 +47,15 @@ mem_block_t* mem_new_block(uint32_t size, uint32_t align) {
 
 	// I did the math and we always have next_aligned >= next.
 	uintptr_t next = (uintptr_t) top + mem_block_size(top);
-	uintptr_t next_aligned = ALIGN(next+header_size, align) - header_size;
+	uintptr_t next_aligned = align_to(next+header_size, align) - header_size;
 
 	mem_block_t* block = (mem_block_t*) next_aligned;
 	block->size = size | 1;
 	block->next = NULL;
 
 	// Insert a free block between top and our aligned block, if there's enough
 	// space. That block is 8-bytes aligned.
-	next = ALIGN(next+header_size, MIN_ALIGN) - header_size;
+	next = align_to(next+header_size, MIN_ALIGN) - header_size;
 	if (next_aligned - next > sizeof(mem_block_t) + MIN_ALIGN) {
 		mem_block_t* filler = (mem_block_t*) next;
 		filler->size = next_aligned - next - sizeof(mem_block_t);
@@ -62,7 +64,6 @@ mem_block_t* mem_new_block(uint32_t size, uint32_t align) {
 		printf("adding filler block (%p, %d)\n", filler->data, filler->size);
 #endif
 
-		filler->next = block;
 		top->next = filler;
 		top = filler;
 	}
@@ -132,7 +133,7 @@ void* kamalloc(uint32_t size, uint32_t align) {
 #ifdef MEM_DEBUG
 	printf("\nkamalloc(0x%X, %d)\n", size, align);
 #endif
-	size = ALIGN(size, 8);
+	size = align_to(size, 8);
 
 	mem_block_t* block = mem_find_block(size, align);
 

kernel/src/mem/paging.c

@@ -9,6 +9,8 @@
 #include <stdlib.h>
 #include <string.h>
 
+#define VIRT_TO_PHYS(x) ((x) - KERNEL_BASE_VIRT)
+
 #define DIRECTORY_INDEX(x) ((x) >> 22)
 #define TABLE_INDEX(x) (((x) >> 12) & 0x3FF)
 
@@ -26,11 +28,11 @@ void init_paging() {
 	kernel_directory[1023] = dir_phys | PAGE_PRESENT | PAGE_RW;
 	paging_invalidate_page(0xFFC00000);
 
-	// Identity map only the first 2 MiB: in boot.S, we identity mapped the first 4 MiB
-	// We take 2 MiB to give room to GRUB modules.
+	// Identity the `0x00-kernel_end` range
+	uint32_t kernel_pages = pmm_get_kernel_end()/0x1000;
 	kernel_directory[0] = 0;
+	paging_map_pages(0x00000000, 0x00000000, kernel_pages, PAGE_RW);
 	paging_invalidate_page(0x00000000);
-	paging_map_pages(0x00000000, 0x00000000, 1024, PAGE_RW);
 	current_page_directory = kernel_directory;
 
 	// Setup the kernel heap

kernel/src/mem/pmm.c

@@ -6,13 +6,17 @@
 #include <stdlib.h>
 #include <stdio.h>
 
+#define PHYS_TO_VIRT(addr) ((addr) + KERNEL_BASE_VIRT)
+#define VIRT_TO_PHYS(addr) ((addr) - KERNEL_BASE_VIRT)
+
 static uint32_t mem_size;
 static uint32_t used_blocks;
 static uint32_t max_blocks;
+static uintptr_t kernel_end;
 
 uint32_t* bitmap;
 
-// Linker-provided symbols
+// Linker-provided symbols. Beware, those don't take into account GRUB's things
 extern uint32_t KERNEL_END;
 extern uint32_t KERNEL_END_PHYS;
 
@@ -22,15 +26,30 @@ uint32_t mmap_test(uint32_t bit);
 uint32_t mmap_find_free();
 uint32_t mmap_find_free_frame(uint32_t num);
 
-// mem is in KiB
 void init_pmm(multiboot_t* boot) {
+	// Compute where the kernel & GRUB modules ends in physical memory
+	mod_t* modules = (mod_t*) boot->mods_addr;
+
+	for (uint32_t i = 0; i < boot->mods_count; i++) {
+		mod_t mod = modules[i];
+		kernel_end = PHYS_TO_VIRT(mod.mod_end);
+	}
+
+	// In the author's case, modules come after the kernel,
+	// but there's no guarantee it'll always be true
+	if (kernel_end < (uintptr_t) KERNEL_END) {
+		kernel_end = (uintptr_t) KERNEL_END;
+	}
+
+	// Setup our bitmap right after the kernel
+	bitmap = (uint32_t*) kernel_end;
 	mem_size = boot->mem_lower + boot->mem_upper;
-	bitmap = &KERNEL_END;
-	max_blocks = (mem_size*1024) / PMM_BLOCK_SIZE;
+	max_blocks = (mem_size*1024) / PMM_BLOCK_SIZE; // `mem_size` is in KiB
 	used_blocks = max_blocks;
 
-	memset(bitmap, 0xFF, max_blocks/8);
+	memset(bitmap, 0xFF, max_blocks/8); // Every block is initially taken
 
+	// Parse the memory map to mark valid areas as available
 	uint64_t available = 0;
 	uint64_t unavailable = 0;
 
@@ -53,7 +72,7 @@ void init_pmm(multiboot_t* boot) {
 	}
 
 	// Protect low memory, our glorious kernel and the PMM itself
-	pmm_deinit_region((uintptr_t) 0, (uint32_t) &KERNEL_END_PHYS + max_blocks/8);
+	pmm_deinit_region(0, pmm_get_kernel_end());
 
 	printf("[PMM] Memory stats: available: \x1B[32m%dMiB", available >> 20);
 	printf("\x1B[37m unavailable: \x1B[32m%dKiB\x1B[37m\n", unavailable >> 10);
@@ -182,7 +201,8 @@ uint32_t mmap_test(uint32_t bit) {
 	return bitmap[bit / 32] & (1 << (bit % 32));
 }
 
-// Returns the first free bit
+/* Returns the index of the first free bit in the bitmap
+ */
 uint32_t mmap_find_free() {
 	for (uint32_t i = 0; i < max_blocks/32; i++) {
 		if (bitmap[i] != 0xFFFFFFFF) {
@@ -197,7 +217,8 @@ uint32_t mmap_find_free() {
 	return 0;
 }
 
-// Returns the first block of frame_size bits
+/* Returns the first block of frame_size bits
+ */
 uint32_t mmap_find_free_frame(uint32_t frame_size) {
 	uint32_t first = 0;
 	uint32_t count = 0;
@@ -228,3 +249,9 @@ uint32_t mmap_find_free_frame(uint32_t frame_size) {
 
 	return 0;
 }
+
+/* Returns the first address after the kernel in physical memory.
+ */
+uintptr_t pmm_get_kernel_end() {
+	return VIRT_TO_PHYS(kernel_end + max_blocks/8);
+}

kernel/src/sys/proc.c

@@ -292,7 +292,7 @@ void* proc_sbrk(intptr_t size) {
 			uint32_t needed_size = size - remaining_bytes;
 			uint32_t num = needed_size / 0x1000 + (needed_size % 0x1000 ? 1 : 0);
 
-			paging_alloc_pages(ALIGN(end, 0x1000), num);
+			paging_alloc_pages(align_to(end, 0x1000), num);
 		}
 	} else if (size < 0) {
 		if (end + size < 0x1000*current_process->code_len) {

modules/src/loop.c

@@ -2,9 +2,7 @@
 
 #include <stdio.h>
 
-// #define RELEASE
-
-#ifdef RELEASE
+#if 0
 #include "bg.h"
 #endif
 

modules/src/test.c

@@ -2,8 +2,7 @@
 #include <string.h>
 
 int main() {
-	snow_sleep(2000);
-	window_t* win = snow_open_window("A static window", 800, 120, WM_NORMAL);
+	window_t* win = snow_open_window("A static window", 200, 120, WM_NORMAL);
 
 	snow_draw_window(win); // Draws the title bar and borders
 	snow_draw_string(win->fb, "Lorem Ipsum", 45, 55, 0x00AA1100);