1use crate::dra::op::Op;
4use crate::error::{Error, Result};
5use crate::limits::Limits;
6
7pub const DRA_VERSION: u8 = 2;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum Authority {
13 Literal,
15 Generated,
17 EntropyChannel,
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct Span {
24 pub start: u64,
26 pub len: u64,
28 pub authority: Authority,
30}
31
32#[derive(Debug, Clone, PartialEq, Eq, Default)]
37pub struct CoverageMap {
38 pub spans: Vec<Span>,
40}
41
42impl CoverageMap {
43 pub fn total_len(&self) -> u64 {
45 self.spans.last().map(|s| s.start + s.len).unwrap_or(0)
46 }
47
48 pub fn validate(&self, declared_len: u64) -> Result<()> {
50 let mut expected = 0u64;
51 for s in &self.spans {
52 if s.start != expected {
53 return Err(Error::coverage_violation(format!(
54 "coverage gap/overlap: expected span at {expected}, found {}",
55 s.start
56 )));
57 }
58 expected = expected
59 .checked_add(s.len)
60 .ok_or_else(|| Error::coverage_violation("coverage length overflow"))?;
61 }
62 if expected != declared_len {
63 return Err(Error::coverage_violation(format!(
64 "coverage covers {expected} bytes but {declared_len} were declared"
65 )));
66 }
67 Ok(())
68 }
69}
70
71#[derive(Debug, Clone, PartialEq, Eq, Default)]
73pub struct Program {
74 pub ops: Vec<Op>,
76}
77
78impl Program {
79 pub fn new(ops: Vec<Op>) -> Self {
81 Program { ops }
82 }
83
84 pub fn encode(&self) -> Result<Vec<u8>> {
86 let count = u32::try_from(self.ops.len())
87 .map_err(|_| Error::resource_limit("too many DRA instructions"))?;
88 let mut out = Vec::with_capacity(5 + self.ops.len() * 5);
89 out.push(DRA_VERSION);
90 out.extend_from_slice(&count.to_le_bytes());
91 for op in &self.ops {
92 op.encode(&mut out)?;
93 }
94 Ok(out)
95 }
96
97 pub fn decode(data: &[u8], limits: Limits) -> Result<Program> {
99 if data.is_empty() {
100 return Err(Error::invalid_graph("empty graph record"));
101 }
102 if data[0] != DRA_VERSION {
103 return Err(Error::unsupported_version(format!(
104 "DRA version {} is not supported",
105 data[0]
106 )));
107 }
108 if data.len() < 5 {
109 return Err(Error::invalid_graph("truncated graph header"));
110 }
111 let count = u32::from_le_bytes([data[1], data[2], data[3], data[4]]);
112 if count > limits.max_graph_ops {
113 return Err(Error::resource_limit(format!(
114 "graph has {count} instructions, limit {}",
115 limits.max_graph_ops
116 )));
117 }
118 let mut pos = 5usize;
119 let mut ops = Vec::with_capacity(count as usize);
120 for _ in 0..count {
121 ops.push(Op::decode(data, &mut pos, limits)?);
122 }
123 if pos != data.len() {
124 return Err(Error::invalid_graph(format!(
125 "graph record has {} trailing bytes",
126 data.len() - pos
127 )));
128 }
129 Ok(Program { ops })
130 }
131
132 pub fn analyze(
137 &self,
138 object_lens: &[u64],
139 channel_lens: &[u64],
140 limits: Limits,
141 ) -> Result<(u64, CoverageMap)> {
142 if self.ops.len() as u64 > limits.max_graph_ops as u64 {
143 return Err(Error::resource_limit("graph instruction limit exceeded"));
144 }
145 let mut spans: Vec<Span> = Vec::new();
146 let mut total: u64 = 0;
147 let mut last_len: u64 = 0;
148 let mut have_last = false;
149
150 for op in &self.ops {
151 match op {
152 Op::EmitObject { object_id } => {
153 let len = *object_lens.get(*object_id as usize).ok_or_else(|| {
154 Error::invalid_graph(format!("graph references missing object {object_id}"))
155 })?;
156 total = total
157 .checked_add(len)
158 .ok_or_else(|| Error::resource_limit("output length overflow"))?;
159 if len > 0 {
160 spans.push(Span {
161 start: total - len,
162 len,
163 authority: Authority::Literal,
164 });
165 }
166 last_len = len;
167 have_last = true;
168 }
169 Op::Inline { bytes } => {
170 let len = bytes.len() as u64;
171 total = total
172 .checked_add(len)
173 .ok_or_else(|| Error::resource_limit("output length overflow"))?;
174 if len > 0 {
175 spans.push(Span {
176 start: total - len,
177 len,
178 authority: Authority::Literal,
179 });
180 }
181 last_len = len;
182 have_last = true;
183 }
184 Op::DecodeChannel { channel_id } => {
185 let len = *channel_lens.get(*channel_id as usize).ok_or_else(|| {
186 Error::invalid_graph(format!(
187 "graph references missing entropy channel {channel_id}"
188 ))
189 })?;
190 total = total
191 .checked_add(len)
192 .ok_or_else(|| Error::resource_limit("output length overflow"))?;
193 if len > 0 {
194 spans.push(Span {
195 start: total - len,
196 len,
197 authority: Authority::EntropyChannel,
198 });
199 }
200 last_len = len;
201 have_last = true;
202 }
203 Op::RepeatLast { count } => {
204 if !have_last {
205 return Err(Error::invalid_graph(
206 "REPEAT_LAST has no preceding literal instruction",
207 ));
208 }
209 if (*count as u64) > limits.max_repeat_count {
210 return Err(Error::resource_limit(format!(
211 "REPEAT_LAST count {count} exceeds limit {}",
212 limits.max_repeat_count
213 )));
214 }
215 let extra = last_len
216 .checked_mul(*count as u64)
217 .ok_or_else(|| Error::resource_limit("repeat length overflow"))?;
218 total = total
219 .checked_add(extra)
220 .ok_or_else(|| Error::resource_limit("output length overflow"))?;
221 if extra > 0 {
222 spans.push(Span {
223 start: total - extra,
224 len: extra,
225 authority: Authority::Generated,
226 });
227 }
228 have_last = false;
231 last_len = 0;
232 }
233 }
234 if total > limits.max_output_bytes {
235 return Err(Error::resource_limit(format!(
236 "predicted output {total} exceeds limit {}",
237 limits.max_output_bytes
238 )));
239 }
240 }
241
242 Ok((total, CoverageMap { spans }))
243 }
244
245 pub fn analyze_inputs(
247 &self,
248 objects: &[Vec<u8>],
249 channels: &[Vec<u8>],
250 limits: Limits,
251 ) -> Result<(u64, CoverageMap)> {
252 let object_lens: Vec<u64> = objects.iter().map(|o| o.len() as u64).collect();
253 let channel_lens: Vec<u64> = channels.iter().map(|c| c.len() as u64).collect();
254 self.analyze(&object_lens, &channel_lens, limits)
255 }
256
257 pub fn analyze_objects(
260 &self,
261 objects: &[Vec<u8>],
262 limits: Limits,
263 ) -> Result<(u64, CoverageMap)> {
264 let lens: Vec<u64> = objects.iter().map(|o| o.len() as u64).collect();
265 self.analyze(&lens, &[], limits)
266 }
267
268 pub fn eval(
270 &self,
271 objects: &[Vec<u8>],
272 channels: &[Vec<u8>],
273 limits: Limits,
274 ) -> Result<Vec<u8>> {
275 let (predicted, _coverage) = self.analyze_inputs(objects, channels, limits)?;
276 let cap = predicted.min(64 * 1024 * 1024) as usize;
277 let mut out: Vec<u8> = Vec::with_capacity(cap);
278 let mut have_last = false;
279 let mut block_len: usize = 0;
280
281 for op in &self.ops {
282 match op {
283 Op::EmitObject { object_id } => {
284 let obj = objects.get(*object_id as usize).ok_or_else(|| {
285 Error::invalid_graph(format!("graph references missing object {object_id}"))
286 })?;
287 block_len = obj.len();
288 out.extend_from_slice(obj);
289 have_last = true;
290 }
291 Op::Inline { bytes } => {
292 block_len = bytes.len();
293 out.extend_from_slice(bytes);
294 have_last = true;
295 }
296 Op::DecodeChannel { channel_id } => {
297 let ch = channels.get(*channel_id as usize).ok_or_else(|| {
298 Error::invalid_graph(format!(
299 "graph references missing entropy channel {channel_id}"
300 ))
301 })?;
302 block_len = ch.len();
303 out.extend_from_slice(ch);
304 have_last = true;
305 }
306 Op::RepeatLast { count } => {
307 if !have_last {
308 return Err(Error::invalid_graph(
309 "REPEAT_LAST has no preceding literal instruction",
310 ));
311 }
312 let start = out.len() - block_len;
317 for _ in 0..*count {
318 out.extend_from_within(start..start + block_len);
319 }
320 have_last = false;
321 block_len = 0;
322 }
323 }
324 if out.len() as u64 > limits.max_output_bytes {
325 return Err(Error::resource_limit(
326 "output exceeds materialization limit",
327 ));
328 }
329 }
330 Ok(out)
331 }
332}
333
334#[cfg(test)]
335mod tests {
336 use super::*;
337
338 fn objs(xs: &[&[u8]]) -> Vec<Vec<u8>> {
339 xs.iter().map(|x| x.to_vec()).collect()
340 }
341
342 #[test]
343 fn literal_concat_and_coverage() {
344 let objects = objs(&[b"hello ", b"world"]);
345 let p = Program::new(vec![
346 Op::EmitObject { object_id: 0 },
347 Op::EmitObject { object_id: 1 },
348 ]);
349 let (len, cov) = p.analyze_objects(&objects, Limits::DEFAULT).unwrap();
350 assert_eq!(len, 11);
351 cov.validate(11).unwrap();
352 assert_eq!(
353 p.eval(&objects, &[], Limits::DEFAULT).unwrap(),
354 b"hello world"
355 );
356 }
357
358 #[test]
359 fn inline_and_repeat() {
360 let objects = objs(&[]);
361 let p = Program::new(vec![
362 Op::Inline {
363 bytes: b"ab".to_vec(),
364 },
365 Op::RepeatLast { count: 2 },
366 ]);
367 let (len, cov) = p.analyze_objects(&objects, Limits::DEFAULT).unwrap();
368 assert_eq!(len, 6);
369 cov.validate(6).unwrap();
370 assert_eq!(p.eval(&objects, &[], Limits::DEFAULT).unwrap(), b"ababab");
371 assert_eq!(
373 cov.spans[0],
374 Span {
375 start: 0,
376 len: 2,
377 authority: Authority::Literal
378 }
379 );
380 assert_eq!(
381 cov.spans[1],
382 Span {
383 start: 2,
384 len: 4,
385 authority: Authority::Generated
386 }
387 );
388 }
389
390 #[test]
391 fn decode_channel_and_repeat() {
392 let objects = objs(&[]);
393 let channels = objs(&[b"abc"]);
394 let p = Program::new(vec![
395 Op::DecodeChannel { channel_id: 0 },
396 Op::RepeatLast { count: 1 },
397 ]);
398 let (len, cov) = p
399 .analyze_inputs(&objects, &channels, Limits::DEFAULT)
400 .unwrap();
401 assert_eq!(len, 6);
402 cov.validate(6).unwrap();
403 assert_eq!(
404 cov.spans[0],
405 Span {
406 start: 0,
407 len: 3,
408 authority: Authority::EntropyChannel,
409 }
410 );
411 assert_eq!(
412 cov.spans[1],
413 Span {
414 start: 3,
415 len: 3,
416 authority: Authority::Generated,
417 }
418 );
419 assert_eq!(
420 p.eval(&objects, &channels, Limits::DEFAULT).unwrap(),
421 b"abcabc"
422 );
423 }
424
425 #[test]
426 fn rejects_missing_channel() {
427 let p = Program::new(vec![Op::DecodeChannel { channel_id: 5 }]);
428 let e = p.analyze_inputs(&[], &[], Limits::DEFAULT).unwrap_err();
429 assert_eq!(e.class(), crate::ErrorClass::InvalidGraph);
430 }
431
432 #[test]
433 fn rejects_missing_object() {
434 let objects = objs(&[]);
435 let p = Program::new(vec![Op::EmitObject { object_id: 3 }]);
436 let e = p.analyze_objects(&objects, Limits::DEFAULT).unwrap_err();
437 assert_eq!(e.class(), crate::ErrorClass::InvalidGraph);
438 }
439
440 #[test]
441 fn rejects_leading_repeat() {
442 let objects = objs(&[]);
443 let p = Program::new(vec![Op::RepeatLast { count: 1 }]);
444 let e = p.analyze_objects(&objects, Limits::DEFAULT).unwrap_err();
445 assert_eq!(e.class(), crate::ErrorClass::InvalidGraph);
446 }
447
448 #[test]
449 fn rejects_consecutive_repeats() {
450 let objects = objs(&[]);
451 let p = Program::new(vec![
452 Op::Inline {
453 bytes: b"x".to_vec(),
454 },
455 Op::RepeatLast { count: 1 },
456 Op::RepeatLast { count: 1 },
457 ]);
458 let e = p.analyze_objects(&objects, Limits::DEFAULT).unwrap_err();
459 assert_eq!(e.class(), crate::ErrorClass::InvalidGraph);
460 }
461
462 #[test]
463 fn enforces_output_limit() {
464 let objects = objs(&[b"abcdefgh"]);
465 let p = Program::new(vec![
466 Op::EmitObject { object_id: 0 },
467 Op::RepeatLast { count: 1000 },
468 ]);
469 let limits = Limits {
470 max_output_bytes: 64,
471 ..Limits::DEFAULT
472 };
473 let e = p.analyze_objects(&objects, limits).unwrap_err();
474 assert_eq!(e.class(), crate::ErrorClass::ResourceLimit);
475 }
476
477 #[test]
478 fn coverage_gap_is_rejected() {
479 let cov = CoverageMap {
480 spans: vec![
481 Span {
482 start: 0,
483 len: 2,
484 authority: Authority::Literal,
485 },
486 Span {
487 start: 3,
488 len: 2,
489 authority: Authority::Literal,
490 },
491 ],
492 };
493 let e = cov.validate(4).unwrap_err();
494 assert_eq!(e.class(), crate::ErrorClass::CoverageViolation);
495 }
496
497 #[test]
498 fn program_roundtrips_via_bytes() {
499 let p = Program::new(vec![
500 Op::EmitObject { object_id: 1 },
501 Op::Inline {
502 bytes: b"hi".to_vec(),
503 },
504 Op::RepeatLast { count: 3 },
505 ]);
506 let enc = p.encode().unwrap();
507 let back = Program::decode(&enc, Limits::DEFAULT).unwrap();
508 assert_eq!(p, back);
509 }
510}