use crate::dra::Program;
use crate::error::{Error, Result};
use crate::limits::Limits;
pub const OBSERVATION_INDEX_VERSION: u8 = 1;
pub const SECTION_OP_TABLE: u8 = 0x01;
pub const SECTION_PDF_SELECTORS: u8 = 0x02;
pub const SECTION_DIGESTS: u8 = 0x04;
const KNOWN_SECTION_FLAGS: u8 = SECTION_OP_TABLE | SECTION_PDF_SELECTORS | SECTION_DIGESTS;
pub const DEP_NONE: u8 = 0;
pub const DEP_OBJECT: u8 = 1;
pub const DEP_CHANNEL: u8 = 2;
pub const SELECTOR_OBJECT: u8 = 0;
pub const SELECTOR_STREAM: u8 = 1;
pub const SELECTOR_REVISION: u8 = 2;
const OP_ENTRY_LEN: usize = 4 + 1 + 4;
const SELECTOR_ENTRY_LEN: usize = 1 + 4 + 4 + 8 + 8;
const DIGEST_ENTRY_LEN: usize = 8 + 8 + 32;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OpEntry {
pub out_len: u32,
pub dep_kind: u8,
pub dep_id: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ObservationSelector {
pub kind: u8,
pub number: u32,
pub generation: u32,
pub out_off: u64,
pub out_len: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ObservationDigest {
pub out_off: u64,
pub out_len: u64,
pub sha256: [u8; 32],
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct ObservationIndex {
pub section_flags: u8,
pub ops: Vec<OpEntry>,
pub selectors: Vec<ObservationSelector>,
pub digests: Vec<ObservationDigest>,
}
impl ObservationIndex {
pub fn encode(&self) -> Result<Vec<u8>> {
if self.section_flags & !KNOWN_SECTION_FLAGS != 0 {
return Err(Error::invalid_container(
"observation index has unknown section flags",
));
}
if self.section_flags & SECTION_OP_TABLE == 0 && !self.ops.is_empty() {
return Err(Error::invalid_container(
"observation index carries ops without the op-table section flag",
));
}
if self.section_flags & SECTION_PDF_SELECTORS == 0 && !self.selectors.is_empty() {
return Err(Error::invalid_container(
"observation index carries selectors without the selector section flag",
));
}
if self.section_flags & SECTION_DIGESTS == 0 && !self.digests.is_empty() {
return Err(Error::invalid_container(
"observation index carries digests without the digest section flag",
));
}
let mut out = Vec::new();
out.push(OBSERVATION_INDEX_VERSION);
out.push(self.section_flags);
if self.section_flags & SECTION_OP_TABLE != 0 {
let count = u32::try_from(self.ops.len())
.map_err(|_| Error::resource_limit("observation index op_count exceeds u32"))?;
out.extend_from_slice(&count.to_le_bytes());
for e in &self.ops {
out.extend_from_slice(&e.out_len.to_le_bytes());
out.push(e.dep_kind);
out.extend_from_slice(&e.dep_id.to_le_bytes());
}
}
if self.section_flags & SECTION_PDF_SELECTORS != 0 {
let count = u32::try_from(self.selectors.len()).map_err(|_| {
Error::resource_limit("observation index selector_count exceeds u32")
})?;
out.extend_from_slice(&count.to_le_bytes());
for s in &self.selectors {
out.push(s.kind);
out.extend_from_slice(&s.number.to_le_bytes());
out.extend_from_slice(&s.generation.to_le_bytes());
out.extend_from_slice(&s.out_off.to_le_bytes());
out.extend_from_slice(&s.out_len.to_le_bytes());
}
}
if self.section_flags & SECTION_DIGESTS != 0 {
let count = u32::try_from(self.digests.len())
.map_err(|_| Error::resource_limit("observation index digest_count exceeds u32"))?;
out.extend_from_slice(&count.to_le_bytes());
for d in &self.digests {
out.extend_from_slice(&d.out_off.to_le_bytes());
out.extend_from_slice(&d.out_len.to_le_bytes());
out.extend_from_slice(&d.sha256);
}
}
Ok(out)
}
pub fn decode(bytes: &[u8], limits: Limits) -> Result<ObservationIndex> {
if bytes.len() < 2 {
return Err(Error::invalid_container(
"truncated OBSERVATION_INDEX header",
));
}
let version = bytes[0];
if version != OBSERVATION_INDEX_VERSION {
return Err(Error::unsupported_version(format!(
"observation index version {version} is not supported"
)));
}
let section_flags = bytes[1];
if section_flags & !KNOWN_SECTION_FLAGS != 0 {
return Err(Error::invalid_container(
"OBSERVATION_INDEX has unknown section flags",
));
}
let mut p = 2usize;
let mut ops = Vec::new();
if section_flags & SECTION_OP_TABLE != 0 {
let count = read_u32(bytes, &mut p)?;
if count > limits.max_graph_ops {
return Err(Error::resource_limit(format!(
"observation index op_count {count} exceeds limit {}",
limits.max_graph_ops
)));
}
let count = count as usize;
require(bytes, p, count, OP_ENTRY_LEN)?;
ops.reserve(count);
for _ in 0..count {
ops.push(OpEntry {
out_len: read_u32(bytes, &mut p)?,
dep_kind: read_u8(bytes, &mut p)?,
dep_id: read_u32(bytes, &mut p)?,
});
}
}
let mut selectors = Vec::new();
if section_flags & SECTION_PDF_SELECTORS != 0 {
let count = read_u32(bytes, &mut p)?;
if count > limits.max_index_selectors {
return Err(Error::resource_limit(format!(
"observation index selector_count {count} exceeds limit {}",
limits.max_index_selectors
)));
}
let count = count as usize;
require(bytes, p, count, SELECTOR_ENTRY_LEN)?;
selectors.reserve(count);
for _ in 0..count {
selectors.push(ObservationSelector {
kind: read_u8(bytes, &mut p)?,
number: read_u32(bytes, &mut p)?,
generation: read_u32(bytes, &mut p)?,
out_off: read_u64(bytes, &mut p)?,
out_len: read_u64(bytes, &mut p)?,
});
}
}
let mut digests = Vec::new();
if section_flags & SECTION_DIGESTS != 0 {
let count = read_u32(bytes, &mut p)?;
let count = count as usize;
require(bytes, p, count, DIGEST_ENTRY_LEN)?;
digests.reserve(count);
for _ in 0..count {
let out_off = read_u64(bytes, &mut p)?;
let out_len = read_u64(bytes, &mut p)?;
let mut sha256 = [0u8; 32];
sha256.copy_from_slice(bytes.get(p..p + 32).ok_or_else(|| {
Error::invalid_container("truncated OBSERVATION_INDEX digest")
})?);
p += 32;
digests.push(ObservationDigest {
out_off,
out_len,
sha256,
});
}
}
if p != bytes.len() {
return Err(Error::invalid_container(format!(
"OBSERVATION_INDEX has {} trailing bytes",
bytes.len() - p
)));
}
Ok(ObservationIndex {
section_flags,
ops,
selectors,
digests,
})
}
pub fn validate(
&self,
program: &Program,
object_lens: &[u64],
channel_lens: &[u64],
limits: Limits,
) -> Result<()> {
let per_op = program.analyze_ops(object_lens, channel_lens, limits)?;
let mut total: u64 = 0;
for len in &per_op {
total = total
.checked_add(*len)
.ok_or_else(|| Error::coverage_violation("analysis output length overflow"))?;
}
if self.section_flags & SECTION_OP_TABLE != 0 {
if self.ops.len() != program.ops.len() {
return Err(Error::coverage_violation(format!(
"observation index op_count {} disagrees with program op count {}",
self.ops.len(),
program.ops.len()
)));
}
for (i, entry) in self.ops.iter().enumerate() {
if u64::from(entry.out_len) != per_op[i] {
return Err(Error::coverage_violation(format!(
"observation index op {i} out_len {} disagrees with analyzed {}",
entry.out_len, per_op[i]
)));
}
match entry.dep_kind {
DEP_NONE => {}
DEP_OBJECT => {
if entry.dep_id as usize >= object_lens.len() {
return Err(Error::coverage_violation(format!(
"observation index op {i} references missing object {}",
entry.dep_id
)));
}
}
DEP_CHANNEL => {
if entry.dep_id as usize >= channel_lens.len() {
return Err(Error::coverage_violation(format!(
"observation index op {i} references missing channel {}",
entry.dep_id
)));
}
}
other => {
return Err(Error::coverage_violation(format!(
"observation index op {i} has invalid dep_kind {other}"
)));
}
}
}
}
if self.section_flags & SECTION_PDF_SELECTORS != 0 {
for (i, sel) in self.selectors.iter().enumerate() {
if sel.kind > SELECTOR_REVISION {
return Err(Error::coverage_violation(format!(
"observation index selector {i} has invalid kind {}",
sel.kind
)));
}
if sel.out_len == 0 {
return Err(Error::coverage_violation(format!(
"observation index selector {i} has an empty range"
)));
}
let end = sel.out_off.checked_add(sel.out_len).ok_or_else(|| {
Error::coverage_violation(format!(
"observation index selector {i} range overflows"
))
})?;
if end > total {
return Err(Error::coverage_violation(format!(
"observation index selector {i} range {}..{} exceeds total {total}",
sel.out_off, end
)));
}
}
}
if self.section_flags & SECTION_DIGESTS != 0 {
for (i, digest) in self.digests.iter().enumerate() {
if digest.out_len == 0 {
return Err(Error::coverage_violation(format!(
"observation index digest {i} has an empty range"
)));
}
let end = digest.out_off.checked_add(digest.out_len).ok_or_else(|| {
Error::coverage_violation(format!(
"observation index digest {i} range overflows"
))
})?;
if end > total {
return Err(Error::coverage_violation(format!(
"observation index digest {i} range {}..{} exceeds total {total}",
digest.out_off, end
)));
}
}
}
Ok(())
}
}
fn require(bytes: &[u8], p: usize, count: usize, entry_len: usize) -> Result<()> {
let need = count
.checked_mul(entry_len)
.ok_or_else(|| Error::invalid_container("OBSERVATION_INDEX section length overflow"))?;
let available = bytes
.len()
.checked_sub(p)
.ok_or_else(|| Error::invalid_container("OBSERVATION_INDEX cursor past end of payload"))?;
if need > available {
return Err(Error::invalid_container(format!(
"OBSERVATION_INDEX section needs {need} bytes but only {available} remain"
)));
}
Ok(())
}
fn read_u8(bytes: &[u8], p: &mut usize) -> Result<u8> {
let v = *bytes
.get(*p)
.ok_or_else(|| Error::invalid_container("truncated OBSERVATION_INDEX payload"))?;
*p += 1;
Ok(v)
}
fn read_u32(bytes: &[u8], p: &mut usize) -> Result<u32> {
let end = p
.checked_add(4)
.ok_or_else(|| Error::invalid_container("OBSERVATION_INDEX cursor overflow"))?;
let slice = bytes
.get(*p..end)
.ok_or_else(|| Error::invalid_container("truncated OBSERVATION_INDEX payload"))?;
*p = end;
Ok(u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]))
}
fn read_u64(bytes: &[u8], p: &mut usize) -> Result<u64> {
let end = p
.checked_add(8)
.ok_or_else(|| Error::invalid_container("OBSERVATION_INDEX cursor overflow"))?;
let slice = bytes
.get(*p..end)
.ok_or_else(|| Error::invalid_container("truncated OBSERVATION_INDEX payload"))?;
*p = end;
Ok(u64::from_le_bytes([
slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
]))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dra::Op;
fn simple_index() -> ObservationIndex {
ObservationIndex {
section_flags: SECTION_OP_TABLE,
ops: vec![
OpEntry {
out_len: 3,
dep_kind: DEP_OBJECT,
dep_id: 0,
},
OpEntry {
out_len: 0,
dep_kind: DEP_NONE,
dep_id: 0,
},
],
selectors: vec![],
digests: vec![],
}
}
#[test]
fn op_table_roundtrips() {
let idx = simple_index();
let bytes = idx.encode().unwrap();
let back = ObservationIndex::decode(&bytes, Limits::DEFAULT).unwrap();
assert_eq!(back, idx);
}
#[test]
fn all_sections_roundtrip() {
let idx = ObservationIndex {
section_flags: SECTION_OP_TABLE | SECTION_PDF_SELECTORS | SECTION_DIGESTS,
ops: vec![OpEntry {
out_len: 3,
dep_kind: DEP_OBJECT,
dep_id: 0,
}],
selectors: vec![ObservationSelector {
kind: SELECTOR_STREAM,
number: 4,
generation: 0,
out_off: 0,
out_len: 3,
}],
digests: vec![ObservationDigest {
out_off: 0,
out_len: 3,
sha256: [9u8; 32],
}],
};
let bytes = idx.encode().unwrap();
let back = ObservationIndex::decode(&bytes, Limits::DEFAULT).unwrap();
assert_eq!(back, idx);
}
#[test]
fn unknown_version_and_flags_fail_closed() {
let mut bytes = simple_index().encode().unwrap();
bytes[0] = 2;
assert_eq!(
ObservationIndex::decode(&bytes, Limits::DEFAULT)
.unwrap_err()
.class(),
crate::ErrorClass::UnsupportedVersion
);
let mut bytes = simple_index().encode().unwrap();
bytes[1] = 0x80;
assert_eq!(
ObservationIndex::decode(&bytes, Limits::DEFAULT)
.unwrap_err()
.class(),
crate::ErrorClass::InvalidContainer
);
}
#[test]
fn truncated_section_is_rejected() {
let mut bytes = simple_index().encode().unwrap();
bytes.truncate(bytes.len() - 1);
assert_eq!(
ObservationIndex::decode(&bytes, Limits::DEFAULT)
.unwrap_err()
.class(),
crate::ErrorClass::InvalidContainer
);
}
#[test]
fn trailing_bytes_are_rejected() {
let mut bytes = simple_index().encode().unwrap();
bytes.push(0);
assert_eq!(
ObservationIndex::decode(&bytes, Limits::DEFAULT)
.unwrap_err()
.class(),
crate::ErrorClass::InvalidContainer
);
}
#[test]
fn validate_accepts_a_correct_index() {
let program = Program::new(vec![
Op::EmitObject { object_id: 0 },
Op::Inline {
bytes: b"ab".to_vec(),
},
]);
let object_lens = [3u64];
let idx = ObservationIndex {
section_flags: SECTION_OP_TABLE | SECTION_PDF_SELECTORS | SECTION_DIGESTS,
ops: vec![
OpEntry {
out_len: 3,
dep_kind: DEP_OBJECT,
dep_id: 0,
},
OpEntry {
out_len: 2,
dep_kind: DEP_NONE,
dep_id: 0,
},
],
selectors: vec![ObservationSelector {
kind: SELECTOR_OBJECT,
number: 1,
generation: 0,
out_off: 0,
out_len: 3,
}],
digests: vec![ObservationDigest {
out_off: 3,
out_len: 2,
sha256: [0u8; 32],
}],
};
idx.validate(&program, &object_lens, &[], Limits::DEFAULT)
.unwrap();
}
#[test]
fn validate_rejects_wrong_op_len_dep_and_selector() {
let program = Program::new(vec![Op::EmitObject { object_id: 0 }]);
let object_lens = [3u64];
let wrong_len = ObservationIndex {
section_flags: SECTION_OP_TABLE,
ops: vec![OpEntry {
out_len: 4,
dep_kind: DEP_OBJECT,
dep_id: 0,
}],
..Default::default()
};
assert_eq!(
wrong_len
.validate(&program, &object_lens, &[], Limits::DEFAULT)
.unwrap_err()
.class(),
crate::ErrorClass::CoverageViolation
);
let bad_dep = ObservationIndex {
section_flags: SECTION_OP_TABLE,
ops: vec![OpEntry {
out_len: 3,
dep_kind: DEP_OBJECT,
dep_id: 7,
}],
..Default::default()
};
assert_eq!(
bad_dep
.validate(&program, &object_lens, &[], Limits::DEFAULT)
.unwrap_err()
.class(),
crate::ErrorClass::CoverageViolation
);
let past_total = ObservationIndex {
section_flags: SECTION_PDF_SELECTORS,
selectors: vec![ObservationSelector {
kind: SELECTOR_OBJECT,
number: 1,
generation: 0,
out_off: 2,
out_len: 5,
}],
..Default::default()
};
assert_eq!(
past_total
.validate(&program, &object_lens, &[], Limits::DEFAULT)
.unwrap_err()
.class(),
crate::ErrorClass::CoverageViolation
);
}
}