katra-trace 0.1.0

//! Trace reader: iterates records from a trace file with validation.

use std::fs::File;
use std::io::{BufReader, Read};
use std::path::Path;

use katra_core::{KatraError, Result};

use crate::header::{FORMAT_VERSION, HostInfo, SCHEMA_HASH, TRACE_MAGIC, TraceHeader};
use crate::record::TraceRecord;

/// Errors surfaced while opening or reading a trace.
#[derive(Debug)]
pub struct TraceReaderError {
    /// Human-readable description.
    pub message: String,
}

/// Reads records from a trace file.
pub struct TraceReader {
    inner: BufReader<File>,
    header: Option<TraceHeader>,
    count: u64,
}

impl TraceReader {
    /// Open and validate the header of a trace file.
    pub fn open(path: &Path) -> Result<Self> {
        let file = File::open(path).map_err(KatraError::from)?;
        let mut r = TraceReader { inner: BufReader::new(file), header: None, count: 0 };
        r.read_header()?;
        Ok(r)
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
        self.inner.read_exact(buf).map_err(|e| KatraError::Corrupt(format!("truncated trace: {e}")))
    }

    fn read_header(&mut self) -> Result<()> {
        let mut magic = [0u8; 8];
        self.read_exact(&mut magic)?;
        if magic != TRACE_MAGIC {
            return Err(KatraError::Corrupt(format!(
                "bad magic: {magic:02x?} (expected {TRACE_MAGIC:02x?})"
            )));
        }
        let mut version_bytes = [0u8; 4];
        self.read_exact(&mut version_bytes)?;
        let version = u32::from_le_bytes(version_bytes);
        if version != FORMAT_VERSION {
            return Err(KatraError::Protocol(format!(
                "trace format version {version} != supported {FORMAT_VERSION}"
            )));
        }
        let mut len_bytes = [0u8; 4];
        self.read_exact(&mut len_bytes)?;
        let len = u32::from_le_bytes(len_bytes) as usize;
        let mut buf = vec![0u8; len];
        self.read_exact(&mut buf)?;
        let header: TraceHeader = bincode::deserialize(&buf)
            .map_err(|e| KatraError::Corrupt(format!("bad header: {e}")))?;
        if header.schema_hash != SCHEMA_HASH {
            return Err(KatraError::Protocol(format!(
                "schema hash {} != expected {}; trace was produced by an incompatible build",
                header.schema_hash, SCHEMA_HASH
            )));
        }
        self.header = Some(header);
        Ok(())
    }

    /// The validated header.
    pub fn header(&self) -> Result<&TraceHeader> {
        self.header.as_ref().ok_or_else(|| KatraError::Protocol("header not read".into()))
    }

    /// Read the next record, or `None` at end of stream.
    pub fn next_record(&mut self) -> Result<Option<TraceRecord>> {
        let mut len_bytes = [0u8; 4];
        match self.inner.read_exact(&mut len_bytes) {
            Ok(()) => {}
            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
            Err(e) => return Err(KatraError::from(e)),
        }
        let len = u32::from_le_bytes(len_bytes) as usize;
        let mut buf = vec![0u8; len];
        self.read_exact(&mut buf)?;
        let record: TraceRecord = bincode::deserialize(&buf)
            .map_err(|e| KatraError::Corrupt(format!("bad record: {e}")))?;
        self.count += 1;
        Ok(Some(record))
    }

    /// Number of records read so far.
    pub fn count(&self) -> u64 {
        self.count
    }
}

/// Collect basic host information for the trace header, from `/proc`
/// (Linux) or defaults elsewhere. Best-effort; never fails.
pub fn host_info() -> HostInfo {
    HostInfo {
        os: Some(std::env::consts::OS.to_string()),
        kernel: read_first_line("/proc/sys/kernel/osrelease"),
        cpu_model: read_proc_cpuinfo_model(),
        cpu_cores: std::thread::available_parallelism().ok().map(|n| n.get() as u32),
        ram_bytes: read_proc_meminfo_total(),
    }
}

fn read_first_line(path: &str) -> Option<String> {
    std::fs::read_to_string(path)
        .ok()
        .and_then(|s| s.lines().next().map(|l| l.trim().to_string()))
        .filter(|s| !s.is_empty())
}

fn read_proc_cpuinfo_model() -> Option<String> {
    let text = std::fs::read_to_string("/proc/cpuinfo").ok()?;
    for line in text.lines() {
        if let Some((key, value)) = line.split_once(':') {
            if key.trim() == "model name" {
                let v = value.trim().to_string();
                if !v.is_empty() {
                    return Some(v);
                }
            }
        }
    }
    None
}

fn read_proc_meminfo_total() -> Option<u64> {
    let text = std::fs::read_to_string("/proc/meminfo").ok()?;
    for line in text.lines() {
        if let Some(rest) = line.strip_prefix("MemTotal:") {
            let kb: u64 = rest.trim().trim_end_matches("kB").trim().parse().ok()?;
            return Some(kb * 1024);
        }
    }
    None
}