baryl 0.0.4

//! Building a command-line runner of your own.
//!
//! The argument types the `baryl` binary is built from. A runner that uses them
//! takes the same flags, spells addresses the same way, and parses
//! `--component` and `--arg` identically.
//!
//! [`CommonArgs`] and [`ComponentArgs`] are `clap::Args`, so flatten them into
//! your own `Parser`. The `parse_*` functions are plain functions and work
//! wherever you have a string.
//!
//! Doc comments on the argument fields below are what `--help` prints.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::str::FromStr;

use anyhow::bail;
use log::LevelFilter;

use crate::run::Options;

// FIXME: Hmmm, why are we defining CLI types here, and also in the baryl-cli crate? they should be
// consolidated

/// The flags every verb shares. Flatten into your own `Parser` and fold the
/// result into `Options` with [`apply`](Self::apply).
#[derive(clap::Args)]
pub struct CommonArgs {
    /// Where this invocation's artifacts are written.
    #[arg(long, global = true, value_name = "DIR")]
    pub state_directory: Option<PathBuf>,

    /// off, error, warn, info, debug, trace.
    #[arg(long, global = true, default_value = "info")]
    pub log_level: LevelFilter,
}

impl CommonArgs {
    /// Set `state_dir` and `log_level` on `options` from these flags.
    ///
    /// # Errors
    ///
    /// `--state-directory` is not UTF-8 or holds an interior NUL.
    pub fn apply(&self, options: Options) -> anyhow::Result<Options> {
        Ok(options
            .state_dir(self.state_directory.as_deref())?
            .log_level(self.log_level))
    }
}

/// Which components to load, and what to pass each of them.
///
/// A component's name is its filename stem, and that is what an `--arg` keys
/// on: `--component ./libhello.so --arg hello:who=world` gives `libhello.so`
/// the argv `["--who=world"]`. Turn one of these into what `Options::components`
/// takes with [`component_args_split`].
#[derive(clap::Args)]
pub struct ComponentArgs {
    /// Load a component. Repeat to load more than one.
    #[arg(long, value_name = "PATH")]
    pub component: Vec<PathBuf>,

    /// One flag for a component's own parser. Repeatable.
    #[arg(long = "arg", value_name = "NAME:KEY[=VALUE]")]
    pub args: Vec<String>,
}

/// Pair each `--component` path with the argv its `--arg` flags built.
///
/// The result goes straight into `Options::components`. Each `--arg` becomes a
/// `--<key>` or `--<key>=<value>` in the named component's argv, in the order
/// they were given.
///
/// # Errors
///
/// An `--arg` that is not `<name>:<key>[=<value>]`; an `--arg` naming no loaded
/// component; or two `--component` paths with the same filename stem, which
/// would make their flags indistinguishable. All three used to surface far from
/// the cause — a missing flag reported by the component itself much later, or
/// two components silently sharing one argv — so they are caught here.
///
/// # Examples
///
/// ```ignore
/// // --component ./libhello.so --component ./libtrace.so
/// // --arg hello:who=world --arg trace:verbose
/// let pairs = component_args_split(&args)?;
/// assert_eq!(pairs[0].1, vec!["--who=world".to_string()]);
/// assert_eq!(pairs[1].1, vec!["--verbose".to_string()]);
///
/// let options = Options::default().components(&pairs)?;
/// ```
pub fn component_args_split(args: &ComponentArgs) -> anyhow::Result<Vec<(PathBuf, Vec<String>)>> {
    let mut keyed: BTreeMap<String, Vec<String>> = BTreeMap::new();
    for entry in &args.args {
        // The first `:` only; a `:` in the value is the value's.
        let Some((name, flag)) = entry.split_once(':') else {
            bail!("`--arg {entry}` is not <name>:<key>[=<value>]");
        };
        keyed
            .entry(name.to_string())
            .or_default()
            .push(format!("--{flag}"));
    }
    check_component_names(&args.component, &keyed)?;
    Ok(args
        .component
        .iter()
        .map(|path| {
            let argv = keyed
                .get(&component_name(path))
                .cloned()
                .unwrap_or_default();
            (path.clone(), argv)
        })
        .collect())
}

/// A component's name: its filename stem, which is what `--arg` keys on.
// FIXME: fold into ComponentArgs as a method.
fn component_name(path: &Path) -> String {
    path.file_stem()
        .unwrap_or_default()
        .to_string_lossy()
        .into_owned()
}

/// Reject a duplicate component stem and an `--arg` naming no component.
fn check_component_names(
    components: &[PathBuf],
    component_args: &BTreeMap<String, Vec<String>>,
) -> anyhow::Result<()> {
    let mut names: Vec<String> = components.iter().map(|p| component_name(p)).collect();
    names.sort();
    if let Some(dup) = names.windows(2).find(|w| w[0] == w[1]) {
        bail!(
            "two --component paths share the name `{}`; \
             their `--arg {}:<key>` flags cannot be told apart",
            dup[0],
            dup[0],
        );
    }
    if let Some(unmatched) = component_args.keys().find(|n| !names.contains(n)) {
        bail!(
            "`--arg {unmatched}:<key>` names no --component{}",
            match names.as_slice() {
                [] => " (none were given)".to_string(),
                loaded => format!("; loaded: {}", loaded.join(", ")),
            }
        );
    }
    Ok(())
}

/// An address as a command line spells it, in the four forms one can take.
///
/// Parse with [`parse_addr_spec`], or through `FromStr` — which is what a clap
/// field of this type does for you. Resolving it is yours: a `Reg` needs the
/// register file, a `Symbol` needs the enlighten subsystem.
///
/// # Examples
///
/// ```ignore
/// // 0xffffffff81000000  -> Virt
/// // p:0x1000            -> Phys
/// // @rip                -> Reg(Rip)
/// // do_sys_openat2      -> Symbol
/// let spec: AddrSpec = "@rsp".parse()?;
/// ```
#[derive(Debug, Clone)]
pub enum AddrSpec {
    /// A bare number: a guest virtual address.
    Virt(u64),
    /// `p:` or `phys:` prefixed: a guest physical address.
    Phys(u64),
    /// `@` prefixed: read the address out of a register.
    Reg(RegName),
    /// Anything else: a name for enlighten to resolve.
    Symbol(String),
}

impl FromStr for AddrSpec {
    type Err = lexopt::Error;

    fn from_str(value: &str) -> Result<AddrSpec, lexopt::Error> {
        parse_addr_spec(value)
    }
}

/// A register named where an address is wanted.
///
/// `@`-prefixed on the command line: `@rip`, `@cr3`. A bare `rip` parses as a
/// symbol. `@gs_base` and `@kernel_gs_base` both spell
/// [`KernelGsBase`](Self::KernelGsBase).
#[derive(Debug, Clone, Copy)]
pub enum RegName {
    Rip,
    Rsp,
    Rbp,
    Rax,
    Rbx,
    Rcx,
    Rdx,
    Rsi,
    Rdi,
    R8,
    R9,
    R10,
    R11,
    R12,
    R13,
    R14,
    R15,
    Cr0,
    Cr2,
    Cr3,
    Cr4,
    KernelGsBase,
    Lstar,
}

impl RegName {
    /// The name as `@<name>` spells it, without the `@`; `None` for a register
    /// this list has not got.
    fn parse(s: &str) -> Option<RegName> {
        Some(match s {
            "rip" => RegName::Rip,
            "rsp" => RegName::Rsp,
            "rbp" => RegName::Rbp,
            "rax" => RegName::Rax,
            "rbx" => RegName::Rbx,
            "rcx" => RegName::Rcx,
            "rdx" => RegName::Rdx,
            "rsi" => RegName::Rsi,
            "rdi" => RegName::Rdi,
            "r8" => RegName::R8,
            "r9" => RegName::R9,
            "r10" => RegName::R10,
            "r11" => RegName::R11,
            "r12" => RegName::R12,
            "r13" => RegName::R13,
            "r14" => RegName::R14,
            "r15" => RegName::R15,
            "cr0" => RegName::Cr0,
            "cr2" => RegName::Cr2,
            "cr3" => RegName::Cr3,
            "cr4" => RegName::Cr4,
            "kernel_gs_base" | "gs_base" => RegName::KernelGsBase,
            "lstar" => RegName::Lstar,
            _ => return None,
        })
    }
}

/// What a search looks for, in the five ways one pattern can be spelled.
///
/// Always `<kind>:<value>`, so a pattern is one word wherever a command takes
/// one. [`bytes`](Self::bytes) turns any of them into the bytes to scan for.
///
/// # Examples
///
/// ```ignore
/// let n: Needle = "hex:4d5a9000".parse()?;
/// assert_eq!(n.bytes(), vec![0x4d, 0x5a, 0x90, 0x00]);
///
/// // Little-endian, so this finds a stored pointer, not its spelling.
/// let p: Needle = "u64:0xffffffff81000000".parse()?;
/// assert_eq!(p.bytes(), 0xffff_ffff_8100_0000u64.to_le_bytes().to_vec());
///
/// let w: Needle = "utf16:C:".parse()?;
/// assert_eq!(w.bytes(), vec![b'C', 0, b':', 0]);
/// ```
#[derive(Debug, Clone)]
pub enum Needle {
    /// `hex:4d5a` — raw bytes, whitespace allowed between them.
    Hex(Vec<u8>),
    /// `ascii:GET /` — the bytes of the string.
    Ascii(String),
    /// `utf16:C:\Windows` — UTF-16LE, which is how Windows stores a string.
    Utf16(String),
    /// `u64:0x...` — eight bytes, little-endian.
    U64(u64),
    /// `u32:0x...` — four bytes, little-endian.
    U32(u32),
}

impl Needle {
    /// The bytes to look for.
    pub fn bytes(&self) -> Vec<u8> {
        match self {
            Needle::Hex(bytes) => bytes.clone(),
            Needle::Ascii(s) => s.clone().into_bytes(),
            Needle::Utf16(s) => s.encode_utf16().flat_map(u16::to_le_bytes).collect(),
            Needle::U64(value) => value.to_le_bytes().to_vec(),
            Needle::U32(value) => value.to_le_bytes().to_vec(),
        }
    }
}

impl FromStr for Needle {
    type Err = lexopt::Error;

    fn from_str(value: &str) -> Result<Needle, lexopt::Error> {
        let Some((kind, body)) = value.split_once(':') else {
            return Err(lexopt::Error::ParsingFailed {
                value: value.to_string(),
                error: "expected hex:, ascii:, utf16:, u64: or u32:".into(),
            });
        };
        Ok(match kind {
            "hex" => Needle::Hex(parse_hex_bytes(body)?),
            "ascii" => Needle::Ascii(body.to_string()),
            "utf16" => Needle::Utf16(body.to_string()),
            "u64" => Needle::U64(parse_addr(body)?),
            "u32" => Needle::U32(parse_addr(body)? as u32),
            _ => {
                return Err(lexopt::Error::ParsingFailed {
                    value: value.to_string(),
                    error: "expected hex:, ascii:, utf16:, u64: or u32:".into(),
                });
            },
        })
    }
}

/// How guest memory is rendered when it is printed.
#[derive(Debug, Clone, Copy, Default)]
pub enum MemFormat {
    /// Offset, sixteen bytes, then their printable spelling.
    #[default]
    Hex,
    /// The bytes themselves, for a pipe.
    Raw,
    /// One little-endian `u64` per line, which is what a stack reads as.
    Words,
    /// The printable spelling alone.
    Ascii,
}

impl FromStr for MemFormat {
    type Err = lexopt::Error;

    fn from_str(value: &str) -> Result<MemFormat, lexopt::Error> {
        Ok(match value {
            "hex" => MemFormat::Hex,
            "raw" => MemFormat::Raw,
            "words" => MemFormat::Words,
            "ascii" => MemFormat::Ascii,
            _ => {
                return Err(lexopt::Error::ParsingFailed {
                    value: value.to_string(),
                    error: "expected hex, raw, words or ascii".into(),
                });
            },
        })
    }
}

/// Whose loaded images to list.
///
/// `all` is the default: "which images does this guest have" is usually the
/// question, and the kernel's alone rarely answer it.
#[derive(Debug, Clone, Copy)]
pub enum ModuleScope {
    /// `kernel` — the kernel image and its modules.
    Kernel,
    /// A bare number — the images mapped into that process.
    Pid(u32),
    /// `all` — both.
    All,
}

impl FromStr for ModuleScope {
    type Err = lexopt::Error;

    fn from_str(value: &str) -> Result<ModuleScope, lexopt::Error> {
        Ok(match value {
            "kernel" => ModuleScope::Kernel,
            "all" => ModuleScope::All,
            pid => ModuleScope::Pid(parse_addr(pid)? as u32),
        })
    }
}

/// The questions an inspect command can ask of a checkpoint.
///
/// Every one is answered from a `&Control` alone, so none of them runs a guest
/// instruction — a checkpoint can be read as many times as you like and is
/// never changed by the reading.
///
/// Where a variant takes `cr3`, that is the address space to work in, and
/// leaving it off uses the one the checkpoint stopped in.
#[derive(Debug, Clone, clap::Subcommand)]
pub enum Query {
    /// What this checkpoint is.
    Info,
    /// The register file.
    Regs {
        /// Add control and segment registers.
        #[arg(long)]
        all: bool,
    },
    /// One address's page walk.
    Translate {
        /// A VA, `p:<pa>`, `@<reg>` or a symbol.
        addr: AddrSpec,
        /// The space to walk; default the checkpoint's.
        #[arg(long, value_parser = parse_addr)]
        cr3: Option<u64>,
    },
    /// Guest memory, rendered.
    Read {
        /// A VA, `p:<pa>`, `@<reg>` or a symbol.
        addr: AddrSpec,
        /// Bytes to read.
        #[arg(long, value_parser = parse_len, default_value = "0x40")]
        len: u64,
        /// The space to read; default the checkpoint's.
        #[arg(long, value_parser = parse_addr)]
        cr3: Option<u64>,
        /// hex, words, ascii or raw.
        #[arg(long, default_value = "hex")]
        format: MemFormat,
    },
    /// Instructions from an address.
    Disasm {
        /// A VA, `p:<pa>`, `@<reg>` or a symbol.
        addr: AddrSpec,
        /// Instructions to print.
        #[arg(long, default_value_t = 10)]
        count: u32,
        /// Bytes to fetch; 0 is however many `count` needs.
        #[arg(long, value_parser = parse_len, default_value = "0")]
        len: u64,
        /// The space to read; default the checkpoint's.
        #[arg(long, value_parser = parse_addr)]
        cr3: Option<u64>,
        /// Decode as 16, 32 or 64-bit code.
        #[arg(long, default_value_t = 64)]
        bits: u32,
    },
    /// Every match of a pattern.
    Search {
        /// A VA, `p:<pa>`, `@<reg>` or a symbol.
        addr: AddrSpec,
        /// hex:, ascii:, utf16:, u64: or u32:.
        needle: Needle,
        /// Bytes to search from `addr`.
        #[arg(long, value_parser = parse_len, conflicts_with = "end")]
        len: Option<u64>,
        /// Search to here instead of `--len`.
        #[arg(long, value_parser = parse_addr)]
        end: Option<u64>,
        /// Visit only addresses this far apart.
        #[arg(long, value_parser = parse_len, default_value = "1")]
        stride: u64,
        /// The space to search; default the checkpoint's.
        #[arg(long, value_parser = parse_addr)]
        cr3: Option<u64>,
        /// Stop after this many hits.
        #[arg(long, default_value_t = u32::MAX, hide_default_value = true)]
        max_hits: u32,
    },
    /// Every address holding one as a pointer.
    Xrefs {
        /// A VA, `p:<pa>`, `@<reg>` or a symbol.
        addr: AddrSpec,
        /// The space to search; default the checkpoint's.
        #[arg(long, value_parser = parse_addr)]
        cr3: Option<u64>,
        /// Stop after this many hits.
        #[arg(long, default_value_t = u32::MAX, hide_default_value = true)]
        max_hits: u32,
    },
    /// Every guest process.
    Ps,
    /// Every image the guest has loaded.
    Modules {
        /// `kernel`, `all`, or one pid.
        #[arg(default_value = "all")]
        scope: ModuleScope,
    },
    /// A name to the address enlighten gives it.
    Resolve {
        /// `name`, `module!name` or `module+0xNN`.
        spec: String,
        /// Resolve against this process; default kernel space.
        #[arg(long)]
        pid: Option<u32>,
    },
}

/// One inspect invocation: which checkpoint, which question, and the shared
/// flags.
#[derive(clap::Args)]
pub struct InspectArgs {
    /// The checkpoint to read.
    pub checkpoint: PathBuf,
    #[command(subcommand)]
    pub query: Query,
    #[command(flatten)]
    pub common: CommonArgs,
}

impl InspectArgs {
    /// The `Options` to open this checkpoint with.
    ///
    /// No machine config — the checkpoint carries its own — and no components,
    /// so a reader loads nothing that could write.
    ///
    /// # Errors
    ///
    /// `--state-directory` is not UTF-8 or holds an interior NUL.
    pub fn options(&self) -> anyhow::Result<Options> {
        self.common.apply(Options::default())
    }
}

/// A byte count: a number with an optional `K`, `M` or `G` suffix, in either
/// case.
///
/// The suffixes are powers of 1024, not 1000.
///
/// # Errors
///
/// The number does not parse, or the count overflows 64 bits once scaled.
///
/// # Examples
///
/// ```ignore
/// assert_eq!(parse_len("64")?, 64);
/// assert_eq!(parse_len("0x40")?, 64);
/// assert_eq!(parse_len("4K")?, 4096);
/// assert_eq!(parse_len("2m")?, 2 * 1024 * 1024);
/// ```
pub fn parse_len(value: &str) -> Result<u64, lexopt::Error> {
    let (digits, scale) = match value.as_bytes().last() {
        Some(b'K' | b'k') => (&value[..value.len() - 1], 1u64 << 10),
        Some(b'M' | b'm') => (&value[..value.len() - 1], 1u64 << 20),
        Some(b'G' | b'g') => (&value[..value.len() - 1], 1u64 << 30),
        _ => (value, 1),
    };
    let n = parse_addr(digits)?;
    n.checked_mul(scale)
        .ok_or_else(|| lexopt::Error::ParsingFailed {
            value: value.to_string(),
            error: "length overflows 64 bits".into(),
        })
}

/// A number: `0x`- or `0X`-prefixed hex, or plain decimal.
///
/// The one number parser here, so every flag taking an address takes it the
/// same way.
///
/// # Errors
///
/// Anything that is not one of those two forms — including a bare hex string
/// with no `0x`, which parses as decimal or not at all.
pub fn parse_addr(value: &str) -> Result<u64, lexopt::Error> {
    let parsed = match value
        .strip_prefix("0x")
        .or_else(|| value.strip_prefix("0X"))
    {
        Some(hex) => u64::from_str_radix(hex, 16),
        None => value.parse(),
    };
    parsed.map_err(|_| lexopt::Error::ParsingFailed {
        value: value.to_string(),
        error: "expected a decimal or 0x-prefixed number".into(),
    })
}

/// A hex string as the bytes it spells. Whitespace anywhere in it is ignored,
/// so `"4d5a"` and `"4d 5a"` are the same two bytes.
///
/// # Errors
///
/// An empty string, an odd number of digits, or a character that is not a hex
/// digit.
pub fn parse_hex_bytes(value: &str) -> Result<Vec<u8>, lexopt::Error> {
    let bad = |what: &'static str| lexopt::Error::ParsingFailed {
        value: value.to_string(),
        error: what.into(),
    };
    let digits: String = value.chars().filter(|c| !c.is_whitespace()).collect();
    if digits.is_empty() || !digits.len().is_multiple_of(2) {
        return Err(bad("expected an even number of hex digits"));
    }
    digits
        .as_bytes()
        .chunks(2)
        .map(|pair| {
            let s = std::str::from_utf8(pair).map_err(|_| bad("expected hex digits"))?;
            u8::from_str_radix(s, 16).map_err(|_| bad("expected hex digits"))
        })
        .collect()
}

/// One address, in whichever of the four forms it was written.
///
/// `@name` is a register, `p:` or `phys:` a physical address, a number a
/// virtual address, and anything else a symbol — so a name that fails to parse
/// as a number is never an error here, it is a symbol for enlighten to resolve.
///
/// # Errors
///
/// Only two things fail: `@` followed by a register this build does not know,
/// and a `p:` or `phys:` prefix followed by something that is not a number.
///
/// # Examples
///
/// ```ignore
/// assert!(matches!(parse_addr_spec("0x1000")?, AddrSpec::Virt(0x1000)));
/// assert!(matches!(parse_addr_spec("p:0x1000")?, AddrSpec::Phys(0x1000)));
/// assert!(matches!(parse_addr_spec("@rip")?, AddrSpec::Reg(RegName::Rip)));
/// assert!(matches!(parse_addr_spec("do_sys_openat2")?, AddrSpec::Symbol(_)));
/// assert!(parse_addr_spec("@nosuchreg").is_err());
/// ```
pub fn parse_addr_spec(value: &str) -> Result<AddrSpec, lexopt::Error> {
    if let Some(name) = value.strip_prefix('@') {
        return RegName::parse(name).map(AddrSpec::Reg).ok_or_else(|| {
            lexopt::Error::ParsingFailed {
                value: value.to_string(),
                error: "no such register; see `baryl inspect --help`".into(),
            }
        });
    }
    if let Some(pa) = value
        .strip_prefix("p:")
        .or_else(|| value.strip_prefix("phys:"))
    {
        return Ok(AddrSpec::Phys(parse_addr(pa)?));
    }
    match parse_addr(value) {
        Ok(va) => Ok(AddrSpec::Virt(va)),
        Err(_) => Ok(AddrSpec::Symbol(value.to_string())),
    }
}