use crate::container::Container;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ShimKind {
NimMain,
NimMainInner,
PreMain,
PreMainInner,
NimMainModule,
}
#[derive(Debug, Clone)]
pub struct EntryShim<'a> {
pub kind: ShimKind,
pub symbol_name: &'a str,
pub address: u64,
}
const SHIM_NAMES: &[(&str, ShimKind)] = &[
("NimMain", ShimKind::NimMain),
("NimMainInner", ShimKind::NimMainInner),
("NimMainModule", ShimKind::NimMainModule),
("PreMain", ShimKind::PreMain),
("PreMainInner", ShimKind::PreMainInner),
];
pub fn scan<'a>(container: &'a Container<'a>) -> Vec<EntryShim<'a>> {
let mut result = Vec::new();
for sym in container.symbols() {
let name = sym.name.as_ref();
let stripped = name.strip_prefix('_').unwrap_or(name);
for &(suffix, kind) in SHIM_NAMES {
if stripped == suffix || stripped.ends_with(suffix) && is_prefix_match(stripped, suffix)
{
result.push(EntryShim {
kind,
symbol_name: name,
address: sym.vm_addr,
});
break;
}
}
}
result
}
fn is_prefix_match(symbol: &str, suffix: &str) -> bool {
if symbol.len() <= suffix.len() {
return false;
}
let prefix = &symbol[..symbol.len() - suffix.len()];
!prefix.contains("__")
}
pub fn detect_prefix<'a>(shims: &'a [EntryShim<'a>]) -> Option<&'a str> {
let nim_main = shims.iter().find(|s| s.kind == ShimKind::NimMain)?;
let name = nim_main.symbol_name;
let stripped = name.strip_prefix('_').unwrap_or(name);
let prefix = stripped.strip_suffix("NimMain")?;
Some(prefix)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn shim_kind_eq() {
assert_eq!(ShimKind::NimMain, ShimKind::NimMain);
assert_ne!(ShimKind::NimMain, ShimKind::PreMain);
}
#[test]
fn prefix_match_rejects_mangled_symbols() {
assert!(!is_prefix_match("something__NimMain", "NimMain"));
}
#[test]
fn prefix_match_accepts_custom_prefix() {
assert!(is_prefix_match("MyLibNimMain", "NimMain"));
assert!(is_prefix_match("AppNimMainInner", "NimMainInner"));
}
#[test]
fn prefix_match_rejects_exact() {
assert!(!is_prefix_match("NimMain", "NimMain"));
}
}