nordisk/heimdall

crates/heimdall-probe/src/learn.rs

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

//! Learning mode (spec §3.1 step 2): run one program under a filter whose
//! default is `SECCOMP_RET_USER_NOTIF`, record every call, and let it run
//! (`FLAG_CONTINUE`). Heimdall is a defensive sandbox: no strace, no ptrace,
//! pure Rust. The decision to continue uses registers only, so the TOCTOU
//! caveat of `Continue` does not apply: the probe allows everything anyway.
//!
//! Kernel: USER_NOTIF >= 5.0 and `FLAG_CONTINUE` >= 5.5. There is no ptrace
//! fallback for older kernels (spec §3.1 allows one; not built: every kernel
//! Heimdall targets has 5.5).

use std::ffi::OsString;
use std::path::PathBuf;
use std::time::{Duration, Instant};

use heimdall_seccomp::launch::{Exit, Launch, Stdio};
use heimdall_seccomp::notif::{Notif, Readiness, Response};
use heimdall_seccomp::{Action, Arch};

use crate::facts::{fixed_rules, spec};
use crate::observe::{Call, Observations};

/// One run of the probe corpus.
#[derive(Clone, Debug)]
pub struct ProbeRun {
    /// The test name the drift check reports.
    pub name: String,
    pub program: PathBuf,
    pub args: Vec<OsString>,
    /// The whole environment (the probe starts from an empty one).
    pub env: Vec<(OsString, OsString)>,
    pub cwd: Option<PathBuf>,
    pub stdin: Option<PathBuf>,
    /// Where stdout and stderr go (`None`: /dev/null).
    pub output: Option<PathBuf>,
    pub error_path: bool,
    pub timeout: Duration,
}

impl ProbeRun {
    pub fn new(name: &str, program: impl Into<PathBuf>) -> ProbeRun {
        ProbeRun {
            name: name.to_string(),
            program: program.into(),
            args: Vec::new(),
            env: Vec::new(),
            cwd: None,
            stdin: None,
            output: None,
            error_path: false,
            timeout: Duration::from_secs(60),
        }
    }
}

/// How a probe run ended.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct RunOutcome {
    pub exit: Exit,
    pub timed_out: bool,
    /// How many calls the probe saw.
    pub calls: u64,
}

/// The call one notification describes (unknown numbers and foreign arches
/// give `Err` with a description).
pub fn call_of(arch: Arch, n: &Notif) -> Result<Call, String> {
    if n.data.arch != arch.audit_arch() {
        return Err(format!("arch {:#x} nr {}", n.data.arch, n.data.nr));
    }
    let Some(name) = u32::try_from(n.data.nr)
        .ok()
        .and_then(|nr| arch.syscall_name(nr))
    else {
        return Err(format!("nr {}", n.data.nr));
    };
    let mut args = Vec::new();
    if let Some(s) = spec(name) {
        for &(i, w) in s.args {
            let v = n.data.args[usize::from(i)];
            let v = match w {
                heimdall_seccomp::Width::W32 => v & 0xffff_ffff,
                heimdall_seccomp::Width::W64 => v,
            };
            args.push((i, v));
        }
    }
    Ok(Call {
        syscall: name.to_string(),
        args,
    })
}

/// Run `run` in learning mode, recording into `obs`.
pub fn learn(run: &ProbeRun, obs: &mut Observations) -> std::io::Result<RunOutcome> {
    let arch = obs.arch;
    let open = |p: &PathBuf, write: bool| -> std::io::Result<Stdio> {
        let f = if write {
            std::fs::OpenOptions::new()
                .create(true)
                .append(true)
                .open(p)?
        } else {
            std::fs::File::open(p)?
        };
        Ok(Stdio::Fd(f.into()))
    };
    let mut l = Launch::new(&run.program)
        .args(&run.args)
        .env_clear()
        .rules(arch, fixed_rules(), Action::UserNotif)
        .listener(true);
    for (k, v) in &run.env {
        l = l.env(k, v);
    }
    if let Some(d) = &run.cwd {
        l = l.cwd(d);
    }
    l = l.stdin(match &run.stdin {
        Some(p) => open(p, false)?,
        None => Stdio::Null,
    });
    match &run.output {
        Some(p) => {
            l = l.stdout(open(p, true)?).stderr(open(p, true)?);
        }
        None => {
            l = l.stdout(Stdio::Null).stderr(Stdio::Null);
        }
    }
    let mut child = l.spawn()?;
    let listener = child
        .take_listener()
        .ok_or_else(|| std::io::Error::other("no listener"))?;
    let deadline = Instant::now() + run.timeout;
    let mut calls = 0u64;
    let mut timed_out = false;
    loop {
        let left = deadline.saturating_duration_since(Instant::now());
        if left.is_zero() {
            timed_out = true;
            let _ = child.signal(libc::SIGKILL);
            break;
        }
        match listener.wait(Some(left.min(Duration::from_secs(1))))? {
            Readiness::Hangup => break,
            Readiness::Timeout => continue,
            Readiness::Ready => {}
        }
        let n = match listener.recv() {
            Ok(n) => n,
            Err(e) if e.raw_os_error() == Some(libc::ENOENT) => continue,
            Err(e) => return Err(e),
        };
        calls += 1;
        match call_of(arch, &n) {
            Ok(c) => obs.record(c, &run.name),
            Err(u) => {
                obs.unknown.entry(u).or_default().insert(run.name.clone());
            }
        }
        match listener.respond(n.id, Response::Continue) {
            Ok(()) => {}
            Err(e) if e.raw_os_error() == Some(libc::ENOENT) => {}
            Err(e) => return Err(e),
        }
    }
    let exit = child.wait()?;
    obs.runs.insert(crate::observe::RunInfo {
        name: run.name.clone(),
        error_path: run.error_path,
        exit: if timed_out {
            "timed out".to_string()
        } else {
            exit.to_string()
        },
    });
    Ok(RunOutcome {
        exit,
        timed_out,
        calls,
    })
}