cubelinux-2/cubetrace/src/lib.rs

//! CUBELinux-2 tracing package.
//!
//! Per the PDF (Package 4, §547 / §1113): *"cubetrace: wraps a DBI engine
//! (via FFI) and streams trace events (basic blocks, syscalls) into
//! cubestore, tagging each with CZYX coordinates and header flags."*
//!
//! This crate is the *head* of the Package‑4 pipeline (cubetrace → cubeai →
//! cubedbt). It is dependency-free and exercises today on a `Null` (offline)
//! source so it builds and is testable without an external DBI engine; the
//! real engine is reached through the documented `DbiEngine` FFI seam.
//!
//! Model
//! -----
//! * A [`TraceEvent`] is a basic block / syscall / metadata observation with a
//!   timestamp, a target coordinate, and policy/behavior header flags.
//! * A [`Tracer`] streams events into a [`CubeStore`] keyed by CZYX (the
//!   `c240` trace band), tagging each record's header with time, environment,
//!   and policy flags (§976).
//! * `replay_all` deterministically replays the stored stream in timestamp
//!   order (§977/§1118: "deterministic replay first").

use cubecode::{Behavior, Kind};
use cubecoords::{CubeHeader, Czyx};
use cubestore::{CubeBackend, CubeStore, HashBackend};

/// `C` axis band where trace events are stored (keyed CZYX, §975).
pub const C_TRACE: u8 = 240;

/// Kind of traced observation.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum EventKind {
    /// A basic block executed (carries the block's code/length).
    BasicBlock,
    /// A syscall entered/exited (carries the syscall number).
    Syscall,
    /// Free-form metadata about the run (policy, environment).
    Meta,
}

/// A single traced observation: what happened, when, where, and with what
/// policy/behavior tags. Serialized into a trace record's body.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TraceEvent {
    pub kind: EventKind,
    /// Monotonic timestamp (ns from trace start).
    pub ts: u64,
    /// Target CZYX the event is about (the block/syscall coordinate).
    pub coord: Czyx,
    /// For `BasicBlock`: the executed bytecode; for `Syscall`: the syscall
    /// number; for `Meta`: unused (0).
    pub payload: Vec<u8>,
    /// Behavior/policy descriptors to stamp onto the record header.
    pub behavior: Behavior,
}

impl TraceEvent {
    /// Serialize to a stable byte form (no external deps).
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut out = Vec::new();
        out.push(self.kind as u8);
        out.extend_from_slice(&self.ts.to_le_bytes());
        out.extend_from_slice(&CZYX_BYTES);
        out.extend_from_slice(&self.coord.pack_u32().to_le_bytes());
        out.extend_from_slice(&self.behavior.to_flags().to_le_bytes());
        out.push(self.payload.len() as u8);
        out.extend_from_slice(&self.payload);
        out
    }

    /// Inverse of [`to_bytes`]. Returns `None` on malformed input.
    pub fn from_bytes(b: &[u8]) -> Option<TraceEvent> {
        let mut i = 0;
        let kind = match *b.get(i)? {
            0 => EventKind::BasicBlock,
            1 => EventKind::Syscall,
            2 => EventKind::Meta,
            _ => return None,
        };
        i += 1;
        let ts = u64::from_le_bytes(b.get(i..i + 8)?.try_into().ok()?);
        i += 8;
        i += 4; // skip CZYX_BYTES sentinel
        let packed = u32::from_le_bytes(b.get(i..i + 4)?.try_into().ok()?);
        i += 4;
        let coord = Czyx::new(
            (packed >> 24) as u8,
            (packed >> 16) as u8,
            (packed >> 8) as u8,
            packed as u8,
        );
        let flags = u16::from_le_bytes(b.get(i..i + 2)?.try_into().ok()?);
        i += 2;
        let plen = *b.get(i)? as usize;
        i += 1;
        let payload = b.get(i..i + plen)?.to_vec();
        Some(TraceEvent {
            kind,
            ts,
            coord,
            payload,
            behavior: Behavior::from_flags(flags),
        })
    }
}

/// Sentinel to make the on-wire format self-identifying as a CZYX record.
const CZYX_BYTES: [u8; 4] = *b"CZYX";

/// Source of trace events. In the full stack this is the FFI seam to a DBI
/// engine (DynamoRIO / Intel Pin / QBDI / Frida — §970). Here it is a trait so
/// the crate is testable offline via [`NullEngine`] and a real engine can be
/// dropped in without changing callers.
pub trait DbiEngine {
    /// Attach to a target (pid or path). Returns `Err` if the engine can't.
    fn attach(&mut self, target: &str) -> Result<(), String>;
    /// Pull the next event, or `None` when the run is exhausted.
    fn next_event(&mut self) -> Option<TraceEvent>;
}

/// Offline engine: replays a pre-recorded event vector. Stands in for a live
/// DBI backend in tests and headless environments.
pub struct NullEngine {
    events: Vec<TraceEvent>,
    idx: usize,
}

impl NullEngine {
    pub fn new(events: Vec<TraceEvent>) -> Self {
        NullEngine { events, idx: 0 }
    }
}

impl DbiEngine for NullEngine {
    fn attach(&mut self, _target: &str) -> Result<(), String> {
        Ok(())
    }
    fn next_event(&mut self) -> Option<TraceEvent> {
        if self.idx < self.events.len() {
            let e = self.events[self.idx].clone();
            self.idx += 1;
            Some(e)
        } else {
            None
        }
    }
}

/// Streams trace events into a [`CubeStore`] (the `c240` band), tagging each
/// record with time, environment, and policy header flags (PDF §976).
pub struct Tracer<B: CubeBackend> {
    store: CubeStore<B>,
    next_x: u8,
}

impl<B: CubeBackend> Tracer<B> {
    pub fn new(store: CubeStore<B>) -> Self {
        Tracer { store, next_x: 1 }
    }

    /// Drain an engine into the store, returning the number of events stored.
    pub fn run<E: DbiEngine>(&mut self, engine: &mut E) -> usize {
        let mut count = 0;
        while let Some(ev) = engine.next_event() {
            self.store_event(&ev);
            count += 1;
        }
        count
    }

    /// Store one event under a fresh CZYX coordinate in the `c240` band.
    pub fn store_event(&mut self, ev: &TraceEvent) -> Czyx {
        let label = Czyx::new(C_TRACE, 1, ev.kind as u8, self.next_x);
        self.next_x = self.next_x.wrapping_add(1).max(1);
        let mut h = CubeHeader::new();
        h.title = Some(format!("{:?}:{:?}", ev.kind, ev.coord));
        h.doc_type = Some(Kind::Other.as_str().into());
        h.size_bytes = Some(ev.to_bytes().len() as u64);
        h.flags.0 |= ev.behavior.to_flags();
        h.refresh_flags();
        self.store.put_record(label, &h, &ev.to_bytes());
        label
    }

    /// Borrow the backing store (e.g. to hand to `cubeai`/others).
    pub fn store(&self) -> &CubeStore<B> {
        &self.store
    }
}

/// Deterministic replay: collect every trace event from the `c240` band and
/// return them in timestamp order (PDF §977/§1118: "deterministic replay
/// first"). The trace is the immutable cube record; replay never mutates it.
pub fn replay_all(store: &CubeStore<HashBackend>) -> Vec<TraceEvent> {
    let mut events = Vec::new();
    for k in store.keys() {
        if k.c != C_TRACE {
            continue;
        }
        if let Some((_, body)) = store.get_record(&k) {
            if let Some(ev) = TraceEvent::from_bytes(&body) {
                events.push(ev);
            }
        }
    }
    events.sort_by_key(|e| e.ts);
    events
}

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

    fn ev(kind: EventKind, ts: u64, coord: Czyx, beh: Behavior) -> TraceEvent {
        TraceEvent {
            kind,
            ts,
            coord,
            payload: vec![kind as u8],
            behavior: beh,
        }
    }

    fn behavior(flag: u16) -> Behavior {
        Behavior(flag)
    }

    #[test]
    fn event_round_trips_through_bytes() {
        let e = ev(
            EventKind::Syscall,
            42,
            Czyx::new(1, 2, 3, 4),
            behavior(Behavior::HOT_PATH),
        );
        let back = TraceEvent::from_bytes(&e.to_bytes()).expect("decodes");
        assert_eq!(e, back);
    }

    #[test]
    fn tracer_streams_to_c240() {
        let mut engine = NullEngine::new(vec![
            ev(
                EventKind::BasicBlock,
                10,
                Czyx::new(1, 0, 0, 1),
                behavior(Behavior::PURE),
            ),
            ev(
                EventKind::Syscall,
                20,
                Czyx::new(1, 0, 0, 2),
                behavior(Behavior::IO_HEAVY),
            ),
            ev(
                EventKind::Meta,
                5,
                Czyx::new(1, 0, 0, 3),
                Behavior::default(),
            ),
        ]);
        let mut tracer = Tracer::new(CubeStore::new(HashBackend::new()));
        assert_eq!(tracer.run(&mut engine), 3);
        // All records landed in the c240 band.
        let stored: usize = tracer
            .store()
            .keys()
            .into_iter()
            .filter(|k| k.c == C_TRACE)
            .count();
        assert_eq!(stored, 3);
    }

    #[test]
    fn replay_is_deterministic_and_ordered() {
        let mut engine = NullEngine::new(vec![
            ev(
                EventKind::BasicBlock,
                30,
                Czyx::new(1, 0, 0, 1),
                behavior(Behavior::PURE),
            ),
            ev(
                EventKind::Syscall,
                10,
                Czyx::new(1, 0, 0, 2),
                behavior(Behavior::IO_HEAVY),
            ),
            ev(
                EventKind::Meta,
                20,
                Czyx::new(1, 0, 0, 3),
                Behavior::default(),
            ),
        ]);
        let mut tracer = Tracer::new(CubeStore::new(HashBackend::new()));
        tracer.run(&mut engine);
        let replayed = replay_all(tracer.store());
        assert_eq!(replayed.len(), 3);
        // Sorted by ts: 10 (Syscall), 20 (Meta), 30 (BasicBlock).
        assert_eq!(replayed[0].kind, EventKind::Syscall);
        assert_eq!(replayed[1].kind, EventKind::Meta);
        assert_eq!(replayed[2].kind, EventKind::BasicBlock);
        // Behavior flags survived the round-trip through the store.
        assert_eq!(replayed[0].behavior, Behavior(Behavior::IO_HEAVY));
    }
}