r2fas 0.2.1

//! Apply [`crate::session::DebugInfo`] to an open radare2 core.

use crate::session::{DebugInfo, SourceExtent, SymbolKind as FasSymbolKind};
use radare2::Core;
use radare2::native::{
    AddrLine as NativeAddrLine, AddrLineSnapshot, BinIdentity, FunctionAdd, FunctionTransaction,
    MapIdentity, Symbol as NativeSymbol, SymbolAdd, SymbolKind as NativeSymbolKind,
    SymbolTransaction, XrefAdd, XrefKind, XrefTransaction, replace_addrlines, resolve_paddr,
    resolve_vaddr,
};
use std::collections::BTreeSet;
use std::path::{Path, PathBuf};

/// True when active IO maps can represent at least one emitted FAS address.
pub fn maps_ready(core: Core, info: &DebugInfo) -> bool {
    let mut resolver = AddressResolver::new(core, info);
    info.source_locations
        .iter()
        .filter(|location| location.primary)
        .any(|location| {
            resolver
                .resolve(
                    location.address,
                    Some(u64::from(location.output_offset)),
                    false,
                )
                .is_some()
        })
        || info
            .symbols
            .iter()
            .filter(|symbol| symbol.debugger_visible)
            .any(|symbol| {
                resolver
                    .resolve(
                        symbol.value,
                        symbol.output_offset.map(u64::from),
                        symbol.kind == FasSymbolKind::Code,
                    )
                    .is_some()
            })
}

/// True when `om` lists a vaddr range that contains `addr`.
///
/// Kept as a parser helper for compatibility and focused tests. Runtime address
/// resolution itself uses header-backed `RIOMap` APIs.
pub fn maps_cover_addr(om: &str, addr: u64) -> bool {
    om.lines()
        .any(|line| parse_om_vaddr_range(line).is_some_and(|(from, to)| addr >= from && addr <= to))
}

/// Parse the `0xFROM - 0xTO` vaddr span from one `om` line.
fn parse_om_vaddr_range(line: &str) -> Option<(u64, u64)> {
    let dash = line.find(" - 0x")?;
    let from_hex = line[..dash]
        .rsplit("0x")
        .next()?
        .chars()
        .take_while(|c| c.is_ascii_hexdigit())
        .collect::<String>();
    let to_hex = line[dash + 5..]
        .chars()
        .take_while(|c| c.is_ascii_hexdigit())
        .collect::<String>();
    let from = u64::from_str_radix(&from_hex, 16).ok()?;
    let to = u64::from_str_radix(&to_hex, 16).ok()?;
    Some((from, to))
}

#[derive(Debug, Clone, Copy)]
struct RuntimeAddress {
    addr: u64,
}

struct AddressResolver {
    core: Core,
    original_baddr: u64,
    current_baddr: u64,
    debugger: bool,
    maps: BTreeSet<MapIdentity>,
}

impl AddressResolver {
    fn new(core: Core, info: &DebugInfo) -> Self {
        Self {
            core,
            original_baddr: info.original_baddr,
            current_baddr: core.baddr(),
            debugger: radare2::io_uri::ptrace_pid(core.cmd_str("o.").trim()).is_some(),
            maps: BTreeSet::new(),
        }
    }

    fn rebased(&self, original: u64) -> u64 {
        if original != 0 && self.current_baddr != 0 && self.original_baddr != 0 {
            original.wrapping_add(self.current_baddr.wrapping_sub(self.original_baddr))
        } else {
            original
        }
    }

    fn accept(
        &mut self,
        addr: u64,
        map: MapIdentity,
        require_executable: bool,
    ) -> Option<RuntimeAddress> {
        if require_executable && !map.is_executable() {
            return None;
        }
        self.maps.insert(map);
        Some(RuntimeAddress { addr })
    }

    fn resolve(
        &mut self,
        original: u64,
        paddr: Option<u64>,
        require_executable: bool,
    ) -> Option<RuntimeAddress> {
        let rebased = self.rebased(original);

        if let Some(paddr) = paddr
            && let Some((mapped, map)) = resolve_paddr(self.core, paddr)
            && (self.original_baddr == 0 || mapped == original || mapped == rebased)
            && let Some(address) = self.accept(mapped, map, require_executable)
        {
            return Some(address);
        }

        for candidate in [rebased, original] {
            if candidate == 0 {
                continue;
            }
            let Some((mapped_paddr, map)) = resolve_vaddr(self.core, candidate) else {
                continue;
            };
            if paddr.is_some_and(|expected| expected != mapped_paddr) && !self.debugger {
                continue;
            }
            if let Some(address) = self.accept(candidate, map, require_executable) {
                return Some(address);
            }
        }
        None
    }

    fn identities(&self) -> Vec<MapIdentity> {
        self.maps.iter().copied().collect()
    }
}

/// Result of a successful apply, used for log lines and unload.
#[derive(Debug, Clone)]
pub struct Applied {
    /// Flag names that were created (for `fas unload`).
    pub flags: Vec<String>,
    /// How many native binary symbols were inserted.
    pub native_symbol_count: usize,
    /// How many native addrline entries were written.
    pub line_count: usize,
    /// How many real analyzed functions were newly created.
    pub function_count: usize,
    /// How many conservative FAS reference xrefs were newly inserted.
    pub xref_count: usize,
    /// Path of the dump that was applied.
    pub fas_path: String,
    /// Binary object that owns all native mutations.
    pub identity: BinIdentity,
    /// IO maps used to resolve runtime metadata.
    pub maps: Vec<MapIdentity>,
    symbol_transaction: SymbolTransaction,
    addrline_snapshot: AddrLineSnapshot,
    function_transaction: FunctionTransaction,
    xref_transaction: XrefTransaction,
}

impl Applied {
    /// Return whether this transaction still targets the current bin and maps.
    pub fn target_is_current(&self, core: Core) -> bool {
        self.identity.is_current(core) && self.maps.iter().all(|map| map.is_current(core))
    }
}

/// Failure while applying native metadata to radare2.
#[derive(Debug, Clone)]
pub struct ApplyError {
    message: String,
}

impl ApplyError {
    fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }
}

impl std::fmt::Display for ApplyError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(&self.message)
    }
}

impl std::error::Error for ApplyError {}

/// Load FAS debug info into `core`.
pub fn apply(core: Core, info: &DebugInfo) -> Result<Applied, ApplyError> {
    let identity = BinIdentity::current(core)
        .ok_or_else(|| ApplyError::new("fas: no current binary object for native metadata"))?;
    let mut symbol_transaction = SymbolTransaction::begin(core, identity)
        .ok_or_else(|| ApplyError::new("fas: cannot access the current binary symbol vector"))?;
    let addrline_snapshot = AddrLineSnapshot::capture(core, identity)
        .ok_or_else(|| ApplyError::new("fas: cannot snapshot current source lines"))?;

    for symbol in native_symbols(info) {
        match symbol_transaction.add(core, &symbol) {
            SymbolAdd::Added | SymbolAdd::Duplicate => {}
            SymbolAdd::Failed => {
                let _ = symbol_transaction.rollback(core);
                return Err(ApplyError::new(format!(
                    "fas: failed to insert native symbol {}",
                    symbol.name
                )));
            }
        }
    }

    let mut resolver = AddressResolver::new(core, info);
    let resolved_lines = resolved_addrlines(info, &mut resolver);
    if !replace_addrlines(core, identity, &resolved_lines) {
        let _ = symbol_transaction.rollback(core);
        return Err(ApplyError::new(
            "fas: failed to replace native source lines",
        ));
    }

    let want_comments = core.cfg_bool("fas.comments", true);
    let want_analyze = core.cfg_bool("fas.analyze", true);
    let mut script = String::with_capacity(4096);
    script.push_str("e asm.dwarf=true\nfs fas\n");

    let mut flags = Vec::new();
    for symbol in info.symbols.iter().filter(|symbol| symbol.debugger_visible) {
        let Some(name) = symbol.alias.as_deref().map(sanitize_flag) else {
            continue;
        };
        if name.is_empty() {
            continue;
        }
        let Some(runtime) = resolver.resolve(
            symbol.value,
            symbol.output_offset.map(u64::from),
            symbol.kind == FasSymbolKind::Code,
        ) else {
            continue;
        };
        let size = u64::from(symbol.size.max(1));
        script.push_str(&format!("f {name} {size} @ 0x{:x}\n", runtime.addr));
        if symbol.kind == FasSymbolKind::Data && symbol.size > 0 {
            script.push_str(&format!("Cd {} @ 0x{:x}\n", symbol.size, runtime.addr));
        }
        flags.push(name);
    }
    script.push_str("fs *\n");

    if want_comments {
        for location in info
            .source_locations
            .iter()
            .filter(|location| location.primary)
        {
            let Some(text) = location_comment_text(location) else {
                continue;
            };
            if let Some(runtime) = resolver.resolve(
                location.address,
                Some(u64::from(location.output_offset)),
                false,
            ) {
                script.push_str(&format!("CCu {text} @ 0x{:x}\n", runtime.addr));
            }
        }
    }
    core.cmd_lines(&script);

    let mut xref_transaction = XrefTransaction::begin();
    if let Err(error) = apply_xrefs(core, info, &mut resolver, &mut xref_transaction) {
        let _ = xref_transaction.rollback(core);
        let _ = addrline_snapshot.restore(core);
        let _ = symbol_transaction.rollback(core);
        remove_flags(core, &flags);
        return Err(error);
    }

    let mut function_transaction = FunctionTransaction::begin();
    if want_analyze {
        analyze_functions(core, info, &mut resolver, &mut function_transaction);
    }

    Ok(Applied {
        flags,
        native_symbol_count: symbol_transaction.added(),
        line_count: resolved_lines.len(),
        function_count: function_transaction.added(),
        xref_count: xref_transaction.added(),
        fas_path: info.fas_path.display().to_string(),
        identity,
        maps: resolver.identities(),
        symbol_transaction,
        addrline_snapshot,
        function_transaction,
        xref_transaction,
    })
}

/// Remove native and analysis metadata created by this plugin.
pub fn unload(core: Core, applied: &Applied) -> bool {
    if !applied.target_is_current(core)
        || !applied.xref_transaction.can_rollback(core)
        || !applied.function_transaction.can_rollback(core)
        || !applied.symbol_transaction.can_rollback(core)
    {
        return false;
    }

    if !applied.xref_transaction.rollback(core)
        || !applied.function_transaction.rollback(core)
        || !applied.addrline_snapshot.restore(core)
        || !applied.symbol_transaction.rollback(core)
    {
        return false;
    }
    remove_flags(core, &applied.flags);
    true
}

fn remove_flags(core: Core, flags: &[String]) {
    core.cmd("fs fas");
    for name in flags {
        core.cmd(&format!("f- {name}"));
    }
    core.cmd("fs *");
}

fn native_symbols(info: &DebugInfo) -> Vec<NativeSymbol> {
    info.symbols
        .iter()
        .filter(|symbol| symbol.debugger_visible)
        .filter_map(|symbol| {
            let name = symbol.original_name.clone()?;
            let kind = match symbol.kind {
                FasSymbolKind::Data => NativeSymbolKind::Object,
                FasSymbolKind::Code => NativeSymbolKind::Function,
                _ => NativeSymbolKind::NoType,
            };
            Some(NativeSymbol {
                name,
                paddr: symbol.output_offset.map(u64::from),
                vaddr: symbol.value,
                size: u32::from(symbol.size),
                ordinal: symbol.id.0.try_into().unwrap_or(u32::MAX),
                kind,
            })
        })
        .collect()
}

fn resolved_addrlines(info: &DebugInfo, resolver: &mut AddressResolver) -> Vec<NativeAddrLine> {
    info.source_locations
        .iter()
        .filter(|location| location.primary)
        .filter_map(|location| {
            let runtime = resolver.resolve(
                location.address,
                Some(u64::from(location.output_offset)),
                false,
            )?;
            let source = resolve_source(&info.source_dir, &location.file);
            Some(NativeAddrLine {
                addr: runtime.addr,
                file: source.to_string_lossy().into_owned(),
                path: None,
                line: location.line,
                column: 0,
            })
        })
        .collect()
}

fn analyze_functions(
    core: Core,
    info: &DebugInfo,
    resolver: &mut AddressResolver,
    transaction: &mut FunctionTransaction,
) {
    for symbol in info.symbols.iter().filter(|symbol| {
        symbol.debugger_visible
            && symbol.kind == FasSymbolKind::Code
            && symbol.output_offset.is_some()
    }) {
        let Some(name) = symbol.alias.as_deref().map(sanitize_flag) else {
            continue;
        };
        if name.is_empty() {
            continue;
        }
        let Some(runtime) =
            resolver.resolve(symbol.value, symbol.output_offset.map(u64::from), true)
        else {
            continue;
        };
        match transaction.add(core, runtime.addr, &name) {
            FunctionAdd::Added | FunctionAdd::Existing | FunctionAdd::Failed => {}
        }
    }
}

fn apply_xrefs(
    core: Core,
    info: &DebugInfo,
    resolver: &mut AddressResolver,
    transaction: &mut XrefTransaction,
) -> Result<(), ApplyError> {
    let mut seen = BTreeSet::new();
    for reference in &info.references {
        let Some(symbol) = info.symbols.get(reference.symbol.0) else {
            continue;
        };
        if !symbol.debugger_visible
            || matches!(
                symbol.kind,
                FasSymbolKind::Constant
                    | FasSymbolKind::External
                    | FasSymbolKind::Marker
                    | FasSymbolKind::Anonymous
            )
        {
            continue;
        }
        let Some(location) = info
            .source_locations
            .iter()
            .find(|location| location.row == reference.row)
        else {
            continue;
        };
        if !matches!(location.extent, SourceExtent::Emitted(_)) {
            continue;
        }
        let Some(from) = resolver.resolve(
            location.address,
            Some(u64::from(location.output_offset)),
            false,
        ) else {
            continue;
        };
        let Some(to) = resolver.resolve(symbol.value, symbol.output_offset.map(u64::from), false)
        else {
            continue;
        };
        if from.addr == to.addr || !seen.insert((from.addr, to.addr)) {
            continue;
        }
        if transaction.add(core, from.addr, to.addr, XrefKind::Data) == XrefAdd::Failed {
            return Err(ApplyError::new(format!(
                "fas: failed to insert xref 0x{:x} -> 0x{:x}",
                from.addr, to.addr
            )));
        }
    }
    Ok(())
}

/// Make a string safe as an r2 flag / command name.
pub fn sanitize_flag(name: &str) -> String {
    let mut output: String = name
        .chars()
        .map(|character| {
            if character.is_ascii_alphanumeric() || matches!(character, '_' | '.' | ':') {
                character
            } else {
                '_'
            }
        })
        .collect();
    if output
        .chars()
        .next()
        .is_some_and(|character| character.is_ascii_digit())
    {
        output.insert(0, '_');
    }
    output
}

fn resolve_source(directory: &Path, file: &str) -> PathBuf {
    let path = directory.join(file);
    if path.exists() {
        path
    } else {
        PathBuf::from(file)
    }
}

fn location_comment_text(location: &crate::session::SourceLocation) -> Option<String> {
    let macro_name = location
        .macro_frames
        .first()
        .and_then(|frame| frame.macro_name.as_deref())?;
    let source = location.text.as_deref().unwrap_or("").trim();
    let text = if source.is_empty() {
        format!("macro {macro_name}")
    } else {
        format!("macro {macro_name}: {source}")
    };
    Some(sanitize_comment(&text))
}

fn sanitize_comment(value: &str) -> String {
    value
        .chars()
        .map(|character| match character {
            '\n' | '\r' | ';' | '@' | '"' | '\'' => ' ',
            character => character,
        })
        .collect::<String>()
        .split_whitespace()
        .collect::<Vec<_>>()
        .join(" ")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn om_range_covers_elf_load_maps_not_file_offsets() {
        let om = "\
* 3 fd: 3 +0x00000000 0x00400000 - 0x00400266 r-x fmap.LOAD0
- 2 fd: 3 +0x00000267 0x00401267 - 0x00401337 r-- fmap.LOAD1
";
        assert!(maps_cover_addr(om, 0x400112));
        assert!(!maps_cover_addr(om, 0x111));
        assert!(!maps_cover_addr("", 0x400112));
        assert_eq!(
            parse_om_vaddr_range("* 3 fd: 3 +0x00000000 0x00400000 - 0x00400266 r-x fmap.LOAD0"),
            Some((0x400000, 0x400266))
        );
    }

    #[test]
    fn om_range_covers_ptrace_whole_as() {
        let om = "* 3 fd: 4 +0x00000000 0x00000000 - 0xffffffffffffffff rwx dbg.ptrace";
        assert!(maps_cover_addr(om, 0x400112));
    }

    #[test]
    fn sanitizers_preserve_names_and_neutralize_commands() {
        assert_eq!(sanitize_flag("12 bad-name"), "_12_bad_name");
        assert_eq!(
            sanitize_comment("macro x: a; b @ \"c\"\n"),
            "macro x: a b c"
        );
    }
}