From f564a042d884e5d57c7d039f5deef84224d16030 Mon Sep 17 00:00:00 2001
From: Alexander Myltsev <a.myltsev@gmail.com>
Date: Fri, 25 Nov 2022 10:42:07 +0000
Subject: [PATCH] Better ELF header parsing.

---
 Makefile      |  2 +-
 drivers/vga.c |  2 +-
 mbr.S         | 24 ++++++++++++++++--------
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/Makefile b/Makefile
index 6670f7d..d139caf 100644
--- a/Makefile
+++ b/Makefile
@@ -8,7 +8,7 @@ GDB=x86_64-elf-gdb
 endif
 
 CFLAGS = -fno-pic -ffreestanding -static -fno-builtin -fno-strict-aliasing \
-		 -O2 -Wall -MD -ggdb -m32 -Werror -fno-omit-frame-pointer
+		 -Os -Wall -MD -ggdb -m32 -Werror -fno-omit-frame-pointer
 CFLAGS += $(shell $(CC) -fno-stack-protector -E -x c /dev/null >/dev/null 2>&1 && echo -fno-stack-protector)
 
 run: image.bin
diff --git a/drivers/vga.c b/drivers/vga.c
index 844892d..beb624f 100644
--- a/drivers/vga.c
+++ b/drivers/vga.c
@@ -1,7 +1,7 @@
 #include "port.h"
 #include "../string.h"
 
-char* const video_memory = (char*) 0xb8000;
+static char* const video_memory = (char*) 0xb8000;
 
 enum colors16 {
     black = 0,
diff --git a/mbr.S b/mbr.S
index fea161f..42e5d60 100644
--- a/mbr.S
+++ b/mbr.S
@@ -30,9 +30,12 @@ get_drive_geometry:
 .equ ELF32_PHENTSIZE_OFFSET, ELF32_PHDR_OFFSET + 14
 .equ ELF32_PHNUM_OFFSET, ELF32_PHENTSIZE_OFFSET + 2
 .equ ELF32_PHDR_P_OFFSET, 4
+.equ ELF32_PHDR_PTYPE_OFFSET, 0
 .equ ELF32_PHDR_FILESZ_OFFSET, 4*4
 .equ KERNEL_OFFSET, 0x1000
 
+.equ PT_LOAD, 1
+
 .equ MBR_SECTORS, 2
 .equ SECTOR_BASE, 1
 .equ ELFHDR_SECTORS, 8
@@ -48,18 +51,23 @@ load_kernel:
     mov KERNEL_OFFSET + ELF32_ENTRY_OFFSET, %si
     mov %si, entry // store entry point
 
-    #if 0
-    mov KERNEL_OFFSET + ELF32_PHNUM_OFFSET, %ax
-    dec %ax  // no offset to the first entry
+    mov KERNEL_OFFSET + ELF32_PHNUM_OFFSET, %si
+read_segment:
+    dec %si  // no offset to the first entry
+    mov %si, %ax
     mulb KERNEL_OFFSET + ELF32_PHENTSIZE_OFFSET
     mov %ax, %di
     add KERNEL_OFFSET + ELF32_PHDR_OFFSET, %di
-    // now di holds offset to the last phentry
-    #else
-    mov KERNEL_OFFSET + ELF32_PHDR_OFFSET, %di
-    // now di holds offset to the first phentry
-    #endif
+    // now di holds offset to the phentry
+    mov KERNEL_OFFSET + ELF32_PHDR_PTYPE_OFFSET(%di), %ax
+    cmp $PT_LOAD, %ax
+    jnz read_segment // not a PT_LOAD segment
     mov KERNEL_OFFSET + ELF32_PHDR_FILESZ_OFFSET(%di), %ax
+    test %ax, %ax
+    jz read_segment // empty segment
+
+    // now di holds offset to the last phentry loaded from file, ax its filesz
+
     add KERNEL_OFFSET + ELF32_PHDR_P_OFFSET(%di), %ax
     sub $0x1000, %ax // we won't load the header
     add $SECTOR_SIZE - 1, %ax