1use crate::dra::Program;
36use crate::error::{Error, Result};
37use crate::limits::Limits;
38
39pub fn primary_dependency(op: &crate::dra::Op) -> (u8, u32) {
45 use crate::dra::op::{DEFLATE_SOURCE_CHANNEL, Op};
46 match op {
47 Op::EmitObject { object_id } => (DEP_OBJECT, *object_id),
48 Op::DecodeChannel { channel_id } => (DEP_CHANNEL, *channel_id),
49 Op::InterleaveChannels { kinds_channel, .. } => (DEP_CHANNEL, *kinds_channel),
50 Op::PackSegments { data_object, .. } => (DEP_OBJECT, *data_object),
51 Op::PackedChannels { data_channel, .. } => (DEP_CHANNEL, *data_channel),
52 Op::DeflateReplay {
53 source_kind,
54 source_id,
55 ..
56 } => match *source_kind {
57 DEFLATE_SOURCE_CHANNEL => (DEP_CHANNEL, *source_id),
58 _ => (DEP_OBJECT, *source_id),
59 },
60 Op::Inline { .. }
61 | Op::MarkOffset { .. }
62 | Op::EmitOffset { .. }
63 | Op::RepeatLast { .. } => (DEP_NONE, 0),
64 }
65}
66
67pub const OBSERVATION_INDEX_VERSION: u8 = 1;
69
70pub const SECTION_OP_TABLE: u8 = 0x01;
72pub const SECTION_PDF_SELECTORS: u8 = 0x02;
74pub const SECTION_DIGESTS: u8 = 0x04;
76
77const KNOWN_SECTION_FLAGS: u8 = SECTION_OP_TABLE | SECTION_PDF_SELECTORS | SECTION_DIGESTS;
79
80pub const DEP_NONE: u8 = 0;
82pub const DEP_OBJECT: u8 = 1;
84pub const DEP_CHANNEL: u8 = 2;
86
87pub const SELECTOR_OBJECT: u8 = 0;
89pub const SELECTOR_STREAM: u8 = 1;
91pub const SELECTOR_REVISION: u8 = 2;
93
94const OP_ENTRY_LEN: usize = 4 + 1 + 4;
95const SELECTOR_ENTRY_LEN: usize = 1 + 4 + 4 + 8 + 8;
96const DIGEST_ENTRY_LEN: usize = 8 + 8 + 32;
97
98#[derive(Debug, Clone, Copy, PartialEq, Eq)]
101pub struct OpEntry {
102 pub out_len: u32,
104 pub dep_kind: u8,
106 pub dep_id: u32,
108}
109
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
112pub struct ObservationSelector {
113 pub kind: u8,
115 pub number: u32,
117 pub generation: u32,
119 pub out_off: u64,
121 pub out_len: u64,
123}
124
125#[derive(Debug, Clone, Copy, PartialEq, Eq)]
127pub struct ObservationDigest {
128 pub out_off: u64,
130 pub out_len: u64,
132 pub sha256: [u8; 32],
134}
135
136#[derive(Debug, Clone, PartialEq, Eq, Default)]
141pub struct ObservationIndex {
142 pub section_flags: u8,
144 pub ops: Vec<OpEntry>,
146 pub selectors: Vec<ObservationSelector>,
148 pub digests: Vec<ObservationDigest>,
150}
151
152impl ObservationIndex {
153 pub fn encode(&self) -> Result<Vec<u8>> {
159 if self.section_flags & !KNOWN_SECTION_FLAGS != 0 {
160 return Err(Error::invalid_container(
161 "observation index has unknown section flags",
162 ));
163 }
164 if self.section_flags & SECTION_OP_TABLE == 0 && !self.ops.is_empty() {
165 return Err(Error::invalid_container(
166 "observation index carries ops without the op-table section flag",
167 ));
168 }
169 if self.section_flags & SECTION_PDF_SELECTORS == 0 && !self.selectors.is_empty() {
170 return Err(Error::invalid_container(
171 "observation index carries selectors without the selector section flag",
172 ));
173 }
174 if self.section_flags & SECTION_DIGESTS == 0 && !self.digests.is_empty() {
175 return Err(Error::invalid_container(
176 "observation index carries digests without the digest section flag",
177 ));
178 }
179
180 let mut out = Vec::new();
181 out.push(OBSERVATION_INDEX_VERSION);
182 out.push(self.section_flags);
183
184 if self.section_flags & SECTION_OP_TABLE != 0 {
185 let count = u32::try_from(self.ops.len())
186 .map_err(|_| Error::resource_limit("observation index op_count exceeds u32"))?;
187 out.extend_from_slice(&count.to_le_bytes());
188 for e in &self.ops {
189 out.extend_from_slice(&e.out_len.to_le_bytes());
190 out.push(e.dep_kind);
191 out.extend_from_slice(&e.dep_id.to_le_bytes());
192 }
193 }
194 if self.section_flags & SECTION_PDF_SELECTORS != 0 {
195 let count = u32::try_from(self.selectors.len()).map_err(|_| {
196 Error::resource_limit("observation index selector_count exceeds u32")
197 })?;
198 out.extend_from_slice(&count.to_le_bytes());
199 for s in &self.selectors {
200 out.push(s.kind);
201 out.extend_from_slice(&s.number.to_le_bytes());
202 out.extend_from_slice(&s.generation.to_le_bytes());
203 out.extend_from_slice(&s.out_off.to_le_bytes());
204 out.extend_from_slice(&s.out_len.to_le_bytes());
205 }
206 }
207 if self.section_flags & SECTION_DIGESTS != 0 {
208 let count = u32::try_from(self.digests.len())
209 .map_err(|_| Error::resource_limit("observation index digest_count exceeds u32"))?;
210 out.extend_from_slice(&count.to_le_bytes());
211 for d in &self.digests {
212 out.extend_from_slice(&d.out_off.to_le_bytes());
213 out.extend_from_slice(&d.out_len.to_le_bytes());
214 out.extend_from_slice(&d.sha256);
215 }
216 }
217 Ok(out)
218 }
219
220 pub fn decode(bytes: &[u8], limits: Limits) -> Result<ObservationIndex> {
226 if bytes.len() < 2 {
227 return Err(Error::invalid_container(
228 "truncated OBSERVATION_INDEX header",
229 ));
230 }
231 let version = bytes[0];
232 if version != OBSERVATION_INDEX_VERSION {
233 return Err(Error::unsupported_version(format!(
234 "observation index version {version} is not supported"
235 )));
236 }
237 let section_flags = bytes[1];
238 if section_flags & !KNOWN_SECTION_FLAGS != 0 {
239 return Err(Error::invalid_container(
240 "OBSERVATION_INDEX has unknown section flags",
241 ));
242 }
243
244 let mut p = 2usize;
245 let mut ops = Vec::new();
246 if section_flags & SECTION_OP_TABLE != 0 {
247 let count = read_u32(bytes, &mut p)?;
248 if count > limits.max_graph_ops {
249 return Err(Error::resource_limit(format!(
250 "observation index op_count {count} exceeds limit {}",
251 limits.max_graph_ops
252 )));
253 }
254 let count = count as usize;
255 require(bytes, p, count, OP_ENTRY_LEN)?;
256 ops.reserve(count);
257 for _ in 0..count {
258 ops.push(OpEntry {
259 out_len: read_u32(bytes, &mut p)?,
260 dep_kind: read_u8(bytes, &mut p)?,
261 dep_id: read_u32(bytes, &mut p)?,
262 });
263 }
264 }
265
266 let mut selectors = Vec::new();
267 if section_flags & SECTION_PDF_SELECTORS != 0 {
268 let count = read_u32(bytes, &mut p)?;
269 if count > limits.max_index_selectors {
270 return Err(Error::resource_limit(format!(
271 "observation index selector_count {count} exceeds limit {}",
272 limits.max_index_selectors
273 )));
274 }
275 let count = count as usize;
276 require(bytes, p, count, SELECTOR_ENTRY_LEN)?;
277 selectors.reserve(count);
278 for _ in 0..count {
279 selectors.push(ObservationSelector {
280 kind: read_u8(bytes, &mut p)?,
281 number: read_u32(bytes, &mut p)?,
282 generation: read_u32(bytes, &mut p)?,
283 out_off: read_u64(bytes, &mut p)?,
284 out_len: read_u64(bytes, &mut p)?,
285 });
286 }
287 }
288
289 let mut digests = Vec::new();
290 if section_flags & SECTION_DIGESTS != 0 {
291 let count = read_u32(bytes, &mut p)?;
292 let count = count as usize;
293 require(bytes, p, count, DIGEST_ENTRY_LEN)?;
294 digests.reserve(count);
295 for _ in 0..count {
296 let out_off = read_u64(bytes, &mut p)?;
297 let out_len = read_u64(bytes, &mut p)?;
298 let mut sha256 = [0u8; 32];
299 sha256.copy_from_slice(bytes.get(p..p + 32).ok_or_else(|| {
300 Error::invalid_container("truncated OBSERVATION_INDEX digest")
301 })?);
302 p += 32;
303 digests.push(ObservationDigest {
304 out_off,
305 out_len,
306 sha256,
307 });
308 }
309 }
310
311 if p != bytes.len() {
312 return Err(Error::invalid_container(format!(
313 "OBSERVATION_INDEX has {} trailing bytes",
314 bytes.len() - p
315 )));
316 }
317
318 Ok(ObservationIndex {
319 section_flags,
320 ops,
321 selectors,
322 digests,
323 })
324 }
325
326 pub fn validate(
332 &self,
333 program: &Program,
334 object_lens: &[u64],
335 channel_lens: &[u64],
336 limits: Limits,
337 ) -> Result<()> {
338 let per_op = program.analyze_ops(object_lens, channel_lens, limits)?;
339 let mut total: u64 = 0;
340 for len in &per_op {
341 total = total
342 .checked_add(*len)
343 .ok_or_else(|| Error::coverage_violation("analysis output length overflow"))?;
344 }
345
346 if self.section_flags & SECTION_OP_TABLE != 0 {
347 if self.ops.len() != program.ops.len() {
348 return Err(Error::coverage_violation(format!(
349 "observation index op_count {} disagrees with program op count {}",
350 self.ops.len(),
351 program.ops.len()
352 )));
353 }
354 for (i, entry) in self.ops.iter().enumerate() {
355 if u64::from(entry.out_len) != per_op[i] {
356 return Err(Error::coverage_violation(format!(
357 "observation index op {i} out_len {} disagrees with analyzed {}",
358 entry.out_len, per_op[i]
359 )));
360 }
361 match entry.dep_kind {
362 DEP_NONE => {}
363 DEP_OBJECT => {
364 if entry.dep_id as usize >= object_lens.len() {
365 return Err(Error::coverage_violation(format!(
366 "observation index op {i} references missing object {}",
367 entry.dep_id
368 )));
369 }
370 }
371 DEP_CHANNEL => {
372 if entry.dep_id as usize >= channel_lens.len() {
373 return Err(Error::coverage_violation(format!(
374 "observation index op {i} references missing channel {}",
375 entry.dep_id
376 )));
377 }
378 }
379 other => {
380 return Err(Error::coverage_violation(format!(
381 "observation index op {i} has invalid dep_kind {other}"
382 )));
383 }
384 }
385 }
386 }
387
388 if self.section_flags & SECTION_PDF_SELECTORS != 0 {
389 for (i, sel) in self.selectors.iter().enumerate() {
390 if sel.kind > SELECTOR_REVISION {
391 return Err(Error::coverage_violation(format!(
392 "observation index selector {i} has invalid kind {}",
393 sel.kind
394 )));
395 }
396 if sel.out_len == 0 {
397 return Err(Error::coverage_violation(format!(
398 "observation index selector {i} has an empty range"
399 )));
400 }
401 let end = sel.out_off.checked_add(sel.out_len).ok_or_else(|| {
402 Error::coverage_violation(format!(
403 "observation index selector {i} range overflows"
404 ))
405 })?;
406 if end > total {
407 return Err(Error::coverage_violation(format!(
408 "observation index selector {i} range {}..{} exceeds total {total}",
409 sel.out_off, end
410 )));
411 }
412 }
413 }
414
415 if self.section_flags & SECTION_DIGESTS != 0 {
416 for (i, digest) in self.digests.iter().enumerate() {
417 if digest.out_len == 0 {
418 return Err(Error::coverage_violation(format!(
419 "observation index digest {i} has an empty range"
420 )));
421 }
422 let end = digest.out_off.checked_add(digest.out_len).ok_or_else(|| {
423 Error::coverage_violation(format!(
424 "observation index digest {i} range overflows"
425 ))
426 })?;
427 if end > total {
428 return Err(Error::coverage_violation(format!(
429 "observation index digest {i} range {}..{} exceeds total {total}",
430 digest.out_off, end
431 )));
432 }
433 }
434 }
435
436 Ok(())
437 }
438}
439
440fn require(bytes: &[u8], p: usize, count: usize, entry_len: usize) -> Result<()> {
442 let need = count
443 .checked_mul(entry_len)
444 .ok_or_else(|| Error::invalid_container("OBSERVATION_INDEX section length overflow"))?;
445 let available = bytes
446 .len()
447 .checked_sub(p)
448 .ok_or_else(|| Error::invalid_container("OBSERVATION_INDEX cursor past end of payload"))?;
449 if need > available {
450 return Err(Error::invalid_container(format!(
451 "OBSERVATION_INDEX section needs {need} bytes but only {available} remain"
452 )));
453 }
454 Ok(())
455}
456
457fn read_u8(bytes: &[u8], p: &mut usize) -> Result<u8> {
458 let v = *bytes
459 .get(*p)
460 .ok_or_else(|| Error::invalid_container("truncated OBSERVATION_INDEX payload"))?;
461 *p += 1;
462 Ok(v)
463}
464
465fn read_u32(bytes: &[u8], p: &mut usize) -> Result<u32> {
466 let end = p
467 .checked_add(4)
468 .ok_or_else(|| Error::invalid_container("OBSERVATION_INDEX cursor overflow"))?;
469 let slice = bytes
470 .get(*p..end)
471 .ok_or_else(|| Error::invalid_container("truncated OBSERVATION_INDEX payload"))?;
472 *p = end;
473 Ok(u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]))
474}
475
476fn read_u64(bytes: &[u8], p: &mut usize) -> Result<u64> {
477 let end = p
478 .checked_add(8)
479 .ok_or_else(|| Error::invalid_container("OBSERVATION_INDEX cursor overflow"))?;
480 let slice = bytes
481 .get(*p..end)
482 .ok_or_else(|| Error::invalid_container("truncated OBSERVATION_INDEX payload"))?;
483 *p = end;
484 Ok(u64::from_le_bytes([
485 slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
486 ]))
487}
488
489#[cfg(test)]
490mod tests {
491 use super::*;
492 use crate::dra::Op;
493
494 fn simple_index() -> ObservationIndex {
495 ObservationIndex {
496 section_flags: SECTION_OP_TABLE,
497 ops: vec![
498 OpEntry {
499 out_len: 3,
500 dep_kind: DEP_OBJECT,
501 dep_id: 0,
502 },
503 OpEntry {
504 out_len: 0,
505 dep_kind: DEP_NONE,
506 dep_id: 0,
507 },
508 ],
509 selectors: vec![],
510 digests: vec![],
511 }
512 }
513
514 #[test]
515 fn op_table_roundtrips() {
516 let idx = simple_index();
517 let bytes = idx.encode().unwrap();
518 let back = ObservationIndex::decode(&bytes, Limits::DEFAULT).unwrap();
519 assert_eq!(back, idx);
520 }
521
522 #[test]
523 fn all_sections_roundtrip() {
524 let idx = ObservationIndex {
525 section_flags: SECTION_OP_TABLE | SECTION_PDF_SELECTORS | SECTION_DIGESTS,
526 ops: vec![OpEntry {
527 out_len: 3,
528 dep_kind: DEP_OBJECT,
529 dep_id: 0,
530 }],
531 selectors: vec![ObservationSelector {
532 kind: SELECTOR_STREAM,
533 number: 4,
534 generation: 0,
535 out_off: 0,
536 out_len: 3,
537 }],
538 digests: vec![ObservationDigest {
539 out_off: 0,
540 out_len: 3,
541 sha256: [9u8; 32],
542 }],
543 };
544 let bytes = idx.encode().unwrap();
545 let back = ObservationIndex::decode(&bytes, Limits::DEFAULT).unwrap();
546 assert_eq!(back, idx);
547 }
548
549 #[test]
550 fn unknown_version_and_flags_fail_closed() {
551 let mut bytes = simple_index().encode().unwrap();
552 bytes[0] = 2;
553 assert_eq!(
554 ObservationIndex::decode(&bytes, Limits::DEFAULT)
555 .unwrap_err()
556 .class(),
557 crate::ErrorClass::UnsupportedVersion
558 );
559
560 let mut bytes = simple_index().encode().unwrap();
561 bytes[1] = 0x80;
562 assert_eq!(
563 ObservationIndex::decode(&bytes, Limits::DEFAULT)
564 .unwrap_err()
565 .class(),
566 crate::ErrorClass::InvalidContainer
567 );
568 }
569
570 #[test]
571 fn truncated_section_is_rejected() {
572 let mut bytes = simple_index().encode().unwrap();
573 bytes.truncate(bytes.len() - 1);
574 assert_eq!(
575 ObservationIndex::decode(&bytes, Limits::DEFAULT)
576 .unwrap_err()
577 .class(),
578 crate::ErrorClass::InvalidContainer
579 );
580 }
581
582 #[test]
583 fn trailing_bytes_are_rejected() {
584 let mut bytes = simple_index().encode().unwrap();
585 bytes.push(0);
586 assert_eq!(
587 ObservationIndex::decode(&bytes, Limits::DEFAULT)
588 .unwrap_err()
589 .class(),
590 crate::ErrorClass::InvalidContainer
591 );
592 }
593
594 #[test]
595 fn validate_accepts_a_correct_index() {
596 let program = Program::new(vec![
598 Op::EmitObject { object_id: 0 },
599 Op::Inline {
600 bytes: b"ab".to_vec(),
601 },
602 ]);
603 let object_lens = [3u64];
604 let idx = ObservationIndex {
605 section_flags: SECTION_OP_TABLE | SECTION_PDF_SELECTORS | SECTION_DIGESTS,
606 ops: vec![
607 OpEntry {
608 out_len: 3,
609 dep_kind: DEP_OBJECT,
610 dep_id: 0,
611 },
612 OpEntry {
613 out_len: 2,
614 dep_kind: DEP_NONE,
615 dep_id: 0,
616 },
617 ],
618 selectors: vec![ObservationSelector {
619 kind: SELECTOR_OBJECT,
620 number: 1,
621 generation: 0,
622 out_off: 0,
623 out_len: 3,
624 }],
625 digests: vec![ObservationDigest {
626 out_off: 3,
627 out_len: 2,
628 sha256: [0u8; 32],
629 }],
630 };
631 idx.validate(&program, &object_lens, &[], Limits::DEFAULT)
632 .unwrap();
633 }
634
635 #[test]
636 fn validate_rejects_wrong_op_len_dep_and_selector() {
637 let program = Program::new(vec![Op::EmitObject { object_id: 0 }]);
638 let object_lens = [3u64];
639
640 let wrong_len = ObservationIndex {
641 section_flags: SECTION_OP_TABLE,
642 ops: vec![OpEntry {
643 out_len: 4,
644 dep_kind: DEP_OBJECT,
645 dep_id: 0,
646 }],
647 ..Default::default()
648 };
649 assert_eq!(
650 wrong_len
651 .validate(&program, &object_lens, &[], Limits::DEFAULT)
652 .unwrap_err()
653 .class(),
654 crate::ErrorClass::CoverageViolation
655 );
656
657 let bad_dep = ObservationIndex {
658 section_flags: SECTION_OP_TABLE,
659 ops: vec![OpEntry {
660 out_len: 3,
661 dep_kind: DEP_OBJECT,
662 dep_id: 7,
663 }],
664 ..Default::default()
665 };
666 assert_eq!(
667 bad_dep
668 .validate(&program, &object_lens, &[], Limits::DEFAULT)
669 .unwrap_err()
670 .class(),
671 crate::ErrorClass::CoverageViolation
672 );
673
674 let past_total = ObservationIndex {
675 section_flags: SECTION_PDF_SELECTORS,
676 selectors: vec![ObservationSelector {
677 kind: SELECTOR_OBJECT,
678 number: 1,
679 generation: 0,
680 out_off: 2,
681 out_len: 5,
682 }],
683 ..Default::default()
684 };
685 assert_eq!(
686 past_total
687 .validate(&program, &object_lens, &[], Limits::DEFAULT)
688 .unwrap_err()
689 .class(),
690 crate::ErrorClass::CoverageViolation
691 );
692 }
693}