Make executable smaller and improve screen output

Author: phil-opp

real_mode/first_stage/src/boot.s

@@ -65,11 +65,13 @@ unreal_mode:
     mov word ptr ds:[eax], bx
 
 check_int13h_extensions:
+    push 'y'    # error code
     mov ah, 0x41
     mov bx, 0x55aa
     # dl contains drive number
     int 0x13
-    jc no_int13h_extensions
+    jc fail
+    pop ax      # pop error code again
 
 rust:
     # push arguments

real_mode/first_stage/src/dap.rs

@@ -41,10 +41,13 @@ impl DiskAddressPacket {
 
     pub unsafe fn perform_load(&self, disk_number: u16) {
         let self_addr = self as *const Self as u16;
-        asm!("mov {1:x}, si",
+        asm!(
+            "push 0x7a", // error code `z`, passed to `fail` on error
+            "mov {1:x}, si",
             "mov si, {0:x}",
             "int 0x13",
-            "jc dap_load_failed",
+            "jc fail",
+            "pop si", // remove error code again
             "mov si, {1:x}",
             in(reg) self_addr,
             out(reg) _,

real_mode/first_stage/src/fail.rs

@@ -0,0 +1,59 @@
+use core::arch::asm;
+
+pub trait UnwrapOrFail {
+    type Out;
+
+    fn unwrap_or_fail(self, code: u8) -> Self::Out;
+}
+
+impl<T> UnwrapOrFail for Option<T> {
+    type Out = T;
+
+    fn unwrap_or_fail(self, code: u8) -> Self::Out {
+        match self {
+            Some(v) => v,
+            None => fail(code),
+        }
+    }
+}
+
+impl<T, E> UnwrapOrFail for Result<T, E> {
+    type Out = T;
+
+    fn unwrap_or_fail(self, code: u8) -> Self::Out {
+        match self {
+            Ok(v) => v,
+            Err(_) => fail(code),
+        }
+    }
+}
+
+#[no_mangle]
+pub extern "C" fn print_char(c: u8) {
+    let ax = u16::from(c) | 0x0e00;
+    unsafe {
+        asm!("int 0x10", in("ax") ax, in("bx") 0);
+    }
+}
+
+#[cold]
+#[inline(never)]
+#[no_mangle]
+pub extern "C" fn fail(code: u8) -> ! {
+    print_char(b'!');
+    print_char(code);
+    loop {
+        hlt()
+    }
+}
+
+fn hlt() {
+    unsafe {
+        asm!("hlt");
+    }
+}
+
+#[panic_handler]
+pub fn panic(_info: &core::panic::PanicInfo) -> ! {
+    fail(b'P');
+}

real_mode/first_stage/src/fat.rs

@@ -1,5 +1,7 @@
 // based on https://github.com/rafalh/rust-fatfs/
 
+use super::split_array_ref;
+
 pub(crate) struct BootSector {
     bootjmp: [u8; 3],
     oem_name: [u8; 8],
@@ -157,13 +159,3 @@ impl BiosParameterBlock {
         self.sectors_per_fat_16 == 0
     }
 }
-
-/// Taken from https://github.com/rust-lang/rust/blob/e100ec5bc7cd768ec17d75448b29c9ab4a39272b/library/core/src/slice/mod.rs#L1673-L1677
-///
-/// TODO replace with `split_array` feature in stdlib as soon as it's stabilized,
-/// see https://github.com/rust-lang/rust/issues/90091
-fn split_array_ref<const N: usize, T>(slice: &[T]) -> (&[T; N], &[T]) {
-    let (a, b) = slice.split_at(N);
-    // SAFETY: a points to [T; N]? Yes it's [T] of length N (checked by split_at)
-    unsafe { (&*(a.as_ptr() as *const [T; N]), b) }
-}

real_mode/first_stage/src/main.rs

@@ -5,11 +5,13 @@ use core::{
     arch::{asm, global_asm},
     slice,
 };
+use fail::{fail, print_char, UnwrapOrFail};
 use mbr::MasterBootRecord;
 
 global_asm!(include_str!("boot.s"));
 
 mod dap;
+mod fail;
 mod fat;
 mod mbr;
 
@@ -23,34 +25,38 @@ fn mbr_start() -> *const u8 {
 
 #[no_mangle]
 pub extern "C" fn first_stage(disk_number: u16) {
+    print_char(b'1');
     let bytes = &unsafe { slice::from_raw_parts(mbr_start(), 512) };
-    let mbr = MasterBootRecord::from_bytes(bytes);
+    let partition = mbr::get_partition(bytes, 0);
 
-    let partition = mbr
-        .partition_table_entries()
-        .get(0)
-        .unwrap_or_else(|| panic!());
-
-    let partition_buf = u16::try_from(mbr_start() as usize).unwrap_or_else(|_| panic!()) + 512;
+    print_char(b'2');
+    let partition_buf = u16::try_from(mbr_start() as usize).unwrap_or_fail(b'a') + 512;
 
     // load first partition into buffer
     // TODO: only load headers
     let dap = dap::DiskAddressPacket::from_lba(
         partition_buf,
         partition.logical_block_address.into(),
-        partition.sector_count.try_into().unwrap(),
+        partition.sector_count.try_into().unwrap_or_fail(b'b'),
     );
     unsafe {
         dap.perform_load(disk_number);
     }
+    if partition.sector_count == 0 {
+        fail(b'c');
+    }
+
+    print_char(b'3');
 
     // try to parse FAT file system
     let fat_slice = unsafe {
         slice::from_raw_parts(
             partition_buf as *const u8,
-            usize::try_from(partition.sector_count).unwrap_or_else(|_| panic!()) * 512,
+            usize::try_from(partition.sector_count).unwrap_or_else(|_| fail(b'a')) * 512,
         )
     };
+
+    print_char(b'4');
     let boot_sector = fat::BootSector::deserialize(fat_slice);
 
     // TODO: get root dir
@@ -81,42 +87,15 @@ fn load_second_stage(
     unsafe { dap.perform_load(disk_number) }
 }
 
-#[no_mangle]
-pub extern "C" fn print_char(c: u8) {
-    let ax = u16::from(c) | 0x0e00;
-    unsafe {
-        asm!("int 0x10", in("ax") ax, in("bx") 0);
-    }
-}
-
-#[no_mangle]
-pub extern "C" fn dap_load_failed() -> ! {
-    err(b'1');
-}
-
-#[no_mangle]
-pub extern "C" fn no_int13h_extensions() -> ! {
-    err(b'2');
-}
-
-#[cold]
-fn err(code: u8) -> ! {
-    for &c in b"Err:" {
-        print_char(c);
-    }
-    print_char(code);
-    loop {
-        hlt()
+/// Taken from https://github.com/rust-lang/rust/blob/e100ec5bc7cd768ec17d75448b29c9ab4a39272b/library/core/src/slice/mod.rs#L1673-L1677
+///
+/// TODO replace with `split_array` feature in stdlib as soon as it's stabilized,
+/// see https://github.com/rust-lang/rust/issues/90091
+fn split_array_ref<const N: usize, T>(slice: &[T]) -> (&[T; N], &[T]) {
+    if N > slice.len() {
+        fail(b'S');
     }
-}
-
-fn hlt() {
-    unsafe {
-        asm!("hlt");
-    }
-}
-
-#[panic_handler]
-pub fn panic(_info: &core::panic::PanicInfo) -> ! {
-    err(b'P');
+    let (a, b) = slice.split_at(N);
+    // SAFETY: a points to [T; N]? Yes it's [T] of length N (checked by split_at)
+    unsafe { (&*(a.as_ptr() as *const [T; N]), b) }
 }

real_mode/first_stage/src/mbr.rs

@@ -1,5 +1,36 @@
 // Based on https://docs.rs/mbr-nostd
 
+use super::fail::{fail, UnwrapOrFail};
+
+pub fn get_partition(buffer: &[u8], index: usize) -> PartitionTableEntry {
+    if buffer.len() < BUFFER_SIZE {
+        fail(b'a');
+    } else if buffer.get(BUFFER_SIZE - SUFFIX_BYTES.len()..BUFFER_SIZE) != Some(&SUFFIX_BYTES[..]) {
+        fail(b'b');
+    }
+
+    let offset = TABLE_OFFSET + index * ENTRY_SIZE;
+    let buffer = buffer.get(offset..).unwrap_or_fail(b'c');
+
+    let partition_type = *buffer.get(4).unwrap_or_fail(b'd');
+
+    let lba = u32::from_le_bytes(
+        buffer
+            .get(8..)
+            .and_then(|s| s.get(..4))
+            .and_then(|s| s.try_into().ok())
+            .unwrap_or_fail(b'e'),
+    );
+    let len = u32::from_le_bytes(
+        buffer
+            .get(12..)
+            .and_then(|s| s.get(..4))
+            .and_then(|s| s.try_into().ok())
+            .unwrap_or_fail(b'f'),
+    );
+    PartitionTableEntry::new(partition_type, lba, len)
+}
+
 /// A struct representing an MBR partition table.
 pub struct MasterBootRecord {
     entries: [PartitionTableEntry; MAX_ENTRIES],
@@ -16,26 +47,38 @@ impl MasterBootRecord {
 
     pub fn from_bytes(buffer: &[u8]) -> MasterBootRecord {
         if buffer.len() < BUFFER_SIZE {
-            panic!();
-        } else if buffer[BUFFER_SIZE - SUFFIX_BYTES.len()..BUFFER_SIZE] != SUFFIX_BYTES[..] {
-            panic!();
+            fail(b'1');
+        } else if buffer.get(BUFFER_SIZE - SUFFIX_BYTES.len()..BUFFER_SIZE)
+            != Some(&SUFFIX_BYTES[..])
+        {
+            fail(b'2');
         }
         let mut entries = [PartitionTableEntry::empty(); MAX_ENTRIES];
+
         for idx in 0..MAX_ENTRIES {
             let offset = TABLE_OFFSET + idx * ENTRY_SIZE;
-            let partition_type = PartitionType::from_mbr_tag_byte(buffer[offset + 4]);
-            if let PartitionType::Unknown(c) = partition_type {
-                panic!();
-            }
-            let lba =
-                u32::from_le_bytes(buffer[offset + 8..].try_into().unwrap_or_else(|_| panic!()));
+            let buffer = buffer.get(offset..).unwrap_or_fail(b'8');
+
+            let partition_type = *buffer.get(4).unwrap_or_fail(b'4');
+
+            let lba = u32::from_le_bytes(
+                buffer
+                    .get(8..)
+                    .and_then(|s| s.get(..4))
+                    .and_then(|s| s.try_into().ok())
+                    .unwrap_or_fail(b'5'),
+            );
             let len = u32::from_le_bytes(
-                buffer[offset + 12..]
-                    .try_into()
-                    .unwrap_or_else(|_| panic!()),
+                buffer
+                    .get(12..)
+                    .and_then(|s| s.get(..4))
+                    .and_then(|s| s.try_into().ok())
+                    .unwrap_or_fail(b'6'),
             );
-            entries[idx] = PartitionTableEntry::new(partition_type, lba, len);
+            *entries.get_mut(idx).unwrap_or_fail(b'7') =
+                PartitionTableEntry::new(partition_type, lba, len);
         }
+
         MasterBootRecord { entries }
     }
 
@@ -45,7 +88,7 @@ impl MasterBootRecord {
 }
 
 /// The type of a particular partition.
-#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
+#[derive(Copy, Clone, Eq, PartialEq)]
 pub enum PartitionType {
     Unused,
     Unknown(u8),
@@ -72,13 +115,22 @@ impl PartitionType {
             _ => PartitionType::Unknown(tag),
         }
     }
+
+    pub fn known_type(tag: u8) -> bool {
+        match tag {
+            0x0 | 0x01 | 0x04 | 0x06 | 0x0e | 0x0b | 0x0c | 0x1b | 0x1c | 0x83 | 0x07 | 0xaf => {
+                true
+            }
+            _ => false,
+        }
+    }
 }
 
 /// An entry in a partition table.
-#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
+#[derive(Copy, Clone, Eq, PartialEq)]
 pub struct PartitionTableEntry {
     /// The type of partition in this entry.
-    pub partition_type: PartitionType,
+    pub partition_type: u8,
 
     /// The index of the first block of this entry.
     pub logical_block_address: u32,
@@ -89,7 +141,7 @@ pub struct PartitionTableEntry {
 
 impl PartitionTableEntry {
     pub fn new(
-        partition_type: PartitionType,
+        partition_type: u8,
         logical_block_address: u32,
         sector_count: u32,
     ) -> PartitionTableEntry {
@@ -101,6 +153,6 @@ impl PartitionTableEntry {
     }
 
     pub fn empty() -> PartitionTableEntry {
-        PartitionTableEntry::new(PartitionType::Unused, 0, 0)
+        PartitionTableEntry::new(0, 0, 0)
     }
 }