crates/heimdall-wasm/src/split.rs
// SPDX-License-Identifier: AGPL-3.0-only
//! The split guest's decoder side (docs/GATLING-WASM.md). Heimdall is a
//! defensive sandbox.
//!
//! A split guest is two components: a VERIFIER (world `verifier`) that never
//! decodes, and a DECODER (world `decoder`) that only decodes. When the
//! verifier calls `decode-chunks`, the host
//!
//! 1. cuts the chunks into UNITS, deterministically from the chunk sizes alone
//! (never from the worker count), so the verdict, the fuel and the
//! refusal codes are the same for every N;
//! 2. runs every unit in a FRESH decoder instance — its own store, linear
//! memory, fuel and deadline (the job's limits, per instance; the deadline is
//! what is left of the job's wall clock) — handed only that unit's frames,
//! lowered straight from the verifier's memory (no copy on the host side);
//! 3. drives the units over N workers with GATLING (`gatling_forkjoin` when
//! only digests come back, `gatling::ordered` when bytes do: in order, with
//! at most 2N units in flight, stopping at the first unit that fails);
//! 4. checks every answer NATIVELY before the verifier sees it: one answer per
//! chunk, a 32-byte digest, exactly the declared size, and BLAKE3(bytes) =
//! digest. A decoder that lies ends the job with
//! [`codes::DECODER_MISMATCH`]; one that traps or hits its limits ends it
//! with that limit's code. The first failing unit in chunk order decides.
//!
//! The verifier then holds every digest to the archive's own checksums and
//! the signed manifest; it never shares memory with a decoder.
use std::ops::Range;
use std::sync::Arc;
use std::time::{Duration, Instant};
use wasmtime::component::{ComponentType, InstancePre, Lift, Lower, WasmList};
use wasmtime::Store;
use crate::codes;
use crate::host::{classify_error, one_line, Limiter, Ticker};
use crate::limits::RunLimits;
/// A unit closes once its frames plus declared sizes reach this many bytes.
pub const UNIT_BYTES: u64 = 8 << 20;
/// A unit holds at most this many chunks.
pub const UNIT_CHUNKS: usize = 64;
/// A decoder's "does not decode" text is cut to this many bytes.
const MAX_DECODER_TEXT: usize = 400;
/// `chunk-ref` as the host lowers it into a decoder instance: the frame is
/// lent, never copied on the host side.
#[derive(ComponentType, Lower, Clone, Copy)]
#[component(record)]
pub(crate) struct ChunkIn<'a> {
frame: &'a [u8],
size: u64,
#[component(name = "want-bytes")]
want_bytes: bool,
}
impl<'a> ChunkIn<'a> {
/// The host always asks a decoder for the bytes: it hashes them itself.
pub(crate) fn new(frame: &'a [u8], size: u64) -> Self {
ChunkIn {
frame,
size,
want_bytes: true,
}
}
}
/// `chunk-ref` as the host lifts it from the verifier.
#[derive(ComponentType, Lift)]
#[component(record)]
pub(crate) struct ChunkRefLifted {
pub(crate) frame: WasmList<u8>,
pub(crate) size: u64,
#[component(name = "want-bytes")]
pub(crate) want_bytes: bool,
}
/// `decoded`: what one chunk decoded to.
#[derive(ComponentType, Lift, Lower, Debug, Clone, PartialEq, Eq)]
#[component(record)]
pub struct DecodedChunk {
/// BLAKE3-256 of `bytes` (checked by the host).
pub digest: Vec<u8>,
/// The decoded bytes; empty when only the digest was asked for.
pub bytes: Vec<u8>,
}
/// One answer per chunk: the decoded chunk, or why its frame does not decode.
pub type ChunkAnswer = Result<DecodedChunk, String>;
/// Why decoding ended the job: a decoder instance trapped, hit one of its own
/// limits, or answered falsely.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DecodeFailure {
/// A stable `heimdall-wasm-<n>` code.
pub code: &'static str,
pub reason: String,
}
/// Every unit's outcome, with the fuel the decoder instances burnt (counted
/// up to and including the first failing unit, so it is the same for any N).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DecodeRun {
pub result: Result<Vec<ChunkAnswer>, DecodeFailure>,
pub fuel_used: u64,
/// How many decoder instances ran to an answer that counted.
pub units: usize,
}
/// Per-instance store data: only the limiter (a decoder has no imports).
pub(crate) struct DState {
pub(crate) limiter: Limiter,
}
/// A verified, compiled decoder component (world `decoder`). Every unit of
/// work runs in a fresh instance of it.
pub struct Decoder {
pub(crate) pre: InstancePre<DState>,
pub(crate) digest_hex: String,
pub(crate) signer: String,
pub(crate) stack_bytes: usize,
pub(crate) epoch_tick: Duration,
pub(crate) _ticker: Arc<Ticker>,
}
/// The unit ranges for chunks of these `(frame length, declared size)`.
pub fn plan_units(chunks: impl IntoIterator<Item = (usize, u64)>) -> Vec<Range<usize>> {
let mut out = Vec::new();
let mut start = 0usize;
let mut cost = 0u64;
let mut i = 0usize;
for (frame, size) in chunks {
cost = cost.saturating_add(frame as u64).saturating_add(size);
i += 1;
if cost >= UNIT_BYTES || i - start >= UNIT_CHUNKS {
out.push(start..i);
start = i;
cost = 0;
}
}
if start < i {
out.push(start..i);
}
out
}
type UnitOut = (Result<Vec<ChunkAnswer>, DecodeFailure>, u64);
impl Decoder {
/// Lowercase hex BLAKE3-256 of the decoder module that was verified.
pub fn digest_hex(&self) -> &str {
&self.digest_hex
}
/// The common name of the certificate that signed the decoder module.
pub fn signer(&self) -> &str {
&self.signer
}
/// Decode `chunks` (`(frame, declared size)`) exactly as a verifier's
/// `decode-chunks` call is served, bytes returned, under `limits` (its
/// `decoders` workers; the wall clock starts now). For tests and fuzzing.
pub fn decode(&self, chunks: &[(&[u8], u64)], limits: &RunLimits) -> DecodeRun {
let ins: Vec<ChunkIn<'_>> = chunks
.iter()
.map(|&(frame, size)| ChunkIn {
frame,
size,
want_bytes: true,
})
.collect();
let wants = vec![true; ins.len()];
self.run(&ins, &wants, limits, Instant::now() + limits.timeout)
}
/// Serve one `decode-chunks` call: plan units, fan them out, check them.
/// `ins` always ask the decoder for bytes (the host hashes them); `wants`
/// says which answers carry their bytes back to the caller.
pub(crate) fn run(
&self,
ins: &[ChunkIn<'_>],
wants: &[bool],
limits: &RunLimits,
deadline: Instant,
) -> DecodeRun {
let units = plan_units(ins.iter().map(|c| (c.frame.len(), c.size)));
let workers = limits.decoders.max(1);
let unit = |u: &Range<usize>| -> UnitOut {
let (r, fuel) = self.run_unit(&ins[u.clone()], limits, deadline, u.start);
let r = r.map(|mut answers| {
for (a, &want) in answers.iter_mut().zip(&wants[u.clone()]) {
if let Ok(d) = a {
if !want {
d.bytes = Vec::new();
}
}
}
answers
});
(r, fuel)
};
let mut outs: Vec<UnitOut> = Vec::with_capacity(units.len());
if wants.iter().any(|&w| w) {
// Bytes come back: in order, bounded in flight, stop at the
// first failing unit (the same unit for every N).
let mut next = 0usize;
let producer = || {
let i = next;
next += 1;
units.get(i).map(|u| (i, u.clone()))
};
let mut sink = |_seq: u64, o: UnitOut| -> anyhow::Result<()> {
let failed = o.0.is_err();
outs.push(o);
if failed {
anyhow::bail!("a decoder unit failed");
}
Ok(())
};
let _ = gatling::gatling::ordered::run_ordered_sink(
producer,
workers,
0,
|_i: usize, u: Range<usize>| unit(&u),
&mut sink,
);
} else {
// Digests only: a plain fork-join over the units, in index order.
outs = gatling::gatling_forkjoin::gatling_for_each(units.len(), workers, |i| {
unit(&units[i])
});
}
let mut answers = Vec::with_capacity(ins.len());
let mut fuel = 0u64;
let mut counted = 0usize;
for (r, f) in outs {
fuel = fuel.saturating_add(f);
counted += 1;
match r {
Ok(a) => answers.extend(a),
Err(e) => {
return DecodeRun {
result: Err(e),
fuel_used: fuel,
units: counted,
}
}
}
}
DecodeRun {
result: Ok(answers),
fuel_used: fuel,
units: counted,
}
}
/// One fresh decoder instance on one unit.
fn run_unit(
&self,
chunks: &[ChunkIn<'_>],
limits: &RunLimits,
deadline: Instant,
first: usize,
) -> UnitOut {
let what = |why: String| {
format!(
"decoder instance for chunks {}..{}: {why}",
first,
first + chunks.len()
)
};
let fail = |code: &'static str, why: String, fuel: u64| -> UnitOut {
(
Err(DecodeFailure {
code,
reason: what(why),
}),
fuel,
)
};
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return fail(
codes::TIMEOUT,
format!("wall-clock timeout after {:?}", limits.timeout),
0,
);
}
let mut store = Store::new(
self.pre.engine(),
DState {
limiter: Limiter::new(limits),
},
);
store.limiter(|s| &mut s.limiter);
if let Err(e) = store.set_fuel(limits.fuel) {
return fail(codes::HOST_FAULT, format!("host fault (fuel: {e})"), 0);
}
let tick = self.epoch_tick.as_nanos().max(1);
let ticks = remaining.as_nanos().div_ceil(tick).max(1);
store.set_epoch_deadline(u64::try_from(ticks).unwrap_or(u64::MAX));
store.epoch_deadline_trap();
let used = |s: &Store<DState>| limits.fuel.saturating_sub(s.get_fuel().unwrap_or(0));
let instance = match self.pre.instantiate(&mut store) {
Ok(i) => i,
Err(e) => {
let (code, why) =
classify_error(&e, None, &store.data().limiter, limits, self.stack_bytes);
let f = used(&store);
return fail(code, why, f);
}
};
type Args<'a> = (&'a [ChunkIn<'a>],);
type Ret = (Vec<ChunkAnswer>,);
let func = match instance.get_typed_func::<Args<'_>, Ret>(&mut store, "decode-chunks") {
Ok(f) => f,
Err(e) => {
let f = used(&store);
return fail(
codes::INTERFACE_MISMATCH,
format!(
"decoder interface mismatch ({})",
one_line(&format!("{e:#}"))
),
f,
);
}
};
let answers = match func.call(&mut store, (chunks,)) {
Ok((a,)) => a,
Err(e) => {
let (code, why) =
classify_error(&e, None, &store.data().limiter, limits, self.stack_bytes);
let f = used(&store);
return fail(code, why, f);
}
};
let fuel = used(&store);
drop(store);
match check_answers(chunks, answers, first) {
Ok(a) => (Ok(a), fuel),
Err(why) => fail(codes::DECODER_MISMATCH, why, fuel),
}
}
}
/// Hold a decoder's answers to what it was asked, natively: one per chunk, a
/// 32-byte digest, exactly the declared size, BLAKE3(bytes) = digest.
fn check_answers(
chunks: &[ChunkIn<'_>],
answers: Vec<ChunkAnswer>,
first: usize,
) -> Result<Vec<ChunkAnswer>, String> {
if answers.len() != chunks.len() {
return Err(format!(
"{} answers for {} chunks",
answers.len(),
chunks.len()
));
}
let mut out = Vec::with_capacity(answers.len());
for (i, (c, a)) in chunks.iter().zip(answers).enumerate() {
let n = first + i;
match a {
Ok(d) => {
if d.digest.len() != 32 {
return Err(format!(
"chunk {n}: a {}-byte digest, not a 32-byte BLAKE3",
d.digest.len()
));
}
if d.bytes.len() as u64 != c.size {
return Err(format!(
"chunk {n}: {} bytes returned, {} declared",
d.bytes.len(),
c.size
));
}
if blake3::hash(&d.bytes).as_bytes()[..] != d.digest[..] {
return Err(format!(
"chunk {n}: the digest is not the BLAKE3 of the bytes returned"
));
}
out.push(Ok(d));
}
Err(why) => {
let mut why = one_line(&why);
if why.len() > MAX_DECODER_TEXT {
let mut cut = MAX_DECODER_TEXT;
while !why.is_char_boundary(cut) {
cut -= 1;
}
why.truncate(cut);
}
out.push(Err(why));
}
}
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn units_depend_on_the_chunk_sizes_only_and_cover_every_chunk_once() {
let sizes: Vec<(usize, u64)> = (0..300).map(|i| (1000 + i, (i as u64 % 7) << 20)).collect();
let u = plan_units(sizes.iter().copied());
assert_eq!(u, plan_units(sizes.iter().copied()));
let mut at = 0;
for r in &u {
assert_eq!(r.start, at);
assert!(r.end > r.start && r.end - r.start <= UNIT_CHUNKS);
at = r.end;
}
assert_eq!(at, sizes.len());
}
#[test]
fn many_small_chunks_close_a_unit_at_the_chunk_cap() {
let u = plan_units((0..130).map(|_| (10usize, 10u64)));
assert_eq!(u, vec![0..64, 64..128, 128..130]);
}
#[test]
fn one_huge_chunk_is_a_unit_of_its_own() {
let u = plan_units([(10usize, 1u64), (10, 64 << 20), (10, 1)]);
assert_eq!(u, vec![0..2, 2..3]);
}
#[test]
fn no_chunks_no_units() {
assert!(plan_units(std::iter::empty()).is_empty());
}
}