r2fas 0.2.0

//! radare2 core-plugin entry points: init, command dispatch, autoload.

use crate::apply::{self, Applied};
use crate::discover;
use crate::session::DebugInfo;
use radare2::Core;
use radare2::plugin::{self as r2plugin, CorePlugin, PluginSession};
use std::ffi::CStr;
use std::os::raw::c_char;
use std::path::PathBuf;

/// Per-core plugin state stored in [`PluginSession::data`].
pub struct PluginState {
    /// Last successful apply, if any.
    pub applied: Option<Applied>,
    /// Guard against `cmd.load` re-entering apply.
    pub busy: bool,
    /// Previous `cmd.load` / `cmd.open` / `cmd.prompt` values, for restore on fini.
    pub prev_hooks: Vec<(String, String)>,
}

/// Command hooks prefixed with `fas.` so autoload retries after debugger attach.
const CMD_HOOKS: &[&str] = &["cmd.load", "cmd.open", "cmd.prompt"];

/// Help text for `fas?`.
const HELP: &[&str] = &[
    "Usage: fas  load  info  unload",
    " fas            autoload a sibling .fas dump",
    " fas load path  load an explicit FAS dump",
    " fas info       summarize the last loaded dump",
    " fas unload     remove owned symbols, functions, xrefs, flags, and source lines",
    " fas?           this help",
];

/// Called once when r2 loads the plugin into a core.
///
/// # Safety
///
/// `session` must be the live plugin session radare2 passed into `init`.
pub unsafe extern "C" fn plugin_init(session: *mut PluginSession) -> bool {
    let Some(core) = (unsafe { Core::from_session(session) }) else {
        return false;
    };

    core.set_cfg_bool(
        "fas.autoload",
        true,
        "automatically load a sibling FASM .fas dump when a binary is opened",
    );
    core.set_cfg_bool(
        "fas.comments",
        true,
        "attach original FASM source lines as comments",
    );
    core.set_cfg_bool(
        "fas.analyze",
        true,
        "create functions at FASM code labels after loading symbols",
    );

    let mut prev_hooks = Vec::new();
    for key in CMD_HOOKS {
        let prev = core.cfg_str(key);
        wrap_cmd_hook(core, key, &prev);
        prev_hooks.push(((*key).to_string(), prev));
    }

    let state = PluginState {
        applied: None,
        busy: false,
        prev_hooks,
    };
    if unsafe { r2plugin::install_state(session, state) }.is_err() {
        return false;
    }

    // File may already be open if the plugin was loaded with `L`.
    autoload(session, true);
    true
}

/// Restore `cmd.load` and free session data.
///
/// # Safety
///
/// `session` must be the live plugin session radare2 passed into `fini`.
pub unsafe extern "C" fn plugin_fini(session: *mut PluginSession) -> bool {
    let Some(core) = (unsafe { Core::from_session(session) }) else {
        return false;
    };
    if let Some(mut state) = unsafe { r2plugin::take_state::<PluginState>(session) } {
        if let Some(applied) = state.applied.take() {
            let _ = apply::unload(core, &applied);
        }
        for (key, prev) in state.prev_hooks {
            core.set_cfg_str(&key, &prev);
        }
    }
    true
}

/// Handle commands that start with `fas`.
///
/// # Safety
///
/// `session` must be live. `input` must be a valid C string or null.
pub unsafe extern "C" fn plugin_call(session: *mut PluginSession, input: *const c_char) -> bool {
    if input.is_null() {
        return false;
    }
    let raw = unsafe { CStr::from_ptr(input).to_string_lossy() };
    let Some(rest) = strip_command(&raw) else {
        return false;
    };
    dispatch(session, rest.trim());
    true
}

/// Return the suffix after `fas` if this command belongs to us.
fn strip_command(input: &str) -> Option<&str> {
    let input = input.trim_start();
    let rest = input.strip_prefix("fas")?;
    match rest.chars().next() {
        None | Some(' ') | Some('.') | Some('?') => Some(rest),
        _ => None,
    }
}

/// True when no dump has been applied yet (or state is missing).
fn applied_is_none(session: *mut PluginSession) -> bool {
    unsafe { r2plugin::state::<PluginState>(session) }.is_none_or(|state| state.applied.is_none())
}

/// Interpret `fas` subcommands.
fn dispatch(session: *mut PluginSession, rest: &str) {
    let Some(core) = (unsafe { Core::from_session(session) }) else {
        return;
    };
    let silent = rest.starts_with('.');
    let body = rest.trim_start_matches('.').trim();

    if body.is_empty() {
        autoload(session, silent);
        return;
    }
    if body == "?" || body.starts_with('?') {
        for line in HELP {
            core.echo(line);
        }
        return;
    }
    let mut parts = body.splitn(2, char::is_whitespace);
    let cmd = parts.next().unwrap_or("");
    let arg = parts.next().map(str::trim).filter(|s| !s.is_empty());
    match cmd {
        "load" => match arg {
            Some(path) => load_path(session, PathBuf::from(path), silent),
            None => autoload(session, silent),
        },
        "info" => print_info(core, session),
        "unload" => unload(session),
        _ => {
            if !silent {
                core.echo("unknown fas subcommand; try fas?");
            }
        }
    }
}

/// Prefix a command hook with `fas` without duplicating it.
fn wrap_cmd_hook(core: Core, key: &str, previous: &str) {
    let prev = previous.trim();
    if prev
        .split(';')
        .any(|p| p.trim() == "fas." || p.trim() == "fas")
    {
        return;
    }
    let wrapped = if prev.is_empty() {
        "fas.".to_string()
    } else {
        format!("fas.;{prev}")
    };
    core.set_cfg_str(key, &wrapped);
}

/// Discover a sibling dump and apply it.
fn autoload(session: *mut PluginSession, silent: bool) {
    if silent && !applied_is_none(session) {
        return;
    }
    let Some(core) = (unsafe { Core::from_session(session) }) else {
        return;
    };
    if !core.cfg_bool("fas.autoload", true) && silent {
        return;
    }
    let Some(file) = core.current_file() else {
        return;
    };
    match discover::find_for_binary(file.as_path()) {
        Some(path) => load_path(session, path, silent),
        None => {
            if !silent {
                core.echo("fas: no sibling .fas dump found");
            }
        }
    }
}

/// Parse `path` and apply it to the current core.
fn load_path(session: *mut PluginSession, path: PathBuf, silent: bool) {
    let Some(core) = (unsafe { Core::from_session(session) }) else {
        return;
    };
    let path_str = path.to_string_lossy().into_owned();
    let state = unsafe { r2plugin::state_mut::<PluginState>(session) };
    if let Some(state) = state {
        if state.busy {
            return;
        }
        if state
            .applied
            .as_ref()
            .is_some_and(|applied| applied.fas_path == path_str && applied.target_is_current(core))
        {
            return;
        }
        state.busy = true;
    }

    let info = match DebugInfo::from_path(&path) {
        Ok(info) => info,
        Err(error) => {
            if let Some(state) = unsafe { r2plugin::state_mut::<PluginState>(session) } {
                state.busy = false;
            }
            if !silent {
                core.echo(&format!("fas: {error}"));
            }
            return;
        }
    };
    if silent && !apply::maps_ready(core, &info) {
        if let Some(state) = unsafe { r2plugin::state_mut::<PluginState>(session) } {
            state.busy = false;
        }
        return;
    }

    let previous = unsafe { r2plugin::state_mut::<PluginState>(session) }
        .and_then(|state| state.applied.take());
    if let Some(previous) = previous {
        if !apply::unload(core, &previous) {
            if let Some(state) = unsafe { r2plugin::state_mut::<PluginState>(session) } {
                state.busy = false;
                state.applied = Some(previous);
            }
            if !silent {
                core.echo("fas: replacement refused because existing metadata was modified");
            }
            return;
        }
    }

    let result = apply::apply(core, &info).map_err(|error| error.to_string());

    if let Some(state) = unsafe {
        session
            .as_mut()
            .and_then(|s| (s.data as *mut PluginState).as_mut())
    } {
        state.busy = false;
        if let Ok(applied) = &result {
            state.applied = Some(applied.clone());
        }
    }

    match result {
        Ok(applied) => core.echo(&format!(
            "fas: loaded {} symbols, {} lines, {} functions, {} xrefs from {}",
            applied.flags.len(),
            applied.line_count,
            applied.function_count,
            applied.xref_count,
            applied.fas_path
        )),
        Err(e) => {
            if !silent {
                core.echo(&format!("fas: {e}"));
            }
        }
    }
}

/// Print stats about the last apply.
fn print_info(core: Core, session: *mut PluginSession) {
    let state = unsafe { r2plugin::state::<PluginState>(session) };
    match state.and_then(|s| s.applied.as_ref()) {
        Some(a) => core.echo(&format!(
            "fas: {}  symbols={}  native_symbols={}  lines={}  functions={}  xrefs={}",
            a.fas_path,
            a.flags.len(),
            a.native_symbol_count,
            a.line_count,
            a.function_count,
            a.xref_count
        )),
        None => core.echo("fas: nothing loaded"),
    }
}

/// Drop flags from the last apply.
fn unload(session: *mut PluginSession) {
    let Some(core) = (unsafe { Core::from_session(session) }) else {
        return;
    };
    let state = unsafe { r2plugin::state_mut::<PluginState>(session) };
    if let Some(state) = state {
        if let Some(applied) = state.applied.take() {
            if apply::unload(core, &applied) {
                core.echo("fas: unloaded");
            } else {
                state.applied = Some(applied);
                core.echo(
                    "fas: unload refused because the binary changed or metadata was modified",
                );
            }
            return;
        }
    }
    core.echo("fas: nothing to unload");
}

/// Static plugin descriptor.
pub static PLUGIN: CorePlugin = CorePlugin {
    meta: radare2::plugin_meta!(
        "fas",
        "FASM .fas symbolic dump loader",
        "Hamid R. K. Pishghadam, Cursor AI",
        "0.1.0",
        "MIT"
    ),
    init: Some(plugin_init),
    fini: Some(plugin_fini),
    call: Some(plugin_call),
};

radare2::export_core_plugin!(&PLUGIN, "r2fas");

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

    #[test]
    fn command_parser_accepts_only_fas_boundaries() {
        assert_eq!(strip_command("fas"), Some(""));
        assert_eq!(strip_command(" fas info"), Some(" info"));
        assert_eq!(strip_command("fas.unload"), Some(".unload"));
        assert_eq!(strip_command("fas?"), Some("?"));
        assert_eq!(strip_command("faster"), None);
        assert_eq!(strip_command("afas"), None);
    }

    #[test]
    fn help_covers_every_user_command() {
        let help = HELP.join("\n");
        for command in ["fas load", "fas info", "fas unload", "fas?"] {
            assert!(help.contains(command), "missing {command}: {help}");
        }
    }
}