refactor(aarch64,riscv64): unify kernel image detection

Cargo.toml

@@ -30,6 +30,7 @@ volatile = { version = "0.6", features = ["derive"] }
 
 [target.'cfg(target_arch = "riscv64")'.dependencies]
 fdt = "0.1"
+goblin = { version = "0.10", default-features = false, features = ["elf64"] }
 naked-function = "0.1"
 sbi-rt = "0.0.3"
 

src/arch/aarch64/mod.rs

@@ -12,15 +12,14 @@ use aarch64_cpu::asm::barrier::{self, NSH, SY, dmb, dsb, isb};
 use aarch64_cpu::registers::{ReadWriteable, SCTLR_EL1, TTBR0_EL1, TTBR1_EL1, Writeable};
 use align_address::Align;
 use fdt::Fdt;
-use goblin::elf::header::header64::{EI_DATA, ELFDATA2LSB, ELFMAG, Header, SELFMAG};
 use hermit_entry::Entry;
 use hermit_entry::boot_info::{BootInfo, HardwareInfo, PlatformInfo, RawBootInfo, SerialPortBase};
 use hermit_entry::elf::LoadedKernel;
 use log::info;
 
-use crate::BootInfoExt;
 use crate::arch::paging::*;
 use crate::os::CONSOLE;
+use crate::{BootInfoExt, FdtExt};
 
 unsafe extern "C" {
 	static mut l0_pgtable: u64;
@@ -57,49 +56,8 @@ pub fn find_kernel() -> &'static [u8] {
 		Fdt::from_ptr(ptr::with_exposed_provenance(DEVICE_TREE as usize))
 			.expect(".fdt file has invalid header")
 	};
-	let module_start = fdt
-		.find_node("/chosen")
-		.unwrap()
-		.children()
-		.find(|node| node.name.starts_with("module@"))
-		.map(|node| {
-			let value = node.name.strip_prefix("module@").unwrap();
-			if let Some(value) = value.strip_prefix("0x") {
-				usize::from_str_radix(value, 16).unwrap()
-			} else if let Some(value) = value.strip_prefix("0X") {
-				usize::from_str_radix(value, 16).unwrap()
-			} else {
-				value.parse().unwrap()
-			}
-		})
-		.unwrap();
-
-	let header = unsafe {
-		&*core::mem::transmute::<*const u8, *const Header>(ptr::with_exposed_provenance(
-			module_start,
-		))
-	};
-
-	if header.e_ident[0..SELFMAG] != ELFMAG[..] {
-		panic!("Didn't find valid ELF file!");
-	}
-
-	let file_size = if header.e_ident[EI_DATA] == ELFDATA2LSB {
-		u64::from_le(header.e_shoff)
-			+ (u16::from_le(header.e_shentsize) as u64 * u16::from_le(header.e_shnum) as u64)
-	} else {
-		u64::from_be(header.e_shoff)
-			+ (u16::from_be(header.e_shentsize) as u64 * u16::from_be(header.e_shnum) as u64)
-	};
 
-	info!("Found ELF file with size {file_size}");
-
-	unsafe {
-		core::slice::from_raw_parts(
-			ptr::with_exposed_provenance(module_start),
-			file_size.try_into().unwrap(),
-		)
-	}
+	fdt.find_kernel().unwrap()
 }
 
 pub unsafe fn boot_kernel(kernel_info: LoadedKernel) -> ! {

src/arch/riscv64/mod.rs

@@ -4,57 +4,21 @@ mod address_range;
 mod start;
 
 use core::arch::asm;
-use core::{mem, ptr, slice};
+use core::ptr;
 
 use address_range::AddressRange;
-use fdt::node::FdtNode;
 use hermit_entry::Entry;
 use hermit_entry::boot_info::{
 	BootInfo, DeviceTreeAddress, HardwareInfo, PlatformInfo, RawBootInfo,
 };
 use hermit_entry::elf::LoadedKernel;
 use log::info;
 
-use crate::BootInfoExt;
-
-fn find_kernel_linux(chosen: &FdtNode<'_, '_>) -> Option<&'static [u8]> {
-	let initrd_start = chosen.property("linux,initrd-start")?.as_usize()?;
-	let initrd_start = ptr::with_exposed_provenance_mut::<u8>(initrd_start);
-	let initrd_end = chosen.property("linux,initrd-end")?.as_usize()?;
-	let initrd_end = ptr::with_exposed_provenance_mut::<u8>(initrd_end);
-	// SAFETY: We trust the raw pointer from the firmware
-	let initrd_len = unsafe { initrd_end.offset_from(initrd_start).try_into().unwrap() };
-
-	// SAFETY: We trust the raw pointer from the firmware
-	Some(unsafe { slice::from_raw_parts(initrd_start, initrd_len) })
-}
-
-fn find_kernel_multiboot(chosen: &FdtNode<'_, '_>) -> Option<&'static [u8]> {
-	let module = chosen
-		.children()
-		.filter(|child| child.name.starts_with("module@"))
-		.find(|child| {
-			child.compatible().is_some_and(|compatible| {
-				compatible
-					.all()
-					.any(|compatible| compatible == "multiboot,ramdisk")
-			})
-		})?;
-	let reg = module.property("reg").unwrap();
-	let addr = usize::from_be_bytes(reg.value[..mem::size_of::<usize>()].try_into().unwrap());
-	let len = usize::from_be_bytes(reg.value[mem::size_of::<usize>()..].try_into().unwrap());
-
-	let initrd_start = ptr::with_exposed_provenance_mut::<u8>(addr);
-	// SAFETY: We trust the raw pointer from the firmware
-	Some(unsafe { slice::from_raw_parts(initrd_start, len) })
-}
+use crate::{BootInfoExt, FdtExt};
 
 pub fn find_kernel() -> &'static [u8] {
 	let fdt = start::get_fdt();
-	let chosen = fdt.find_node("/chosen").unwrap();
-	find_kernel_linux(&chosen)
-		.or_else(|| find_kernel_multiboot(&chosen))
-		.expect("could not find kernel")
+	fdt.find_kernel().expect("could not find kernel")
 }
 
 pub unsafe fn get_memory(memory_size: u64) -> u64 {

src/main.rs

@@ -6,6 +6,8 @@
 #![allow(clippy::missing_safety_doc)]
 
 use ::log::info;
+#[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
+use goblin::elf::header::header64::{EI_DATA, ELFDATA2LSB, ELFMAG, Header, SELFMAG};
 use hermit_entry::boot_info::{BootInfo, RawBootInfo};
 
 #[macro_use]
@@ -39,6 +41,55 @@ impl BootInfoExt for BootInfo {
 	}
 }
 
+#[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
+trait FdtExt {
+	fn find_kernel(&self) -> Option<&'static [u8]>;
+}
+
+#[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
+impl FdtExt for fdt::Fdt<'_> {
+	fn find_kernel(&self) -> Option<&'static [u8]> {
+		let chosen = self.find_node("/chosen").unwrap();
+
+		// The reg size of the module nodes is always 0, so we cannot trust them and
+		// instead need to parse the ELF header
+		let find_module_start = || {
+			let module = chosen
+				.children()
+				.find(|node| node.name.starts_with("module@"))?;
+			let start_ptr = module.reg().unwrap().next().unwrap().starting_address;
+			let header = unsafe { &*core::mem::transmute::<*const u8, *const Header>(start_ptr) };
+
+			if header.e_ident[0..SELFMAG] != ELFMAG[..] {
+				return None;
+			}
+
+			let len = if header.e_ident[EI_DATA] == ELFDATA2LSB {
+				u64::from_le(header.e_shoff)
+					+ (u16::from_le(header.e_shentsize) as u64
+						* u16::from_le(header.e_shnum) as u64)
+			} else {
+				u64::from_be(header.e_shoff)
+					+ (u16::from_be(header.e_shentsize) as u64
+						* u16::from_be(header.e_shnum) as u64)
+			};
+
+			Some(unsafe { core::slice::from_raw_parts(start_ptr, len.try_into().unwrap()) })
+		};
+
+		let find_linux_initrd = || {
+			let start = chosen.property("linux,initrd-start")?.as_usize()?;
+			let end = chosen.property("linux,initrd-end")?.as_usize()?;
+			let start_ptr = core::ptr::with_exposed_provenance::<u8>(start);
+			let end_ptr = core::ptr::with_exposed_provenance::<u8>(end);
+			let len = unsafe { end_ptr.offset_from(start_ptr).try_into().unwrap() };
+			Some(unsafe { core::slice::from_raw_parts(start_ptr, len) })
+		};
+
+		find_module_start().or_else(find_linux_initrd)
+	}
+}
+
 #[doc(hidden)]
 fn _print(args: core::fmt::Arguments<'_>) {
 	use core::fmt::Write;

xtask/src/ci/u-boot/boot.txt

@@ -7,7 +7,7 @@ fdt get value hermit_load /images/ramdisk-1 load
 fdt addr ${fdt_addr}
 fdt mknode /chosen module@${hermit_load}
 fdt set /chosen/module@${hermit_load} compatible "multiboot,module\0multiboot,ramdisk"
-fdt set /chosen/module@${hermit_load} reg <${hermit_load} 0>
+fdt set /chosen/module@${hermit_load} reg <0x0 ${hermit_load} 0x0 0x0>
 fdt print
 
 bootm