exfiltrate 0.3.0

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Implements the native `memory` command.
#![cfg(not(target_arch = "wasm32"))]
use crate::command::{Command, Response};

/// The `memory` command.
///
/// Reports resident set size and whatever else the platform is willing to say.
/// The set of available statistics is genuinely platform-dependent, so this
/// command reports what it found and names what it could not find, rather than
/// printing zeroes that read like measurements.
pub(crate) struct Memory;

impl Command for Memory {
    fn name(&self) -> &'static str {
        "memory"
    }

    fn short_description(&self) -> &'static str {
        "Reports resident set size and other memory statistics the platform exposes."
    }

    fn full_description(&self) -> &'static str {
        "Reports the process's memory use.

What is available depends entirely on the platform:

* Linux reads /proc/self/status and /proc/self/statm, so RSS, peak RSS, virtual size
  and the shared/private split are all reported.
* macOS and Windows have no equivalent file, and reading their APIs would mean a
  platform crate. Those targets report what std can see and say plainly that the rest
  is unavailable, rather than printing a zero that looks like a measurement.

Numbers are a snapshot at the moment of the call. Two calls a second apart tell you
far more than one."
    }

    fn execute(&self, _args: Vec<String>) -> Result<Response, Response> {
        let stats = collect();
        if stats.is_empty() {
            return Ok(format!(
                "no memory statistics are available on {}.\n\
                 This target has no /proc; reading its native API would require a platform crate.",
                std::env::consts::OS
            )
            .into());
        }
        let width = stats
            .iter()
            .map(|(label, _)| label.len())
            .max()
            .unwrap_or(0);
        let mut out = String::new();
        for (label, value) in stats {
            out.push_str(&format!("{label:<width$}  {value}\n", width = width));
        }
        Ok(out.into())
    }
}

/// The statistics this platform can report, as label/value pairs.
fn collect() -> Vec<(String, String)> {
    #[cfg(target_os = "linux")]
    {
        linux_stats()
    }
    #[cfg(not(target_os = "linux"))]
    {
        Vec::new()
    }
}

#[cfg(target_os = "linux")]
fn linux_stats() -> Vec<(String, String)> {
    let Ok(status) = std::fs::read_to_string("/proc/self/status") else {
        return Vec::new();
    };
    parse_proc_status(&status)
}

/// Extracts the memory lines from `/proc/self/status`.
///
/// Split out from the file read so the parsing can be tested on every platform,
/// including the ones with no `/proc` to read.
#[cfg_attr(not(target_os = "linux"), allow(dead_code))]
fn parse_proc_status(status: &str) -> Vec<(String, String)> {
    const WANTED: &[(&str, &str)] = &[
        ("VmRSS", "resident (RSS)"),
        ("VmHWM", "peak resident"),
        ("VmSize", "virtual size"),
        ("VmPeak", "peak virtual"),
        ("RssAnon", "anonymous"),
        ("RssFile", "file-backed"),
        ("RssShmem", "shared"),
        ("Threads", "threads"),
    ];
    let mut stats = Vec::new();
    for (key, label) in WANTED {
        for line in status.lines() {
            let Some((name, value)) = line.split_once(':') else {
                continue;
            };
            if name.trim() != *key {
                continue;
            }
            let value = value.trim();
            let pretty = match value
                .strip_suffix(" kB")
                .and_then(|kb| kb.parse::<u64>().ok())
            {
                Some(kib) => format!("{} ({kib} kB)", human_bytes(kib * 1024)),
                None => value.to_string(),
            };
            stats.push((label.to_string(), pretty));
            break;
        }
    }
    stats
}

/// Formats a byte count the way a human reads one.
#[cfg_attr(not(target_os = "linux"), allow(dead_code))]
fn human_bytes(bytes: u64) -> String {
    const UNITS: &[&str] = &["B", "KiB", "MiB", "GiB", "TiB"];
    let mut value = bytes as f64;
    let mut unit = 0;
    while value >= 1024.0 && unit + 1 < UNITS.len() {
        value /= 1024.0;
        unit += 1;
    }
    if unit == 0 {
        format!("{bytes} B")
    } else {
        format!("{value:.1} {}", UNITS[unit])
    }
}

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

    const SAMPLE: &str = "Name:\tdemo\n\
        VmPeak:\t 2097152 kB\n\
        VmSize:\t 1048576 kB\n\
        VmHWM:\t   65536 kB\n\
        VmRSS:\t   32768 kB\n\
        RssAnon:\t   16384 kB\n\
        RssFile:\t   16384 kB\n\
        Threads:\t7\n";

    #[test]
    fn the_parser_picks_out_the_memory_lines_in_a_stable_order() {
        let stats = parse_proc_status(SAMPLE);
        let labels: Vec<&str> = stats.iter().map(|(label, _)| label.as_str()).collect();
        assert_eq!(
            labels,
            [
                "resident (RSS)",
                "peak resident",
                "virtual size",
                "peak virtual",
                "anonymous",
                "file-backed",
                "threads"
            ]
        );
    }

    #[test]
    fn kilobyte_values_are_rendered_in_human_units_without_losing_the_raw_number() {
        let stats = parse_proc_status(SAMPLE);
        let rss = &stats
            .iter()
            .find(|(label, _)| label == "resident (RSS)")
            .unwrap()
            .1;
        assert_eq!(rss, "32.0 MiB (32768 kB)");
    }

    #[test]
    fn a_non_kilobyte_value_is_passed_through_untouched() {
        let stats = parse_proc_status(SAMPLE);
        let threads = &stats
            .iter()
            .find(|(label, _)| label == "threads")
            .unwrap()
            .1;
        assert_eq!(threads, "7");
    }

    #[test]
    fn a_status_file_without_the_memory_lines_yields_nothing_rather_than_zeroes() {
        assert!(parse_proc_status("Name:\tdemo\n").is_empty());
    }

    #[test]
    fn the_command_answers_on_every_native_target() {
        let text = Memory.execute(Vec::new()).unwrap().into_string();
        assert!(!text.is_empty());
        #[cfg(target_os = "linux")]
        assert!(text.contains("resident (RSS)"), "{text}");
        #[cfg(not(target_os = "linux"))]
        assert!(text.contains("no memory statistics"), "{text}");
    }

    #[test]
    fn byte_formatting_crosses_units_where_expected() {
        assert_eq!(human_bytes(512), "512 B");
        assert_eq!(human_bytes(1024), "1.0 KiB");
        assert_eq!(human_bytes(1024 * 1024 * 3 / 2), "1.5 MiB");
    }
}