From 24f1f412177e52486a824767f9f23d9cc0e3762b Mon Sep 17 00:00:00 2001
From: Mark Poliakov <mark@alnyan.me>
Date: Sat, 25 Jan 2025 02:27:57 +0200
Subject: [PATCH] tools: write riscv image size to header in gentables

---
 kernel/tools/gentables/src/main.rs | 89 ++++++++++++++++++++++++++----
 1 file changed, 77 insertions(+), 12 deletions(-)

diff --git a/kernel/tools/gentables/src/main.rs b/kernel/tools/gentables/src/main.rs
index 4e00f017..122d4d79 100644
--- a/kernel/tools/gentables/src/main.rs
+++ b/kernel/tools/gentables/src/main.rs
@@ -56,6 +56,10 @@ struct Args {
     symbol_out: PathBuf,
 }
 
+pub enum ImageHeader {
+    Riscv(u64),
+}
+
 pub struct GenData {
     pub kernel_start: u64,
     pub kernel_end: u64,
@@ -65,6 +69,13 @@ pub struct GenData {
     pub kernel_virt_offset: u64,
 }
 
+pub struct BuiltTables {
+    pub image_size: u64,
+    pub image_header: Option<ImageHeader>,
+    pub tables: Option<(AnyTables, u64)>,
+    pub symbol_table: HashMap<String, usize>,
+}
+
 fn kernel_image_range<F: Read + Seek>(
     elf: &mut ElfStream<AnyEndian, F>,
     kernel_virt_offset: u64,
@@ -113,16 +124,41 @@ fn kernel_virt_offset<F: Read + Seek>(elf: &mut ElfStream<AnyEndian, F>) -> Resu
     Err(GenError::MissingSymbol("KERNEL_VIRT_OFFSET"))
 }
 
-fn find_tables<F: Read + Seek>(elf: &mut ElfStream<AnyEndian, F>) -> Result<(u64, u64), GenError> {
+fn find_tables<F: Read + Seek>(
+    elf: &mut ElfStream<AnyEndian, F>,
+) -> Result<(Option<ImageHeader>, u64, u64), GenError> {
     let section_size = match elf.ehdr.e_machine {
         EM_AARCH64 => size_of::<memtables::aarch64::FixedTables>(),
         EM_X86_64 => size_of::<memtables::x86_64::FixedTables>(),
         EM_RISCV => size_of::<memtables::riscv64::FixedTables>(),
         _ => unimplemented!(),
     };
+
+    let image_header = if let Some(text_entry) = elf.section_header_by_name(".text.entry")? {
+        let header = text_entry.clone();
+        let section_offset = header.sh_offset;
+        let (data, _) = elf.section_data(&header)?;
+        if data.len() >= 64 {
+            let version = u32::from_le_bytes(data[32..36].try_into().unwrap());
+            let magic0 = &data[48..56];
+            let magic1 = &data[56..60];
+
+            match (version, magic0, magic1) {
+                (2, b"RISCV\x00\x00\x00", b"RSC\x05") => Some(ImageHeader::Riscv(section_offset)),
+                (_, _, _) => None,
+            }
+        } else {
+            None
+        }
+    } else {
+        None
+    };
+
     let (shdrs, strtab) = elf.section_headers_with_strtab()?;
     let strtab = strtab.ok_or_else(|| GenError::MissingSection(".strtab"))?;
 
+    let mut tables = None;
+
     for shdr in shdrs {
         let name = strtab.get(shdr.sh_name as _)?;
 
@@ -135,11 +171,13 @@ fn find_tables<F: Read + Seek>(elf: &mut ElfStream<AnyEndian, F>) -> Result<(u64
             }
 
             // TODO section checks
-            return Ok((shdr.sh_offset, shdr.sh_addr));
+            tables = Some((shdr.sh_offset, shdr.sh_addr));
         }
     }
 
-    Err(GenError::MissingSection(".data.tables"))
+    let (tables_offset, tables_addr) = tables.ok_or(GenError::MissingSection(".data.tables"))?;
+
+    Ok((image_header, tables_offset, tables_addr))
 }
 
 fn extract_symbols<F: Read + Seek>(
@@ -194,20 +232,23 @@ fn into_any<T: Into<AnyTables>, U, V>((x, y, z): (T, U, V)) -> (AnyTables, U, V)
     (x.into(), y, z)
 }
 
-fn build_tables<F: Read + Seek>(
-    file: F,
-) -> Result<(Option<(AnyTables, u64)>, HashMap<String, usize>), GenError> {
+fn build_tables<F: Read + Seek>(file: F) -> Result<BuiltTables, GenError> {
     let mut elf = ElfStream::<AnyEndian, F>::open_stream(file)?;
 
     if elf.ehdr.e_machine == EM_386 {
         // Locate symbol table
         let symbol_table = extract_symbols(&mut elf)?;
-        return Ok((None, symbol_table));
+        return Ok(BuiltTables {
+            image_size: 0,
+            image_header: None,
+            tables: None,
+            symbol_table,
+        });
     }
 
     let kernel_virt_offset = kernel_virt_offset(&mut elf)?;
     let (kernel_start, kernel_end) = kernel_image_range(&mut elf, kernel_virt_offset)?;
-    let (table_offset, table_virt_addr) = find_tables(&mut elf)?;
+    let (image_header, table_offset, table_virt_addr) = find_tables(&mut elf)?;
     let table_physical_address = table_virt_addr
         .checked_sub(kernel_virt_offset)
         .ok_or_else(|| GenError::IncorrectTablesPlacement(table_virt_addr))?;
@@ -230,7 +271,12 @@ fn build_tables<F: Read + Seek>(
         _ => todo!(),
     }?;
 
-    Ok((Some((tables, table_offset)), symbol_table))
+    Ok(BuiltTables {
+        image_size: kernel_end - kernel_start,
+        image_header,
+        tables: Some((tables, table_offset)),
+        symbol_table,
+    })
 }
 
 fn write_tables<F: Write + Seek>(
@@ -243,6 +289,22 @@ fn write_tables<F: Write + Seek>(
     Ok(())
 }
 
+fn write_image_header<F: Write + Seek>(
+    file: &mut F,
+    header: ImageHeader,
+    image_size: u64,
+) -> Result<(), GenError> {
+    match header {
+        ImageHeader::Riscv(offset) => {
+            let size_bytes = image_size.to_le_bytes();
+            println!("Writing RISC-V image header: image_size={image_size}");
+            file.seek(SeekFrom::Start(offset + 16))?;
+            file.write_all(&size_bytes)?;
+            Ok(())
+        }
+    }
+}
+
 fn write_symbol_table(
     out: impl AsRef<Path>,
     table: HashMap<String, usize>,
@@ -268,9 +330,12 @@ fn gentables(image: impl AsRef<Path>, symbol_out: impl AsRef<Path>) -> Result<()
         .truncate(false)
         .open(image)?;
 
-    let (tables, symbol_table) = build_tables(&mut file)?;
-    write_symbol_table(symbol_out, symbol_table)?;
-    if let Some((tables, file_offset)) = tables {
+    let built = build_tables(&mut file)?;
+    write_symbol_table(symbol_out, built.symbol_table)?;
+    if let Some(header) = built.image_header {
+        write_image_header(&mut file, header, built.image_size)?;
+    }
+    if let Some((tables, file_offset)) = built.tables {
         write_tables(file, file_offset, tables)?;
     }