nordisk/heimdall

crates/heimdall-trace/src/record.rs

// SPDX-License-Identifier: AGPL-3.0-only
// Copyright (C) 2026 Ignalina ApS and Vetra Automation AB

//! The bounded recorder: adds entries to a [`Trace`] while the encoded
//! document stays within its `max_bytes`. Heimdall is a defensive sandbox; a
//! package that makes millions of distinct calls fills the bound, the rest is
//! counted in `dropped` and the trace says [`Status::Truncated`], never
//! silently complete.

use crate::doc::{ExecEvent, FileEvent, NetEvent, Status, SyscallFact, Trace};

/// Bytes kept free for the header fields set at the end (exit, counters,
/// the incomplete reason).
pub const HEADER_RESERVE: u64 = 2048;

/// A trace under construction.
#[derive(Debug)]
pub struct Recorder {
    trace: Trace,
    used: u64,
}

impl Recorder {
    /// Start a trace of `argv` bounded by `max_bytes` (at least
    /// [`crate::doc::MIN_MAX_BYTES`]).
    pub fn new(argv: Vec<String>, max_bytes: u64) -> Recorder {
        let trace = Trace::new(argv, max_bytes);
        let used = trace.to_canonical().len() as u64 + HEADER_RESERVE;
        Recorder { trace, used }
    }

    fn drop_one(&mut self) {
        self.trace.dropped += 1;
        if self.trace.status == Status::Complete {
            self.trace.status = Status::Truncated;
        }
    }

    /// Room for `n` more bytes; when there is none the entry is dropped.
    fn room(&mut self, n: u64) -> bool {
        if self.used.saturating_add(n) > self.trace.max_bytes {
            self.drop_one();
            return false;
        }
        self.used += n;
        true
    }

    fn cost<T: serde::Serialize + ?Sized>(v: &T) -> u64 {
        // The entry plus its comma.
        serde_json::to_vec(v)
            .map(|b| b.len() as u64 + 1)
            .unwrap_or(u64::MAX / 4)
    }

    pub fn exec(&mut self, e: ExecEvent) {
        if !self.trace.exec.contains(&e) && self.room(Self::cost(&e)) {
            self.trace.exec.insert(e);
        }
    }

    pub fn net(&mut self, e: NetEvent) {
        if !self.trace.net.contains(&e) && self.room(Self::cost(&e)) {
            self.trace.net.insert(e);
        }
    }

    pub fn file(&mut self, e: FileEvent) {
        if !self.trace.files.contains(&e) && self.room(Self::cost(&e)) {
            self.trace.files.insert(e);
        }
    }

    pub fn syscall(&mut self, program: &str, f: SyscallFact) {
        let n = match self.trace.syscalls.get(program) {
            Some(set) if set.contains(&f) => return,
            Some(_) => Self::cost(&f),
            // The program's key and its set's brackets too.
            None => Self::cost(&f) + Self::cost(program) + 3,
        };
        if self.room(n) {
            self.trace
                .syscalls
                .entry(program.to_string())
                .or_default()
                .insert(f);
        }
    }

    pub fn unknown(&mut self, u: String) {
        if !self.trace.unknown.contains(&u) && self.room(Self::cost(&u)) {
            self.trace.unknown.insert(u);
        }
    }

    pub fn notification(&mut self) {
        self.trace.notifications += 1;
    }

    /// The trace so far, to set the end fields on.
    pub fn trace_mut(&mut self) -> &mut Trace {
        &mut self.trace
    }

    /// Mark the trace incomplete (the reason cut to 500 characters).
    pub fn incomplete(&mut self, why: &str) {
        let mut w: String = why.chars().take(500).collect();
        if w.is_empty() {
            w = "unknown".into();
        }
        self.trace.status = Status::Incomplete;
        self.trace.incomplete = w;
    }

    /// The finished document: canonical and within `max_bytes` (entries are
    /// dropped from the largest set until it fits, should the reserve not
    /// have been enough).
    pub fn finish(mut self) -> (Trace, Vec<u8>) {
        loop {
            let b = self.trace.to_canonical();
            if b.len() as u64 <= self.trace.max_bytes {
                return (self.trace, b);
            }
            if !self.pop_largest() {
                // Only the header is left and it is still too large: the argv.
                let t = &mut self.trace;
                t.argv.truncate(1);
                t.argv[0] = t.argv[0].chars().take(200).collect();
                let b = self.trace.to_canonical();
                return (self.trace, b);
            }
            self.drop_one();
        }
    }

    fn pop_largest(&mut self) -> bool {
        let t = &mut self.trace;
        let sizes = [
            t.files.len(),
            t.exec.len(),
            t.net.len(),
            t.syscalls.values().map(|s| s.len()).sum(),
            t.unknown.len(),
        ];
        let largest = (0..sizes.len()).max_by_key(|&i| sizes[i]).unwrap_or(0);
        if sizes[largest] == 0 {
            return false;
        }
        match largest {
            0 => t.files.pop_last().is_some(),
            1 => t.exec.pop_last().is_some(),
            2 => t.net.pop_last().is_some(),
            3 => {
                let Some(k) = t
                    .syscalls
                    .iter()
                    .max_by_key(|(_, s)| s.len())
                    .map(|(k, _)| k.clone())
                else {
                    return false;
                };
                let s = t.syscalls.get_mut(&k).expect("present");
                s.pop_last();
                if s.is_empty() {
                    t.syscalls.remove(&k);
                }
                true
            }
            _ => t.unknown.pop_last().is_some(),
        }
    }
}