viking: Check calls to unimplemented functions

viking/src/checks.rs

@@ -131,6 +131,7 @@ pub enum MismatchCause {
     DataReference(ReferenceDiff),
     Immediate,
     Unknown,
+    InternalError(String),
 }
 
 impl std::fmt::Display for MismatchCause {
@@ -144,6 +145,7 @@ impl std::fmt::Display for MismatchCause {
             Self::DataReference(diff) => write!(f, "wrong data reference\n{diff}"),
             Self::Immediate => write!(f, "wrong immediate"),
             Self::Unknown => write!(f, "unknown reason"),
+            Self::InternalError(msg) => write!(f, "internal error: {}", msg),
         }
     }
 }
@@ -415,26 +417,67 @@ impl<'a, 'functions, 'orig_elf, 'decomp_elf>
         Ok(None)
     }
 
+    fn plt_name_to_addr(elf: &elf::OwnedElf, func_name: &str) -> Option<u64> {
+        // find plt relocation for given function name
+        let (reloc_idx, _) = elf.pltrelocs.iter().enumerate().find(|(_, reloc)| {
+            elf.dynsyms.get(reloc.r_sym).map(|s| elf.dynstrtab.get_at(s.st_name)).flatten() == Some(func_name)
+        })?;
+
+        let plt_section = elf::find_section(elf, ".plt").ok()?;
+        Some(plt_section.sh_addr + (reloc_idx as u64 + 2) * 0x10)
+    }
+
+    fn plt_addr_to_name(elf: &elf::OwnedElf, addr: u64) -> Option<&str> {
+        let plt_section = elf::find_section(elf, ".plt").ok()?;
+        let reloc_idx = (addr - plt_section.sh_addr) / 16 - 2;
+        let reloc = elf.pltrelocs.get(reloc_idx as usize)?;
+        let name = elf.dynstrtab.get_at(elf.dynsyms.get(reloc.r_sym)?.st_name)?;
+        Some(name)
+    }
+
     /// Returns None on success and a MismatchCause on failure.
     fn check_function_call(&self, orig_addr: u64, decomp_addr: u64) -> Option<MismatchCause> {
-        let info = *self.known_functions.get(&(orig_addr as u32))?;
+        let Some(info) = self.known_functions.get(&(orig_addr as u32)) else {
+            return Some(MismatchCause::InternalError(
+                format!("failed to resolve orig function at address {:x}", orig_addr)
+            ))
+        };
         let name = info.name();
-        let decomp_symbol = self.decomp_symtab.get(name)?;
-        let expected = decomp_symbol.st_value;
-
-        if decomp_addr == expected {
-            None
-        } else {
+        if name.is_empty() {
+            // called some function that is not named in the binary or function list
             let actual_symbol_name = self.translate_decomp_addr_to_name(decomp_addr);
+            return Some(MismatchCause::FunctionCall(ReferenceDiff {
+                referenced_symbol: orig_addr,
+                expected_ref_in_decomp: 0,
+                actual_ref_in_decomp: decomp_addr,
+                expected_symbol_name: "<unnamed function>".to_string(),
+                actual_symbol_name: actual_symbol_name.unwrap_or("unknown").to_string(),
+            }))
+        }
+        let Some(expected) = self.decomp_symtab.get(name).map(|sym| sym.st_value)
+                .or_else(|| Self::plt_name_to_addr(self.decomp_elf, name)) else {
+            let actual_symbol_name = self.translate_decomp_addr_to_name(decomp_addr);
+            return Some(MismatchCause::FunctionCall(ReferenceDiff {
+                referenced_symbol: orig_addr,
+                expected_ref_in_decomp: 0,
+                actual_ref_in_decomp: decomp_addr,
+                expected_symbol_name: name.to_string(),
+                actual_symbol_name: actual_symbol_name.unwrap_or("unknown").to_string(),
+            }))
+        };
 
-            Some(MismatchCause::FunctionCall(ReferenceDiff {
+        if decomp_addr != expected {
+            let actual_symbol_name = self.translate_decomp_addr_to_name(decomp_addr);
+            return Some(MismatchCause::FunctionCall(ReferenceDiff {
                 referenced_symbol: orig_addr,
                 expected_ref_in_decomp: expected,
                 actual_ref_in_decomp: decomp_addr,
                 expected_symbol_name: name.to_string(),
                 actual_symbol_name: actual_symbol_name.unwrap_or("unknown").to_string(),
             }))
         }
+
+        None
     }
 
     /// Returns None on success and a MismatchCause on failure.
@@ -511,6 +554,6 @@ impl<'a, 'functions, 'orig_elf, 'decomp_elf>
             let map = elf::make_addr_to_name_map(self.decomp_elf).ok();
             map.unwrap_or_default()
         });
-        map.get(&decomp_addr).copied()
+        map.get(&decomp_addr).copied().or_else(|| Self::plt_addr_to_name(self.decomp_elf, decomp_addr))
     }
 }

viking/src/elf.rs

@@ -22,7 +22,7 @@ use memmap::{Mmap, MmapOptions};
 use owning_ref::OwningHandle;
 use rustc_hash::FxHashMap;
 
-use crate::repo;
+use crate::{repo, functions::{AddressLabel, Info, Status}};
 
 pub type OwnedElf = OwningHandle<Box<(File, Mmap)>, Box<Elf<'static>>>;
 pub type SymbolTableByName<'a> = HashMap<&'a str, goblin::elf::Sym>;
@@ -55,7 +55,7 @@ fn make_goblin_ctx() -> container::Ctx {
 
 /// A stripped down version of `goblin::elf::Elf::parse`, parsing only the sections that we need.
 ///
-/// *Warning*: In particular, `strtab`, `dynstrtab`, `soname` and `libraries` are **not** parsed.
+/// *Warning*: In particular, `strtab`, `soname` and `libraries` are **not** parsed.
 fn parse_elf_faster(bytes: &[u8]) -> Result<Elf<'_>> {
     let header = Elf::parse_header(bytes)?;
     let mut elf = Elf::lazy_parse(header)?;
@@ -97,6 +97,12 @@ fn parse_elf_faster(bytes: &[u8]) -> Result<Elf<'_>> {
         let is_rela = dyn_info.pltrel == dynamic::DT_RELA;
         elf.pltrelocs =
             RelocSection::parse(bytes, dyn_info.jmprel, dyn_info.pltrelsz, is_rela, ctx)?;
+
+        let hash_offset = dyn_info.hash.context("no hash")? as usize;
+        let num_syms = u32::from_le_bytes(bytes[hash_offset+4..hash_offset+8].try_into()?) as usize;
+        elf.dynsyms = Symtab::parse(bytes, dyn_info.symtab, num_syms, ctx)?;
+
+        elf.dynstrtab = Strtab::parse(bytes, dyn_info.strtab, dyn_info.strsz, 0x0)?;
     }
 
     Ok(elf)
@@ -296,6 +302,24 @@ pub fn build_glob_data_table(elf: &OwnedElf) -> Result<GlobDataTable> {
     Ok(table)
 }
 
+pub fn get_plt_functions(elf: &OwnedElf) -> Result<Vec<crate::functions::Info>> {
+    let plt_section = find_section(elf, ".plt")?;
+    let mut functions = Vec::with_capacity(elf.pltrelocs.len());
+
+    for (idx, reloc) in elf.pltrelocs.iter().enumerate() {
+        // each PLT entry is 0x10 bytes, and the first (reserved) entry is twice the size
+        let addr = plt_section.sh_addr + (idx as u64 + 2) * 0x10;
+        let sym = elf.dynsyms.get(reloc.r_sym).context("Failed to get dynsym")?;
+        let name = elf.dynstrtab.get_at(sym.st_name).context("Failed to get dynstr")?;
+        functions.push(Info {
+            offset: addr as u32, size: 0x10, label: AddressLabel::Single(name.to_string()), 
+            status: Status::Library, lazy: false, guess: false,
+        });
+    }
+
+    Ok(functions)
+}
+
 pub fn get_offset_in_file(elf: &OwnedElf, addr: u64) -> Result<usize> {
     let addr = addr as usize;
     for segment in elf.program_headers.iter() {

viking/src/functions.rs

@@ -134,9 +134,6 @@ pub fn make_known_function_map(functions: &[Info]) -> FxHashMap<u32, &Info> {
         FxHashMap::with_capacity_and_hasher(functions.len(), Default::default());
 
     for function in functions {
-        if function.name().is_empty() {
-            continue;
-        }
         known_functions.insert(function.offset, function);
     }
 

viking/src/tools/check.rs

@@ -67,6 +67,7 @@ fn main() -> Result<()> {
     let mut decomp_symtab = None;
     let mut decomp_glob_data_table = None;
     let mut file_list = None;
+    let mut plt_functions = None;
 
     rayon::scope(|s| {
         s.spawn(|_| decomp_symtab = Some(elf::make_symbol_map_by_name(&decomp_elf)));
@@ -76,19 +77,20 @@ fn main() -> Result<()> {
                 get_file_list_path(version).as_path(),
             ));
         });
+        s.spawn(|_| plt_functions = Some(elf::get_plt_functions(&orig_elf)));
     });
 
     let decomp_symtab = decomp_symtab
         .unwrap()
         .context("failed to make symbol map")?;
-
     let decomp_glob_data_table = decomp_glob_data_table
         .unwrap()
         .context("failed to make global data table")?;
-
     let file_list = file_list.unwrap().context("failed to load file list")?;
+    let plt_functions = plt_functions.unwrap().context("failed to load plt functions")?;
 
-    let functions = functions::get_functions(&file_list);
+    let file_functions = functions::get_functions(&file_list);
+    let functions = vec![file_functions, plt_functions].concat();
 
     let checker = FunctionChecker::new(
         &orig_elf,