use anyhow::{Context, anyhow, ensure};
use std::path::Path;
pub const WPROF_PB_MIN_BYTES: usize = 4096;
pub const PERFETTO_TRACE_PACKETS_TAG: u8 = 0x0a;
pub(crate) fn reassemble_wprof_trace(entries: &[crate::vmm::wire::ShmEntry]) -> Option<Vec<u8>> {
use crate::vmm::wire::MsgType;
let mut out: Vec<u8> = Vec::new();
let mut saw_terminal = false;
for e in entries {
match MsgType::from_wire(e.msg_type) {
Some(MsgType::WprofTraceChunk) => {
if !e.crc_ok {
return None;
}
out.extend_from_slice(&e.payload);
}
Some(MsgType::WprofTrace) => {
if !e.crc_ok {
return None;
}
out.extend_from_slice(&e.payload);
saw_terminal = true;
break;
}
_ => {}
}
}
(saw_terminal && !out.is_empty()).then_some(out)
}
pub fn assert_wprof_pb_shape(path: &Path) -> anyhow::Result<()> {
let bytes = std::fs::read(path).with_context(|| {
format!(
"wprof .pb missing at {}. Chain: #[ktstr_test(wprof)] → \
KtstrTestEntry::wprof → primary VM builder.wprof(Some(config)) \
at src/test_support/eval/mod.rs → KTSTR_WPROF_ARGS cmdline → \
guest spawn_wprof_if_configured → send_wprof_trace → host \
MsgType::WprofTrace arm → \
sidecar_dir.join(\"<name>-<variant_hash:016x>.wprof.pb\")",
path.display(),
)
})?;
ensure!(
bytes.len() >= WPROF_PB_MIN_BYTES,
"wprof .pb at {} is only {} bytes — expected >= {WPROF_PB_MIN_BYTES}. \
wprof's init_pb_trace emits a ~4 KB interned-string table on every \
capture; a smaller file means wprof either aborted before \
init_pb_trace or the .pb write/transport truncated.",
path.display(),
bytes.len(),
);
let first = bytes[0];
if first != PERFETTO_TRACE_PACKETS_TAG {
return Err(anyhow!(
"wprof .pb at {} first byte {first:#04x} — expected \
{PERFETTO_TRACE_PACKETS_TAG:#04x} (field=1, wire_type=2, the \
Perfetto `Trace.packets` repeated TracePacket tag). File may \
be truncated, in a different format, or corrupted.",
path.display(),
));
}
Ok(())
}
#[cfg(all(test, feature = "wprof"))]
mod tests {
use super::reassemble_wprof_trace;
use crate::vmm::guest_comms::wprof_trace_frames;
use crate::vmm::wire::{MsgType, ShmEntry};
fn roundtrip(buf: &[u8], cap: usize) -> Option<Vec<u8>> {
let entries: Vec<ShmEntry> = wprof_trace_frames(buf, cap)
.into_iter()
.map(|(msg_type, chunk)| ShmEntry {
msg_type,
payload: chunk.to_vec(),
crc_ok: true,
})
.collect();
reassemble_wprof_trace(&entries)
}
#[test]
fn frames_tag_all_but_last_as_chunk() {
let buf: Vec<u8> = (0..10u8).collect();
let frames = wprof_trace_frames(&buf, 4);
assert_eq!(frames.len(), 3);
assert_eq!(frames[0].0, MsgType::WprofTraceChunk.wire_value());
assert_eq!(frames[1].0, MsgType::WprofTraceChunk.wire_value());
assert_eq!(frames[2].0, MsgType::WprofTrace.wire_value());
assert_eq!(frames[0].1, &[0, 1, 2, 3]);
assert_eq!(frames[1].1, &[4, 5, 6, 7]);
assert_eq!(frames[2].1, &[8, 9]);
}
#[test]
fn single_frame_is_lone_terminal() {
let buf: Vec<u8> = (0..4u8).collect();
let frames = wprof_trace_frames(&buf, 4);
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].0, MsgType::WprofTrace.wire_value());
}
#[test]
fn roundtrip_small_cap_byte_identical() {
for len in [0usize, 1, 3, 4, 5, 8, 9, 100] {
let buf: Vec<u8> = (0..len).map(|i| (i % 251) as u8).collect();
let out = roundtrip(&buf, 4);
if buf.is_empty() {
assert_eq!(out, None, "empty trace reassembles to None");
} else {
assert_eq!(out.as_deref(), Some(buf.as_slice()), "len={len}");
}
}
}
#[test]
fn boundary_at_real_cap() {
let cap = crate::vmm::bulk::MAX_BULK_FRAME_PAYLOAD as usize;
let exact = vec![0xABu8; cap];
let frames = wprof_trace_frames(&exact, cap);
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].0, MsgType::WprofTrace.wire_value());
assert_eq!(roundtrip(&exact, cap).as_deref(), Some(exact.as_slice()));
let mut over = vec![0xCDu8; cap];
over.push(0xEE);
let frames = wprof_trace_frames(&over, cap);
assert_eq!(frames.len(), 2);
assert_eq!(frames[0].0, MsgType::WprofTraceChunk.wire_value());
assert_eq!(frames[0].1.len(), cap);
assert_eq!(frames[1].0, MsgType::WprofTrace.wire_value());
assert_eq!(frames[1].1, &[0xEE]);
assert_eq!(roundtrip(&over, cap).as_deref(), Some(over.as_slice()));
}
#[test]
fn torn_chunk_bails_to_none() {
let entries = vec![
ShmEntry {
msg_type: MsgType::WprofTraceChunk.wire_value(),
payload: vec![1, 2, 3],
crc_ok: false,
},
ShmEntry {
msg_type: MsgType::WprofTrace.wire_value(),
payload: vec![4, 5],
crc_ok: true,
},
];
assert_eq!(reassemble_wprof_trace(&entries), None);
}
#[test]
fn torn_terminal_bails_to_none() {
let entries = vec![
ShmEntry {
msg_type: MsgType::WprofTraceChunk.wire_value(),
payload: vec![1, 2, 3],
crc_ok: true,
},
ShmEntry {
msg_type: MsgType::WprofTrace.wire_value(),
payload: vec![4, 5],
crc_ok: false,
},
];
assert_eq!(reassemble_wprof_trace(&entries), None);
}
#[test]
fn no_terminal_returns_none() {
let entries = vec![ShmEntry {
msg_type: MsgType::WprofTraceChunk.wire_value(),
payload: vec![1, 2, 3],
crc_ok: true,
}];
assert_eq!(reassemble_wprof_trace(&entries), None);
}
#[test]
fn unrelated_crc_bad_frame_ignored() {
let entries = vec![
ShmEntry {
msg_type: MsgType::Stdout.wire_value(),
payload: vec![9, 9, 9],
crc_ok: false,
},
ShmEntry {
msg_type: MsgType::WprofTrace.wire_value(),
payload: vec![4, 5],
crc_ok: true,
},
];
assert_eq!(
reassemble_wprof_trace(&entries).as_deref(),
Some(&[4u8, 5][..])
);
}
fn chunk(payload: &[u8], crc_ok: bool) -> ShmEntry {
ShmEntry {
msg_type: MsgType::WprofTraceChunk.wire_value(),
payload: payload.to_vec(),
crc_ok,
}
}
fn terminal(payload: &[u8], crc_ok: bool) -> ShmEntry {
ShmEntry {
msg_type: MsgType::WprofTrace.wire_value(),
payload: payload.to_vec(),
crc_ok,
}
}
#[test]
fn torn_middle_chunk_in_n_gt_2_bails() {
let entries = vec![
chunk(&[1, 2], true),
chunk(&[3, 4], false), chunk(&[5, 6], true),
terminal(&[7, 8], true),
];
assert_eq!(reassemble_wprof_trace(&entries), None);
}
#[test]
fn interleaved_crc_good_nonwprof_frame_skipped_preserves_order() {
let entries = vec![
chunk(&[1, 2], true),
ShmEntry {
msg_type: MsgType::Stdout.wire_value(),
payload: vec![9, 9],
crc_ok: true,
},
chunk(&[3, 4], true),
terminal(&[5], true),
];
assert_eq!(
reassemble_wprof_trace(&entries).as_deref(),
Some(&[1, 2, 3, 4, 5][..])
);
}
#[test]
fn duplicate_terminal_stops_at_first() {
let entries = vec![
chunk(&[1], true),
terminal(&[2], true),
terminal(&[3], true),
];
assert_eq!(
reassemble_wprof_trace(&entries).as_deref(),
Some(&[1u8, 2][..])
);
}
#[test]
fn chunk_after_terminal_ignored() {
let entries = vec![chunk(&[1], true), terminal(&[2], true), chunk(&[3], true)];
assert_eq!(
reassemble_wprof_trace(&entries).as_deref(),
Some(&[1u8, 2][..])
);
}
}