misc: Add measures to prevent stage 2 from accidentally relying on stage 3 before the latter is loaded

stage23/Makefile

@@ -61,6 +61,9 @@ limine.sys: limine.elf
 	$(OBJCOPY) -O binary $< $@
 
 limine_nomap.elf: $(OBJ)
+	$(LD) $(OBJ) $(LDFLAGS) $(INTERNAL_LDFLAGS) -Wl,--gc-sections -Tlinker_stage2only.ld -o $@ || \
+		( echo "This error means that stage2 was trying to use stage3 symbols before loading stage 3" && \
+		  false )
 	$(LD) $(OBJ) $(LDFLAGS) $(INTERNAL_LDFLAGS) -Tlinker_nomap.ld -o $@
 
 limine.elf: $(OBJ) limine.map.o

stage23/lib/blib.c

@@ -6,9 +6,10 @@
 #include <lib/print.h>
 #include <lib/trace.h>
 #include <lib/real.h>
+#include <fs/file.h>
 
-uint64_t build_id = BUILD_ID;
-stage3_data uint64_t stage3_build_id = BUILD_ID;
+__attribute__((section(".stage3_build_id")))
+uint64_t stage3_build_id = BUILD_ID;
 
 uint8_t boot_drive;
 int     boot_partition = -1;
@@ -17,6 +18,39 @@ bool booted_from_pxe = false;
 bool booted_from_cd = false;
 bool stage3_loaded = false;
 
+extern symbol stage3_addr;
+extern symbol limine_sys_size;
+
+__attribute__((noreturn))
+void (*stage3)(void) = (void *)stage3_addr;
+
+bool stage3_init(struct volume *part) {
+    struct file_handle stage3;
+
+    if (fopen(&stage3, part, "/limine.sys")
+     && fopen(&stage3, part, "/boot/limine.sys")) {
+        return false;
+    }
+
+    if (stage3.size != (size_t)limine_sys_size) {
+        print("limine.sys size incorrect.\n");
+        return false;
+    }
+
+    fread(&stage3, stage3_addr,
+          (uintptr_t)stage3_addr - 0x8000,
+          stage3.size - ((uintptr_t)stage3_addr - 0x8000));
+
+    if (BUILD_ID != stage3_build_id) {
+        print("limine.sys build ID mismatch.\n");
+        return false;
+    }
+
+    stage3_loaded = true;
+
+    return true;
+}
+
 bool parse_resolution(int *width, int *height, int *bpp, const char *buf) {
     int res[3] = {0};
 

stage23/lib/blib.h

@@ -4,6 +4,7 @@
 #include <stddef.h>
 #include <stdint.h>
 #include <stdbool.h>
+#include <fs/file.h>
 
 extern uint8_t boot_drive;
 extern int     boot_partition;
@@ -48,7 +49,7 @@ typedef char symbol[];
 #define stage3_text __attribute__((section(".stage3_text")))
 #define stage3_data __attribute__((section(".stage3_data")))
 
-extern uint64_t build_id;
-extern uint64_t stage3_build_id;
+__attribute__((noreturn)) extern void (*stage3)(void);
+bool stage3_init(struct volume *part);
 
 #endif

stage23/lib/config.c

@@ -16,28 +16,7 @@ bool config_ready = false;
 
 static char *config_addr;
 
-extern symbol stage3_addr;
-
 int init_config_disk(struct volume *part) {
-    if (!stage3_loaded) {
-        struct file_handle stage3;
-
-        if (fopen(&stage3, part, "/limine.sys")
-         && fopen(&stage3, part, "/boot/limine.sys")) {
-            panic("Could not open stage 3");
-        }
-
-        fread(&stage3, stage3_addr,
-              (uintptr_t)stage3_addr - 0x8000,
-              stage3.size - ((uintptr_t)stage3_addr - 0x8000));
-
-        if (build_id != stage3_build_id) {
-            panic("Limine build ID mismatch. Use the correct limine.sys.");
-        }
-
-        stage3_loaded = true;
-    }
-
     struct file_handle f;
 
     if (fopen(&f, part, "/limine.cfg")
@@ -56,12 +35,6 @@ int init_config_disk(struct volume *part) {
 }
 
 int init_config_pxe(void) {
-    struct tftp_file_handle stage3;
-    if (tftp_open(&stage3, 0, 69, "limine.sys")) {
-        panic("Could not open stage 3");
-    }
-    tftp_read(&stage3, stage3_addr, 0, stage3.file_size);
-
     struct tftp_file_handle cfg;
     if (tftp_open(&cfg, 0, 69, "limine.cfg")
      && tftp_open(&cfg, 0, 69, "tomatboot.cfg")) {

stage23/lib/term.c

@@ -4,6 +4,7 @@
 #include <lib/term.h>
 #include <lib/real.h>
 #include <lib/image.h>
+#include <lib/blib.h>
 #include <drivers/vga_textmode.h>
 #include <drivers/vbe.h>
 
@@ -27,6 +28,7 @@ void (*term_double_buffer_flush)(void);
 
 int term_rows, term_cols;
 
+stage3_text
 void term_vbe(uint32_t *colours, int margin, int margin_gradient, struct image *background) {
     term_deinit();
     if (!vbe_tty_init(&term_rows, &term_cols, colours, margin, margin_gradient, background)) {

stage23/linker.ld

@@ -6,11 +6,11 @@ SECTIONS
     . = 0x8000;
 
     .entry : {
-        KEEP(*(.entry*))
+        *(.entry*)
     }
 
     .realmode : {
-        KEEP(*(.realmode*))
+        *(.realmode*)
     }
 
     .stage2.text : {
@@ -24,6 +24,8 @@ SECTIONS
 
     .stage3.text : {
         stage3_addr = .;
+        *(.stage3_entry*)
+        *(.stage3_build_id*)
         *(.stage3_text*)
     }
 
@@ -32,7 +34,8 @@ SECTIONS
     }
 
     .map : {
-        KEEP(*(.map*))
+        *(.map*)
+        limine_sys_size = . - 0x8000;
     }
 
     .bss : {

stage23/linker_nomap.ld

@@ -6,11 +6,11 @@ SECTIONS
     . = 0x8000;
 
     .entry : {
-        KEEP(*(.entry*))
+        *(.entry*)
     }
 
     .realmode : {
-        KEEP(*(.realmode*))
+        *(.realmode*)
     }
 
     .stage2.text : {
@@ -24,6 +24,8 @@ SECTIONS
 
     .stage3.text : {
         stage3_addr = .;
+        *(.stage3_entry*)
+        *(.stage3_build_id*)
         *(.stage3_text*)
     }
 
@@ -33,6 +35,7 @@ SECTIONS
 
     .map : {
         limine_map = .;
+        limine_sys_size = .;
     }
 
     .bss : {

stage23/linker_stage2only.ld

@@ -0,0 +1,39 @@
+OUTPUT_FORMAT(elf32-i386)
+ENTRY(_start)
+
+SECTIONS
+{
+    . = 0x8000;
+
+    .entry : {
+        *(.entry*)
+    }
+
+    .realmode : {
+        *(.realmode*)
+    }
+
+    .stage2.text : {
+        *(.text*)
+    }
+
+    .stage2.data : {
+        *(.data*)
+        *(.rodata*)
+        stage3_addr = .;
+        *(.stage3_build_id*)
+        limine_map = .;
+        limine_sys_size = .;
+    }
+
+    .bss : {
+        bss_begin = .;
+        *(COMMON)
+        *(.bss*)
+        bss_end = .;
+    }
+
+    /DISCARD/ : {
+        *(*)
+    }
+}

stage23/main.c

@@ -26,6 +26,7 @@ enum {
 	BOOT_FROM_CD
 };
 
+__attribute__((noreturn))
 void entry(uint8_t _boot_drive, int boot_from) {
     boot_drive = _boot_drive;
 
@@ -44,31 +45,39 @@ void entry(uint8_t _boot_drive, int boot_from) {
     init_e820();
     init_memmap();
 
-    struct volume part;
     volume_create_index();
 
     switch (boot_from) {
-    case BOOT_FROM_HDD:
+    case BOOT_FROM_HDD: {
         print("Boot drive: %x\n", boot_drive);
-        // Look for config file.
-        print("Searching for config file...\n");
+        struct volume boot_volume;
+        volume_get_by_coord(&boot_volume, boot_drive, -1);
+        struct volume part = boot_volume;
+        bool stage3_loaded = false, config_loaded = false;
         for (int i = 0; ; i++) {
-            int ret = volume_get_by_coord(&part, boot_drive, i);
+            if (!stage3_loaded && stage3_init(&part)) {
+                stage3_loaded = true;
+                print("Stage 3 found and loaded.\n");
+            }
+            if (!config_loaded && !init_config_disk(&part)) {
+                config_loaded = true;
+                print("Config file found and loaded.\n");
+                boot_partition = i - 1;
+            }
+            int ret = part_get(&part, &boot_volume, i);
             switch (ret) {
                 case INVALID_TABLE:
-                    panic("Partition table of boot drive is invalid.");
                 case END_OF_TABLE:
+                    goto break2;
+            }
+        }
+break2:
+        if (!stage3_loaded)
+            panic("Stage 3 not loaded.");
+        if (!config_loaded)
             panic("Config file not found.");
-                case NO_PARTITION:
-                    continue;
-            }
-            if (!init_config_disk(&part)) {
-                print("Config file found and loaded.\n");
-                boot_partition = i;
         break;
     }
-        }
-        break;
 
     case BOOT_FROM_PXE:
         pxe_init();
@@ -77,15 +86,14 @@ void entry(uint8_t _boot_drive, int boot_from) {
         }
         print("Config loaded via PXE\n");
         break;
-
-    case BOOT_FROM_CD:
-        boot_partition = -1;  // raw device
-        volume_get_by_coord(&part, boot_drive, boot_partition);
-        if (init_config_disk(&part))
-            panic("Failed to load config file");
-        break;
     }
 
+    stage3();
+}
+
+__attribute__((noreturn))
+__attribute__((section(".stage3_entry")))
+void stage3_entry(void) {
     char *cmdline;
     char *config = menu(&cmdline);
 
@@ -105,4 +113,6 @@ void entry(uint8_t _boot_drive, int boot_from) {
     } else {
         panic("Invalid protocol specified");
     }
+
+    for (;;);
 }