procutils-pmap 0.2.0

use clap::Parser;
use cols::Cols;
use procfs::process::{MMPermissions, MMapPath, MemoryMap, Process};
use procutils_common::{MAX_TERM_WIDTH, man::ManContent};
use std::process::ExitCode;

pub const MAN: ManContent = ManContent {
    description: Some(include_str!("../man/description.man")),
    extra_sections: &[
        (
            "FIELD DESCRIPTIONS",
            include_str!("../man/field_descriptions.man"),
        ),
        ("EXAMPLES", include_str!("../man/examples.man")),
        ("NOTES", include_str!("../man/notes.man")),
        ("DIVERGENCES", include_str!("../man/divergences.man")),
        ("SEE ALSO", include_str!("../man/see_also.man")),
    ],
};

/// Report memory map of a process.
#[derive(Parser)]
#[command(
    name = "pmap",
    version,
    about,
    max_term_width = MAX_TERM_WIDTH,
    override_usage = "pmap [options] PID [PID ...]"
)]
pub struct Args {
    /// Show the extended format.
    #[arg(short = 'x', long)]
    extended: bool,

    /// Show the device format.
    #[arg(short, long)]
    device: bool,

    /// Do not display some header or footer lines.
    #[arg(short, long)]
    quiet: bool,

    /// Show full path to files in the mapping column.
    #[arg(short = 'p', long)]
    show_path: bool,

    /// Show kernel-supplied mapping names (`[vdso]`, `[vsyscall]`,
    /// `[stack]`, `[heap]`, `[vvar]`) verbatim, without our default
    /// padded form (`[ vdso ]` etc).
    #[arg(short = 'k', long = "use-kernel-name")]
    use_kernel_name: bool,

    /// Limit results to the given address range. Accepts a single
    /// address or a comma-separated `LOW,HIGH` pair. Hex with or
    /// without `0x` prefix.
    #[arg(short = 'A', long = "range", value_name = "LOW[,HIGH]")]
    range: Option<String>,

    /// Process IDs.
    #[arg(required = true)]
    pid: Vec<i32>,
}

/// Parse `-A LOW` or `-A LOW,HIGH`. A single value sets HIGH = LOW
/// (interpret as "any mapping containing this address").
fn parse_range(s: &str) -> Result<(u64, u64), String> {
    fn parse_addr(a: &str) -> Result<u64, String> {
        let stripped = a.trim_start_matches("0x").trim_start_matches("0X");
        u64::from_str_radix(stripped, 16)
            .map_err(|_| format!("invalid address: {a}"))
    }
    match s.split_once(',') {
        Some((lo, hi)) => Ok((parse_addr(lo)?, parse_addr(hi)?)),
        None => {
            let v = parse_addr(s)?;
            Ok((v, v))
        }
    }
}

/// True if `[map_start, map_end)` overlaps the inclusive range `[low, high]`.
fn in_range(map_start: u64, map_end: u64, low: u64, high: u64) -> bool {
    map_start <= high && map_end > low
}

#[derive(Cols)]
struct DefaultRow {
    #[column(header = "Address")]
    address: String,
    #[column(right, header = "Kbytes")]
    kbytes: String,
    #[column(header = "Mode")]
    mode: String,
    #[column(header = "Mapping")]
    mapping: String,
}

#[derive(Cols)]
struct ExtendedRow {
    #[column(header = "Address")]
    address: String,
    #[column(right, header = "Kbytes")]
    kbytes: String,
    #[column(right, header = "RSS")]
    rss: String,
    #[column(right, header = "Dirty")]
    dirty: String,
    #[column(header = "Mode")]
    mode: String,
    #[column(header = "Mapping")]
    mapping: String,
}

#[derive(Cols)]
struct DeviceRow {
    #[column(header = "Address")]
    address: String,
    #[column(right, header = "Kbytes")]
    kbytes: String,
    #[column(header = "Mode")]
    mode: String,
    #[column(header = "Offset")]
    offset: String,
    #[column(header = "Device")]
    device: String,
    #[column(header = "Mapping")]
    mapping: String,
}

fn format_perms_pmap(perms: MMPermissions) -> String {
    let r = if perms.contains(MMPermissions::READ) {
        'r'
    } else {
        '-'
    };
    let w = if perms.contains(MMPermissions::WRITE) {
        'w'
    } else {
        '-'
    };
    let x = if perms.contains(MMPermissions::EXECUTE) {
        'x'
    } else {
        '-'
    };
    let s = if perms.contains(MMPermissions::SHARED) {
        's'
    } else {
        '-'
    };
    // pmap uses 5-char format: "rwx-s" or "rwx--"
    format!("{r}{w}{x}{s}-")
}

fn format_mapping(
    path: &MMapPath,
    show_path: bool,
    kernel_name: bool,
) -> String {
    match path {
        MMapPath::Path(p) => {
            if show_path {
                p.display().to_string()
            } else {
                p.file_name()
                    .map(|f| f.to_string_lossy().into_owned())
                    .unwrap_or_else(|| p.display().to_string())
            }
        }
        MMapPath::Heap if kernel_name => "[heap]".into(),
        MMapPath::Stack if kernel_name => "[stack]".into(),
        MMapPath::TStack(tid) if kernel_name => format!("[stack:{tid}]"),
        MMapPath::Vdso if kernel_name => "[vdso]".into(),
        MMapPath::Vvar if kernel_name => "[vvar]".into(),
        MMapPath::Vsyscall if kernel_name => "[vsyscall]".into(),
        MMapPath::Anonymous if kernel_name => String::new(),
        MMapPath::Vsys(id) if kernel_name => format!("[sysv:{id}]"),
        MMapPath::Rollup if kernel_name => "[rollup]".into(),
        MMapPath::Other(s) if kernel_name => format!("[{s}]"),

        MMapPath::Heap => "[ anon ]".into(),
        MMapPath::Stack => "[ stack ]".into(),
        MMapPath::TStack(tid) => format!("[ stack:{tid} ]"),
        MMapPath::Vdso => "[ vdso ]".into(),
        MMapPath::Vvar => "[ vvar ]".into(),
        MMapPath::Vsyscall => "[ anon ]".into(),
        MMapPath::Anonymous => "[ anon ]".into(),
        MMapPath::Vsys(_) => "[ sysv ]".into(),
        MMapPath::Rollup => "[ rollup ]".into(),
        MMapPath::Other(s) => format!("[ {s} ]"),
    }
}

fn kbytes(map: &MemoryMap) -> u64 {
    (map.address.1 - map.address.0) / 1024
}

pub fn run(args: Args) -> ExitCode {
    let range = match args.range.as_deref().map(parse_range) {
        Some(Ok(r)) => Some(r),
        Some(Err(e)) => {
            eprintln!("pmap: {e}");
            return ExitCode::from(2);
        }
        None => None,
    };

    let mut not_found = false;

    for (i, &pid) in args.pid.iter().enumerate() {
        if i > 0 {
            println!();
        }

        let proc = match Process::new(pid) {
            Ok(p) => p,
            Err(e) => {
                eprintln!("pmap: {pid}: {e}");
                not_found = true;
                continue;
            }
        };

        let cmdline = proc.cmdline().unwrap_or_default().join(" ");

        println!("{pid}:   {cmdline}");

        if args.extended {
            show_extended(&proc, &args, range, pid, &mut not_found);
        } else if args.device {
            show_device(&proc, &args, range, pid, &mut not_found);
        } else {
            show_default(&proc, &args, range, pid, &mut not_found);
        }
    }

    if not_found {
        ExitCode::from(42)
    } else {
        ExitCode::SUCCESS
    }
}

fn show_default(
    proc: &Process,
    args: &Args,
    range: Option<(u64, u64)>,
    pid: i32,
    not_found: &mut bool,
) {
    let maps = match proc.maps() {
        Ok(m) => m,
        Err(e) => {
            eprintln!("pmap: {pid}: {e}");
            *not_found = true;
            return;
        }
    };

    let filtered: Vec<&MemoryMap> = maps
        .0
        .iter()
        .filter(|m| match range {
            Some((lo, hi)) => in_range(m.address.0, m.address.1, lo, hi),
            None => true,
        })
        .collect();

    let rows: Vec<DefaultRow> = filtered
        .iter()
        .map(|m| DefaultRow {
            address: format!("{:016x}", m.address.0),
            kbytes: format!("{}K", kbytes(m)),
            mode: format_perms_pmap(m.perms),
            mapping: format_mapping(
                &m.pathname,
                args.show_path,
                args.use_kernel_name,
            ),
        })
        .collect();

    let total_kb: u64 = filtered.iter().copied().map(kbytes).sum();

    print_rows(&rows);

    if !args.quiet {
        println!(" total          {total_kb:>5}K");
    }
}

fn show_extended(
    proc: &Process,
    args: &Args,
    range: Option<(u64, u64)>,
    pid: i32,
    not_found: &mut bool,
) {
    let maps = match proc.smaps() {
        Ok(m) => m,
        Err(e) => {
            eprintln!("pmap: {pid}: {e}");
            *not_found = true;
            return;
        }
    };

    let mut rows: Vec<ExtendedRow> = Vec::new();
    let mut total_kb = 0u64;
    let mut total_rss = 0u64;
    let mut total_dirty = 0u64;

    for m in &maps.0 {
        if let Some((lo, hi)) = range
            && !in_range(m.address.0, m.address.1, lo, hi)
        {
            continue;
        }
        let kb = kbytes(m);
        let rss = m.extension.map.get("Rss").copied().unwrap_or(0) / 1024;
        let dirty =
            (m.extension.map.get("Private_Dirty").copied().unwrap_or(0)
                + m.extension.map.get("Shared_Dirty").copied().unwrap_or(0))
                / 1024;

        total_kb += kb;
        total_rss += rss;
        total_dirty += dirty;

        rows.push(ExtendedRow {
            address: format!("{:016x}", m.address.0),
            kbytes: format!("{kb}"),
            rss: format!("{rss}"),
            dirty: format!("{dirty}"),
            mode: format_perms_pmap(m.perms),
            mapping: format_mapping(
                &m.pathname,
                args.show_path,
                args.use_kernel_name,
            ),
        });
    }

    print_rows(&rows);

    if !args.quiet {
        println!("{} ------- ------- ------- ", "-".repeat(16),);
        println!(
            "total kB {:>14} {:>7} {:>7}",
            total_kb, total_rss, total_dirty,
        );
    }
}

fn show_device(
    proc: &Process,
    args: &Args,
    range: Option<(u64, u64)>,
    pid: i32,
    not_found: &mut bool,
) {
    let maps = match proc.maps() {
        Ok(m) => m,
        Err(e) => {
            eprintln!("pmap: {pid}: {e}");
            *not_found = true;
            return;
        }
    };

    let mut rows: Vec<DeviceRow> = Vec::new();
    let mut total_kb = 0u64;
    let mut total_writable_private = 0u64;
    let mut total_shared = 0u64;

    for m in &maps.0 {
        if let Some((lo, hi)) = range
            && !in_range(m.address.0, m.address.1, lo, hi)
        {
            continue;
        }
        let kb = kbytes(m);
        total_kb += kb;

        if m.perms.contains(MMPermissions::SHARED) {
            total_shared += kb;
        }
        if m.perms.contains(MMPermissions::WRITE)
            && m.perms.contains(MMPermissions::PRIVATE)
        {
            total_writable_private += kb;
        }

        rows.push(DeviceRow {
            address: format!("{:016x}", m.address.0),
            kbytes: format!("{kb}"),
            mode: format_perms_pmap(m.perms),
            offset: format!("{:016x}", m.offset),
            device: format!("{:03}:{:05}", m.dev.0, m.dev.1),
            mapping: format_mapping(
                &m.pathname,
                args.show_path,
                args.use_kernel_name,
            ),
        });
    }

    print_rows(&rows);

    if !args.quiet {
        println!(
            "mapped: {total_kb}K    writeable/private: {total_writable_private}K    shared: {total_shared}K",
        );
    }
}

fn print_rows<T: Cols>(rows: &[T]) {
    let table = T::to_table(rows);
    cols::print_table(&table, &mut std::io::stdout().lock()).unwrap();
}