sig new-elf: New subcommand to create function signatures from ELF files

cli/src/analysis/signature.rs

@@ -1,13 +1,13 @@
 use std::{borrow::Cow, path::Path};
 
-use anyhow::{Result, anyhow};
+use anyhow::{Context, Result, anyhow, bail};
 use base64::{Engine as _, engine::general_purpose::STANDARD};
 use ds_decomp::{
     analysis::functions::Function,
     config::{module::Module, relocations::RelocationKind, symbol::SymbolMaps},
 };
 use serde::{Deserialize, Serialize};
-use unarm::{ArmVersion, Endian, ParseFlags, ParseMode, Parser};
+use unarm::{ArmVersion, Endian, Ins, ParseFlags, ParseMode, Parser, arm, thumb};
 
 use crate::{config::program::Program, util::io::read_to_string};
 
@@ -54,6 +54,13 @@ fn is_zero(value: &i32) -> bool {
     *value == 0
 }
 
+pub struct SignatureRelocationInfo {
+    /// Name of the object this relocation points to.
+    pub name: String,
+    pub kind: RelocationKind,
+    pub addend: i32,
+}
+
 pub enum ApplyResult {
     /// The signature was successfully applied.
     Applied,
@@ -138,6 +145,62 @@ impl Signatures {
         })
     }
 
+    pub fn from_function_raw<GetRelocCb>(
+        function: &Function,
+        function_code: &[u8],
+        get_relocation: GetRelocCb,
+    ) -> Result<Self>
+    where
+        GetRelocCb: Fn(u32, Option<Ins>) -> Result<Option<SignatureRelocationInfo>>,
+    {
+        let mut bitmask = vec![0xff; function_code.len()];
+        let mut pattern = function_code.to_vec();
+        let mut relocations = Vec::new();
+
+        for (&address, call) in function.function_calls() {
+            let offset = (address - function.start_address()) as usize;
+            let (ins_bitmask, code) = match call.ins {
+                Ins::Arm(ins) => match ins.op {
+                    arm::Opcode::B | arm::Opcode::Bl | arm::Opcode::BlxI => (0xff000000, ins.code),
+                    op => bail!("Unexpected ARM opcode {op:?} for function call"),
+                },
+                Ins::Thumb(ins) => match ins.op {
+                    thumb::Opcode::Bl | thumb::Opcode::BlxI => (0xf800f800, ins.code),
+                    op => bail!("Unexpected Thumb opcode {op:?} for function call"),
+                },
+                Ins::Data => bail!("Unexpected data word for function call"),
+            };
+            let ins_pattern = code & ins_bitmask;
+            bitmask[offset..offset + 4].copy_from_slice(&ins_bitmask.to_le_bytes());
+            pattern[offset..offset + 4].copy_from_slice(&ins_pattern.to_le_bytes());
+
+            let info = get_relocation(address, Some(call.ins)).with_context(|| {
+                format!("No relocation found for function call at address {:#x}", address)
+            })?;
+            let Some(SignatureRelocationInfo { name, kind, addend }) = info else { continue };
+            relocations.push(SignatureRelocation { offset, name, kind, addend });
+        }
+
+        for &address in function.pool_constants() {
+            let offset = (address - function.start_address()) as usize;
+            bitmask[offset..offset + 4].fill(0);
+            pattern[offset..offset + 4].fill(0);
+
+            let info = get_relocation(address, None).with_context(|| {
+                format!("No relocation found for pool constant at address {:#x}", address)
+            })?;
+            let Some(SignatureRelocationInfo { name, kind, addend }) = info else {
+                continue;
+            };
+            relocations.push(SignatureRelocation { offset, name, kind, addend });
+        }
+
+        Ok(Self {
+            name: function.name().to_string(),
+            signatures: vec![Signature { mask: SignatureMask { bitmask, pattern }, relocations }],
+        })
+    }
+
     pub fn name(&self) -> &str {
         &self.name
     }

cli/src/cmd/sig/mod.rs

@@ -1,11 +1,13 @@
 mod apply;
 mod list;
 mod new;
+mod new_elf;
 
 pub use apply::*;
 use clap::{Args, Subcommand};
 pub use list::*;
 pub use new::*;
+pub use new_elf::*;
 
 /// Subcommands for creating/applying signatures.
 #[derive(Args)]
@@ -20,6 +22,7 @@ impl SigArgs {
             SigCommand::New(new) => new.run(),
             SigCommand::Apply(apply) => apply.run(),
             SigCommand::List(list) => list.run(),
+            SigCommand::NewElf(new_elf) => new_elf.run(),
         }
     }
 }
@@ -29,4 +32,5 @@ enum SigCommand {
     New(NewSignature),
     Apply(ApplySignature),
     List(ListSignatures),
+    NewElf(NewElfSignature),
 }

cli/src/cmd/sig/new_elf.rs

@@ -0,0 +1,117 @@
+use std::{collections::BTreeMap, path::PathBuf};
+
+use anyhow::{Context, Result, bail};
+use clap::Args;
+use ds_decomp::{
+    self,
+    analysis::functions::{Function, FunctionParseOptions, ParseFunctionOptions},
+    config::relocations::RelocationKind,
+};
+use object::{
+    LittleEndian, Object, ObjectSection, ObjectSymbol, RelocationFlags, RelocationTarget,
+};
+use unarm::arm;
+
+use crate::{
+    analysis::signature::{SignatureRelocationInfo, Signatures},
+    config::relocation::RelocationKindExt,
+    util::io,
+};
+
+#[derive(Args)]
+pub struct NewElfSignature {
+    /// ELF to extract function from
+    #[arg(long, short = 'i')]
+    input: PathBuf,
+
+    /// Function name to create the signature for.
+    #[arg(long, short = 'f')]
+    function: String,
+}
+
+impl NewElfSignature {
+    pub fn run(&self) -> Result<()> {
+        let object_bytes = io::read_file(&self.input)?;
+        let object = object::read::elf::ElfFile32::<LittleEndian>::parse(object_bytes.as_slice())?;
+        let function_symbol = object
+            .symbol_by_name(&self.function)
+            .with_context(|| format!("Symbol '{}' not found", self.function))?;
+        let function_section_index = function_symbol.section_index().unwrap();
+        let function_section = object.section_by_index(function_section_index).unwrap();
+        let section_data = function_section.uncompressed_data()?;
+        let start_address = function_symbol.address() as u32;
+        let end_address = (function_symbol.address() + function_symbol.size()) as u32;
+        let function_code = &section_data[start_address as usize..end_address as usize];
+
+        let function = Function::parse_function(FunctionParseOptions {
+            name: self.function.clone(),
+            start_address: start_address as u32,
+            base_address: start_address as u32,
+            module_code: function_code,
+            known_end_address: Some(end_address as u32),
+            module_start_address: start_address as u32,
+            module_end_address: end_address as u32,
+            existing_functions: None,
+            check_defs_uses: false,
+            parse_options: ParseFunctionOptions { thumb: None },
+        })?;
+
+        let relocations = function_section
+            .relocations()
+            .map(|(offset, reloc)| (offset as u32, reloc))
+            .filter(|(offset, _)| *offset >= start_address && *offset < end_address)
+            .collect::<BTreeMap<_, _>>();
+
+        let signature = Signatures::from_function_raw(&function, function_code, |address, ins| {
+            let Some(relocation) = relocations.get(&address) else {
+                if ins.is_some() {
+                    bail!("Relocation not found for instruction");
+                } else {
+                    // It's just data, could be a large constant or a link-time constant
+                    return Ok(None);
+                }
+            };
+            let target_symbol = match relocation.target() {
+                RelocationTarget::Symbol(symbol_index) => {
+                    object.symbol_by_index(symbol_index).unwrap()
+                }
+                target => bail!("Invalid relocation target {target:?}"),
+            };
+            let name = target_symbol.name().unwrap().to_string();
+            let kind = match relocation.flags() {
+                RelocationFlags::Elf { r_type } => {
+                    // FIXME: mwcc-generated ELFs do not seem to set bit zero to 1 for Thumb
+                    // functions, so this ends up being `false` all the time.
+                    let dest_thumb = (target_symbol.address() & 1) != 0;
+                    // `ins` may be `None` if the relocation is R_ARM_ABS32
+                    // `is_branch` only applies to ARM tail calls as there are no recorded instances
+                    // of Thumb branches used for tail calls
+                    let is_branch = match ins {
+                        Some(unarm::Ins::Arm(ins)) => ins.op == arm::Opcode::B,
+                        _ => false,
+                    };
+                    RelocationKind::from_elf_relocation_type(r_type, dest_thumb, is_branch).unwrap()
+                }
+                flags => bail!("Invalid relocation flags {flags:?}"),
+            };
+            let addend = match kind {
+                RelocationKind::ArmCall | RelocationKind::ArmCallThumb => {
+                    relocation.addend() as i32 + 8
+                }
+                RelocationKind::ThumbCall | RelocationKind::ThumbCallArm => {
+                    relocation.addend() as i32 + 4
+                }
+                RelocationKind::ArmBranch => relocation.addend() as i32,
+                RelocationKind::Load
+                | RelocationKind::OverlayId
+                | RelocationKind::LinkTimeConst(_) => relocation.addend() as i32,
+            };
+            Ok(Some(SignatureRelocationInfo { name, kind, addend }))
+        })?;
+
+        let signature_yaml = serde_saphyr::to_string(&signature)?;
+        print!("{signature_yaml}");
+
+        Ok(())
+    }
+}

cli/src/config/relocation.rs

@@ -9,9 +9,10 @@ use ds_decomp::config::{
 use ds_rom::rom::raw::AutoloadKind;
 use object::elf::{R_ARM_ABS32, R_ARM_PC24, R_ARM_THM_PC22};
 
-pub trait RelocationKindExt {
+pub trait RelocationKindExt: Sized {
     fn as_obj_symbol_kind(&self) -> object::SymbolKind;
     fn as_elf_relocation_type(&self) -> u32;
+    fn from_elf_relocation_type(r_type: u32, dest_thumb: bool, is_branch: bool) -> Option<Self>;
 }
 
 impl RelocationKindExt for RelocationKind {
@@ -42,6 +43,29 @@ impl RelocationKindExt for RelocationKind {
             Self::LinkTimeConst(_) => R_ARM_ABS32,
         }
     }
+
+    fn from_elf_relocation_type(r_type: u32, dest_thumb: bool, is_branch: bool) -> Option<Self> {
+        match r_type {
+            R_ARM_PC24 => {
+                if is_branch {
+                    Some(Self::ArmBranch)
+                } else if dest_thumb {
+                    Some(Self::ArmCallThumb)
+                } else {
+                    Some(Self::ArmCall)
+                }
+            }
+            R_ARM_THM_PC22 => {
+                if dest_thumb {
+                    Some(Self::ThumbCall)
+                } else {
+                    Some(Self::ThumbCallArm)
+                }
+            }
+            R_ARM_ABS32 => Some(Self::Load),
+            _ => None,
+        }
+    }
 }
 
 pub trait RelocationModuleExt