1#![allow(missing_docs)]
34
35use std::collections::VecDeque;
36use std::fmt;
37
38use crate::canonical::remap_node;
39use crate::encode::{encode_with, hash_image};
40use crate::error::{Error, Result};
41use crate::ir::{Annotation, AnnotationKind, InfinityModule, NodeId, Provenance};
42use crate::layout::{
43 align8, read_u32, read_u64, EncodeOpts, Section, HEADER_SIZE, SECTION_COUNT, STREAM_MAGIC,
44};
45use crate::validate::Validated;
46
47pub const STREAM_VERSION: u8 = 1;
49pub const FRAME_HEADER_SIZE: usize = 32;
51pub const DEFAULT_CHUNK: usize = 64 << 10;
53pub const MAX_FRAME_PAYLOAD: u32 = 64 << 20;
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum FrameKind {
59 Begin,
61 Batch,
63 Delta,
65 End,
67}
68
69impl FrameKind {
70 pub const fn code(self) -> u8 {
71 match self {
72 FrameKind::Begin => 1,
73 FrameKind::Batch => 2,
74 FrameKind::Delta => 3,
75 FrameKind::End => 4,
76 }
77 }
78
79 pub const fn from_code(code: u8) -> Option<Self> {
80 match code {
81 1 => Some(FrameKind::Begin),
82 2 => Some(FrameKind::Batch),
83 3 => Some(FrameKind::Delta),
84 4 => Some(FrameKind::End),
85 _ => None,
86 }
87 }
88}
89
90#[derive(Debug, Clone, PartialEq, Eq)]
92pub struct Frame {
93 pub kind: FrameKind,
94 pub sequence: u64,
95 pub payload: Vec<u8>,
96}
97
98impl Frame {
99 pub fn new(kind: FrameKind, sequence: u64, payload: Vec<u8>) -> Self {
100 Frame {
101 kind,
102 sequence,
103 payload,
104 }
105 }
106
107 pub fn to_bytes(&self) -> Vec<u8> {
109 let mut out = Vec::with_capacity(FRAME_HEADER_SIZE + self.payload.len());
110 out.extend_from_slice(&STREAM_MAGIC);
111 out.push(STREAM_VERSION);
112 out.push(self.kind.code());
113 out.extend_from_slice(&0u16.to_le_bytes());
114 out.extend_from_slice(&self.sequence.to_le_bytes());
115 out.extend_from_slice(&(self.payload.len() as u32).to_le_bytes());
116 out.extend_from_slice(&0u32.to_le_bytes());
117 let checksum = frame_checksum(&self.payload);
118 out.extend_from_slice(&checksum[..8]);
119 out.extend_from_slice(&self.payload);
120 out
121 }
122
123 pub fn parse(buf: &[u8]) -> Result<(Frame, usize)> {
125 if buf.len() < 4 || buf[0..4] != STREAM_MAGIC {
126 return Err(Error::StreamMagic);
127 }
128 if buf.len() < FRAME_HEADER_SIZE {
129 return Err(Error::Truncated {
130 needed: FRAME_HEADER_SIZE as u64,
131 found: buf.len() as u64,
132 });
133 }
134 let version = buf[4];
135 if version != STREAM_VERSION {
136 return Err(Error::StreamVersion(version));
137 }
138 let kind = FrameKind::from_code(buf[5]).ok_or(Error::StreamFrameKind(buf[5]))?;
139 let flags = u16::from_le_bytes([buf[6], buf[7]]);
140 if flags != 0 {
141 return Err(Error::UnknownFlags(flags));
142 }
143 let sequence = u64::from_le_bytes(buf[8..16].try_into().expect("8 bytes"));
144 let len = u32::from_le_bytes(buf[16..20].try_into().expect("4 bytes"));
145 if len > MAX_FRAME_PAYLOAD {
146 return Err(Error::StreamFrameLen {
147 declared: len,
148 available: buf.len() as u64,
149 });
150 }
151 let reserved = u32::from_le_bytes(buf[20..24].try_into().expect("4 bytes"));
152 if reserved != 0 {
153 return Err(Error::ReservedNonZero {
154 section: "stream-frame",
155 offset: 20,
156 });
157 }
158 let payload_start = FRAME_HEADER_SIZE;
159 let payload_end = payload_start + len as usize;
160 if payload_end > buf.len() {
161 return Err(Error::StreamFrameLen {
162 declared: len,
163 available: (buf.len() - payload_start) as u64,
164 });
165 }
166 let payload = &buf[payload_start..payload_end];
167 let expected = frame_checksum(payload);
168 if buf[24..32] != expected[..8] {
169 return Err(Error::HashMismatch);
170 }
171 Ok((
172 Frame {
173 kind,
174 sequence,
175 payload: payload.to_vec(),
176 },
177 payload_end,
178 ))
179 }
180}
181
182pub fn frame_checksum(payload: &[u8]) -> [u8; 32] {
184 *blake3::hash(payload).as_bytes()
185}
186
187pub fn encode_stream(module: &InfinityModule, chunk: usize) -> Result<Vec<Frame>> {
193 encode_stream_with(module, chunk, EncodeOpts::default())
194}
195
196pub fn encode_stream_with(
198 module: &InfinityModule,
199 chunk: usize,
200 opts: EncodeOpts,
201) -> Result<Vec<Frame>> {
202 if chunk == 0 {
203 return Err(Error::InvalidArgument("chunk size must be non-zero"));
204 }
205 let image = encode_with(module, opts)?;
206 let header = crate::layout::Header::parse(&image)?;
207 let mut frames = Vec::new();
208 let mut sequence = 0u64;
209 let mut begin = Vec::with_capacity(24);
210 begin.extend_from_slice(&header.major.to_le_bytes());
211 begin.extend_from_slice(&header.minor.to_le_bytes());
212 begin.extend_from_slice(&header.flags.to_le_bytes());
213 begin.extend_from_slice(&0u16.to_le_bytes());
214 begin.extend_from_slice(&header.dialect_hash.to_le_bytes());
215 begin.extend_from_slice(&header.total_size.to_le_bytes());
216 frames.push(Frame::new(FrameKind::Begin, sequence, begin));
217 sequence += 1;
218
219 for section in Section::ALL {
222 let span = header.span(section);
223 let start = span.offset as usize;
224 let len = section_body_len(&image, &header, section)?;
225 let body = &image[start..start + len];
226 if body.is_empty() {
227 let mut payload = Vec::with_capacity(9);
228 payload.push(section as u8);
229 payload.extend_from_slice(&span.count.to_le_bytes());
230 payload.extend_from_slice(&0u32.to_le_bytes());
231 frames.push(Frame::new(FrameKind::Batch, sequence, payload));
232 sequence += 1;
233 continue;
234 }
235 for piece in body.chunks(chunk) {
236 let mut payload = Vec::with_capacity(9 + piece.len());
237 payload.push(section as u8);
238 payload.extend_from_slice(&span.count.to_le_bytes());
239 payload.extend_from_slice(&(piece.len() as u32).to_le_bytes());
240 payload.extend_from_slice(piece);
241 frames.push(Frame::new(FrameKind::Batch, sequence, payload));
242 sequence += 1;
243 }
244 }
245
246 let mut end = Vec::with_capacity(40);
247 let hash = hash_image(&image)?;
248 end.extend_from_slice(&hash);
249 end.extend_from_slice(&(image.len() as u64).to_le_bytes());
250 frames.push(Frame::new(FrameKind::End, sequence, end));
251 Ok(frames)
252}
253
254fn section_body_len(
255 image: &[u8],
256 header: &crate::layout::Header,
257 section: Section,
258) -> crate::Result<usize> {
259 let span = header.span(section);
260 let start = span.offset as u64;
261 let count = span.count as u64;
262 let read_at = |offset: u64| -> crate::Result<usize> {
263 usize::try_from(offset).map_err(|_| Error::TooLarge)
264 };
265 let end = match section {
266 Section::Strings => {
267 if count == 0 {
268 start
269 } else {
270 let index = start
271 .checked_add(count.checked_mul(4).ok_or(Error::TooLarge)?)
272 .ok_or(Error::TooLarge)?;
273 let heap = read_u32(image, read_at(index)?).unwrap_or(0) as u64;
274 start
275 .checked_add(
276 count
277 .checked_add(1)
278 .ok_or(Error::TooLarge)?
279 .checked_mul(4)
280 .ok_or(Error::TooLarge)?,
281 )
282 .and_then(|value| value.checked_add(heap))
283 .ok_or(Error::TooLarge)?
284 }
285 }
286 Section::Blobs => {
287 let mut heap = 0u64;
288 for index in 0..count {
289 let offset = start
290 .checked_add(index.checked_mul(40).ok_or(Error::TooLarge)?)
291 .and_then(|value| value.checked_add(36))
292 .ok_or(Error::TooLarge)?;
293 heap = heap
294 .checked_add(u64::from(read_u32(image, read_at(offset)?).unwrap_or(0)))
295 .ok_or(Error::TooLarge)?;
296 }
297 start
298 .checked_add(count.checked_mul(40).ok_or(Error::TooLarge)?)
299 .and_then(|value| value.checked_add(heap))
300 .ok_or(Error::TooLarge)?
301 }
302 Section::Provenance => start.checked_add(count).ok_or(Error::TooLarge)?,
303 Section::Secrets => {
304 let mut heap = 0u64;
305 for index in 0..count {
306 let offset = start
307 .checked_add(index.checked_mul(8).ok_or(Error::TooLarge)?)
308 .and_then(|value| value.checked_add(4))
309 .ok_or(Error::TooLarge)?;
310 heap = heap
311 .checked_add(u64::from(read_u32(image, read_at(offset)?).unwrap_or(0)))
312 .ok_or(Error::TooLarge)?;
313 }
314 start
315 .checked_add(count.checked_mul(8).ok_or(Error::TooLarge)?)
316 .and_then(|value| value.checked_add(heap))
317 .ok_or(Error::TooLarge)?
318 }
319 Section::Extensions => {
320 let directory = count.checked_mul(16).ok_or(Error::TooLarge)?;
321 let mut end = start.checked_add(directory).ok_or(Error::TooLarge)?;
322 for index in 0..count {
323 let dir = start
324 .checked_add(index.checked_mul(16).ok_or(Error::TooLarge)?)
325 .ok_or(Error::TooLarge)?;
326 let offset = u64::from(
327 read_u32(image, read_at(dir.checked_add(8).ok_or(Error::TooLarge)?)?)
328 .unwrap_or(0),
329 );
330 let len = u64::from(
331 read_u32(image, read_at(dir.checked_add(12).ok_or(Error::TooLarge)?)?)
332 .unwrap_or(0),
333 );
334 end = end.max(offset.checked_add(len).ok_or(Error::TooLarge)?);
335 }
336 end
337 }
338 Section::Nodes => start
339 .checked_add(count.checked_mul(32).ok_or(Error::TooLarge)?)
340 .ok_or(Error::TooLarge)?,
341 Section::Children => start
342 .checked_add(count.checked_mul(4).ok_or(Error::TooLarge)?)
343 .ok_or(Error::TooLarge)?,
344 Section::Fields => start
345 .checked_add(count.checked_mul(16).ok_or(Error::TooLarge)?)
346 .ok_or(Error::TooLarge)?,
347 Section::Parts => start
348 .checked_add(count.checked_mul(8).ok_or(Error::TooLarge)?)
349 .ok_or(Error::TooLarge)?,
350 Section::Edges => start
351 .checked_add(count.checked_mul(8).ok_or(Error::TooLarge)?)
352 .ok_or(Error::TooLarge)?,
353 Section::Annotations => start
354 .checked_add(count.checked_mul(16).ok_or(Error::TooLarge)?)
355 .ok_or(Error::TooLarge)?,
356 Section::Units => start
357 .checked_add(count.checked_mul(8).ok_or(Error::TooLarge)?)
358 .ok_or(Error::TooLarge)?,
359 Section::Types => start
360 .checked_add(count.checked_mul(16).ok_or(Error::TooLarge)?)
361 .ok_or(Error::TooLarge)?,
362 Section::Aux => start
363 .checked_add(count.checked_mul(8).ok_or(Error::TooLarge)?)
364 .ok_or(Error::TooLarge)?,
365 Section::Numbers => start
366 .checked_add(count.checked_mul(8).ok_or(Error::TooLarge)?)
367 .ok_or(Error::TooLarge)?,
368 Section::Hashes => start
369 .checked_add(count.checked_mul(32).ok_or(Error::TooLarge)?)
370 .ok_or(Error::TooLarge)?,
371 Section::Bindings => start
372 .checked_add(count.checked_mul(8).ok_or(Error::TooLarge)?)
373 .ok_or(Error::TooLarge)?,
374 };
375 usize::try_from(end.saturating_sub(start)).map_err(|_| Error::TooLarge)
376}
377
378#[derive(Clone, PartialEq)]
380pub enum DeltaOp {
381 SetRoot(Option<NodeId>),
383 PutNode {
385 target: NodeId,
386 fragment: Box<InfinityModule>,
387 },
388 SetBinding { name: String, node: NodeId },
390 RemoveBinding { name: String },
392 SetAnnotation { node: NodeId, kind: AnnotationKind },
394 ClearAnnotations,
396 SetProvenance(Option<Provenance>),
398}
399
400impl fmt::Debug for DeltaOp {
401 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
402 match self {
403 DeltaOp::SetRoot(root) => write!(f, "SetRoot({root:?})"),
404 DeltaOp::PutNode { target, fragment } => write!(
405 f,
406 "PutNode {{ target: {target}, fragment: <{} nodes> }}",
407 fragment.node_count()
408 ),
409 DeltaOp::SetBinding { node, .. } => {
410 write!(f, "SetBinding {{ name: <redacted>, node: {node} }}")
411 }
412 DeltaOp::RemoveBinding { .. } => f.write_str("RemoveBinding { name: <redacted> }"),
413 DeltaOp::SetAnnotation { node, kind } => {
414 write!(f, "SetAnnotation {{ node: {node}, kind: {kind:?} }}")
415 }
416 DeltaOp::ClearAnnotations => f.write_str("ClearAnnotations"),
417 DeltaOp::SetProvenance(provenance) => write!(f, "SetProvenance({provenance:?})"),
418 }
419 }
420}
421
422impl DeltaOp {
423 fn encode(&self, out: &mut Vec<u8>) -> Result<()> {
424 match self {
425 DeltaOp::SetRoot(root) => {
426 out.push(0);
427 out.extend_from_slice(&root.map(|id| id.0).unwrap_or(u32::MAX).to_le_bytes());
428 }
429 DeltaOp::PutNode { target, fragment } => {
430 out.push(1);
431 out.extend_from_slice(&target.0.to_le_bytes());
432 let image = encode_with(fragment, EncodeOpts::default())?;
433 out.extend_from_slice(&(image.len() as u32).to_le_bytes());
434 out.extend_from_slice(&image);
435 }
436 DeltaOp::SetBinding { name, node } => {
437 out.push(3);
438 let bytes = name.as_bytes();
439 out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
440 out.extend_from_slice(bytes);
441 out.extend_from_slice(&node.0.to_le_bytes());
442 }
443 DeltaOp::RemoveBinding { name } => {
444 out.push(4);
445 let bytes = name.as_bytes();
446 out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
447 out.extend_from_slice(bytes);
448 }
449 DeltaOp::SetAnnotation { node, kind } => {
450 out.push(5);
451 out.extend_from_slice(&node.0.to_le_bytes());
452 match kind {
453 AnnotationKind::Intent(text) => {
454 out.push(0);
455 write_string(out, text)?;
456 }
457 AnnotationKind::Constraint { source, expr } => {
458 out.push(1);
459 write_string(out, source)?;
460 out.extend_from_slice(&expr.0.to_le_bytes());
461 }
462 AnnotationKind::AiImmutable(flag) => {
463 out.push(2);
464 out.push(u8::from(*flag));
465 }
466 }
467 }
468 DeltaOp::ClearAnnotations => out.push(6),
469 DeltaOp::SetProvenance(prov) => {
470 out.push(7);
471 match prov {
472 None => out.push(0),
473 Some(prov) => {
474 out.push(1);
475 write_opt_string(out, prov.producer.as_deref())?;
476 write_opt_string(out, prov.model.as_deref())?;
477 write_opt_string(out, prov.tool.as_deref())?;
478 out.extend_from_slice(&prov.timestamp_ms.to_le_bytes());
479 match prov.parent {
480 None => out.push(0),
481 Some(hash) => {
482 out.push(1);
483 out.extend_from_slice(&hash);
484 }
485 }
486 out.extend_from_slice(&(prov.entries.len() as u32).to_le_bytes());
487 for entry in &prov.entries {
488 write_string(out, &entry.key)?;
489 match &entry.value {
490 crate::ir::ProvValue::Text(text) => {
491 out.push(0);
492 write_string(out, text)?;
493 }
494 crate::ir::ProvValue::Int(value) => {
495 out.push(1);
496 out.extend_from_slice(&value.to_le_bytes());
497 }
498 crate::ir::ProvValue::Flag(flag) => {
499 out.push(2);
500 out.push(u8::from(*flag));
501 }
502 crate::ir::ProvValue::Hash(hash) => {
503 out.push(3);
504 out.extend_from_slice(hash);
505 }
506 }
507 }
508 }
509 }
510 }
511 }
512 Ok(())
513 }
514
515 fn decode(buf: &[u8], cursor: &mut usize, limits: crate::layout::Limits) -> Result<DeltaOp> {
516 let tag = take_u8(buf, cursor)?;
517 let op = match tag {
518 0 => {
519 let id = take_u32(buf, cursor)?;
520 DeltaOp::SetRoot(if id == u32::MAX {
521 None
522 } else {
523 Some(NodeId(id))
524 })
525 }
526 1 => {
527 let target = NodeId(take_u32(buf, cursor)?);
528 let fragment = Box::new(decode_fragment(buf, cursor, limits)?);
529 DeltaOp::PutNode { target, fragment }
530 }
531 3 => {
532 let name = take_string(buf, cursor)?;
533 let node = NodeId(take_u32(buf, cursor)?);
534 DeltaOp::SetBinding { name, node }
535 }
536 4 => {
537 let name = take_string(buf, cursor)?;
538 DeltaOp::RemoveBinding { name }
539 }
540 5 => {
541 let node = NodeId(take_u32(buf, cursor)?);
542 let kind_tag = take_u8(buf, cursor)?;
543 let kind = match kind_tag {
544 0 => AnnotationKind::Intent(take_string(buf, cursor)?),
545 1 => {
546 let source = take_string(buf, cursor)?;
547 let expr = NodeId(take_u32(buf, cursor)?);
548 AnnotationKind::Constraint { source, expr }
549 }
550 2 => AnnotationKind::AiImmutable(take_u8(buf, cursor)? != 0),
551 _ => return Err(Error::StreamDelta("unknown annotation kind")),
552 };
553 DeltaOp::SetAnnotation { node, kind }
554 }
555 6 => DeltaOp::ClearAnnotations,
556 7 => {
557 if take_u8(buf, cursor)? == 0 {
558 DeltaOp::SetProvenance(None)
559 } else {
560 let producer = take_opt_string(buf, cursor)?;
561 let model = take_opt_string(buf, cursor)?;
562 let tool = take_opt_string(buf, cursor)?;
563 let timestamp_ms = take_u64(buf, cursor)?;
564 let parent = if take_u8(buf, cursor)? == 0 {
565 None
566 } else {
567 Some(take_hash(buf, cursor)?)
568 };
569 let count = take_u32(buf, cursor)?;
570 let mut entries = Vec::with_capacity(count.min(4096) as usize);
571 for _ in 0..count {
572 let key = take_string(buf, cursor)?;
573 let value = match take_u8(buf, cursor)? {
574 0 => crate::ir::ProvValue::Text(take_string(buf, cursor)?),
575 1 => crate::ir::ProvValue::Int(take_i64(buf, cursor)?),
576 2 => crate::ir::ProvValue::Flag(take_u8(buf, cursor)? != 0),
577 3 => crate::ir::ProvValue::Hash(take_hash(buf, cursor)?),
578 _ => return Err(Error::StreamDelta("unknown provenance value kind")),
579 };
580 entries.push(crate::ir::ProvEntry { key, value });
581 }
582 DeltaOp::SetProvenance(Some(Provenance {
583 producer,
584 timestamp_ms,
585 model,
586 tool,
587 parent,
588 entries,
589 }))
590 }
591 }
592 _ => return Err(Error::StreamDelta("unknown delta operation")),
593 };
594 Ok(op)
595 }
596}
597
598pub fn apply_delta(module: &InfinityModule, ops: &[DeltaOp]) -> Result<InfinityModule> {
600 let mut out = module.clone();
601 for op in ops {
602 match op {
603 DeltaOp::SetRoot(root) => {
604 if let Some(root) = root {
605 if out.get(*root).is_none() {
606 return Err(Error::StreamDelta("root node does not exist"));
607 }
608 }
609 out.set_root(*root);
610 }
611 DeltaOp::PutNode { target, fragment } => {
612 if out.get(*target).is_none() {
613 return Err(Error::StreamDelta("target node does not exist"));
614 }
615 if fragment.node_count() == 0 {
616 return Err(Error::StreamDelta("fragment is empty"));
617 }
618 let replacement = merge_fragment(&mut out, fragment)?;
619 let content = out
620 .get(replacement)
621 .cloned()
622 .ok_or(Error::StreamDelta("fragment root missing"))?;
623 *out.node_mut(*target)? = content;
624 }
625 DeltaOp::SetBinding { name, node } => {
626 if out.get(*node).is_none() {
627 return Err(Error::StreamDelta("binding target does not exist"));
628 }
629 out.bind(name.clone(), *node);
630 }
631 DeltaOp::RemoveBinding { name } => {
632 out.remove_binding(name);
633 }
634 DeltaOp::SetAnnotation { node, kind } => {
635 if out.get(*node).is_none() {
636 return Err(Error::StreamDelta("annotation subject does not exist"));
637 }
638 out.annotate(*node, kind.clone());
639 }
640 DeltaOp::ClearAnnotations => out.clear_annotations(),
641 DeltaOp::SetProvenance(prov) => out.set_provenance(prov.clone()),
642 }
643 }
644 Ok(out)
645}
646
647fn merge_fragment(base: &mut InfinityModule, fragment: &InfinityModule) -> Result<NodeId> {
648 let base_count = base.node_count();
649 let map: Vec<NodeId> = (0..fragment.node_count())
650 .map(|index| NodeId(base_count as u32 + index as u32))
651 .collect();
652 for node in fragment.nodes() {
653 base.push(remap_node(node, &map)?);
654 }
655 Ok(map[fragment.root().unwrap_or(NodeId(0)).index()])
656}
657
658#[derive(Debug, Clone, PartialEq)]
660pub struct StreamedModule {
661 pub module: InfinityModule,
662 pub image: Vec<u8>,
663 pub content_hash: [u8; 32],
664}
665
666#[derive(Debug, Default)]
668pub struct StreamDecoder {
669 limits: crate::layout::Limits,
670 sections: [Vec<u8>; SECTION_COUNT],
671 counts: [u32; SECTION_COUNT],
672 seen: [bool; SECTION_COUNT],
673 begun: bool,
674 ended: bool,
675 expected: u64,
676 flags: u16,
677 total_size: Option<u64>,
678 module: Option<InfinityModule>,
679 image: Option<Vec<u8>>,
680 pending_deltas: Vec<DeltaOp>,
681 frames: u64,
682 partial: Vec<u8>,
683 cursor: usize,
684}
685
686impl StreamDecoder {
687 pub fn new() -> Self {
688 Self {
689 limits: crate::layout::Limits::default(),
690 sections: Default::default(),
691 counts: [0; SECTION_COUNT],
692 seen: [false; SECTION_COUNT],
693 begun: false,
694 ended: false,
695 expected: 0,
696 flags: 0,
697 total_size: None,
698 module: None,
699 image: None,
700 pending_deltas: Vec::new(),
701 frames: 0,
702 partial: Vec::new(),
703 cursor: 0,
704 }
705 }
706
707 pub fn with_limits(mut self, limits: crate::layout::Limits) -> Self {
710 self.limits = limits;
711 self
712 }
713
714 pub fn push_bytes(&mut self, buf: &[u8]) -> Result<()> {
717 self.partial.extend_from_slice(buf);
718 loop {
719 if self.partial.len() < FRAME_HEADER_SIZE {
720 return Ok(());
721 }
722 match Frame::parse(&self.partial) {
723 Ok((frame, consumed)) => {
724 self.push(&frame)?;
725 self.partial.drain(..consumed);
726 }
727 Err(Error::Truncated { .. }) => return Ok(()),
731 Err(Error::StreamFrameLen { declared, .. }) if declared <= MAX_FRAME_PAYLOAD => {
732 return Ok(())
733 }
734 Err(error) => return Err(error),
735 }
736 }
737 }
738
739 pub fn push(&mut self, frame: &Frame) -> Result<()> {
741 if frame.sequence != self.expected {
742 return Err(Error::StreamSequence {
743 expected: self.expected,
744 found: frame.sequence,
745 });
746 }
747 self.expected += 1;
748 self.frames += 1;
749 match frame.kind {
750 FrameKind::Begin => {
751 if self.begun {
752 return Err(Error::StreamState("duplicate begin frame"));
753 }
754 if self.ended {
755 return Err(Error::StreamState("begin after end"));
756 }
757 if frame.payload.len() != 24 {
758 return Err(Error::StreamFrameLen {
759 declared: 24,
760 available: frame.payload.len() as u64,
761 });
762 }
763 let major = u16::from_le_bytes([frame.payload[0], frame.payload[1]]);
764 let minor = u16::from_le_bytes([frame.payload[2], frame.payload[3]]);
765 if major != crate::layout::FORMAT_MAJOR || minor > crate::layout::FORMAT_MINOR {
766 return Err(Error::VersionMajor(major));
767 }
768 let flags = u16::from_le_bytes([frame.payload[4], frame.payload[5]]);
769 if flags & !crate::layout::FLAG_KNOWN_MASK != 0 {
770 return Err(Error::UnknownFlags(flags & !crate::layout::FLAG_KNOWN_MASK));
771 }
772 if frame.payload[6..8] != [0, 0] {
773 return Err(Error::ReservedNonZero {
774 section: "stream-begin",
775 offset: 6,
776 });
777 }
778 let dialect = read_u64(&frame.payload, 8)?;
779 if dialect != crate::layout::dialect_hash() {
780 return Err(Error::DialectMismatch(dialect));
781 }
782 let size = read_u64(&frame.payload, 16)?;
783 if size > self.limits.max_total_bytes {
786 return Err(Error::Limit {
787 kind: "total bytes",
788 limit: self.limits.max_total_bytes,
789 });
790 }
791 self.total_size = Some(size);
792 self.flags = flags;
793 self.begun = true;
794 }
795 FrameKind::Batch => {
796 if !self.begun || self.ended {
797 return Err(Error::StreamState("batch outside a generation"));
798 }
799 if frame.payload.len() < 9 {
800 return Err(Error::StreamFrameLen {
801 declared: 9,
802 available: frame.payload.len() as u64,
803 });
804 }
805 let section = frame.payload[0] as usize;
806 if section >= SECTION_COUNT {
807 return Err(Error::StreamState("unknown section id"));
808 }
809 let count = read_u32(&frame.payload, 1)?;
810 let len = read_u32(&frame.payload, 5)? as usize;
811 if frame.payload.len() != 9 + len {
812 return Err(Error::StreamFrameLen {
813 declared: len as u32,
814 available: (frame.payload.len() - 9) as u64,
815 });
816 }
817 if section != self.cursor && section != self.cursor + 1 {
818 return Err(Error::StreamState("sections out of order"));
819 }
820 if section == self.cursor + 1 {
821 self.cursor = section;
822 }
823 if self.seen[section] && self.counts[section] != count {
824 return Err(Error::StreamState("section count changed mid-stream"));
825 }
826 self.seen[section] = true;
827 self.counts[section] = count;
828 self.sections[section].extend_from_slice(&frame.payload[9..]);
829 let accumulated: usize = self.sections.iter().map(|section| section.len()).sum();
833 let declared = self
834 .total_size
835 .map(|size| size.min(usize::MAX as u64) as usize)
836 .unwrap_or(usize::MAX);
837 if accumulated > declared {
838 return Err(Error::Limit {
839 kind: "stream bytes",
840 limit: declared as u64,
841 });
842 }
843 if accumulated as u64 > self.limits.max_total_bytes {
844 return Err(Error::Limit {
845 kind: "stream bytes",
846 limit: self.limits.max_total_bytes,
847 });
848 }
849 }
850 FrameKind::Delta => {
851 if !self.ended {
852 return Err(Error::StreamState("delta before end"));
853 }
854 let mut cursor = 0usize;
855 let count = take_u32(&frame.payload, &mut cursor)?;
856 if self.pending_deltas.len().saturating_add(count as usize)
859 > self.limits.max_aux as usize
860 {
861 return Err(Error::Limit {
862 kind: "delta operations",
863 limit: self.limits.max_aux as u64,
864 });
865 }
866 for _ in 0..count {
867 self.pending_deltas.push(DeltaOp::decode(
868 &frame.payload,
869 &mut cursor,
870 self.limits,
871 )?);
872 }
873 if cursor != frame.payload.len() {
874 return Err(Error::StreamDelta("trailing bytes in delta payload"));
875 }
876 }
877 FrameKind::End => {
878 if !self.begun {
879 return Err(Error::StreamState("end before begin"));
880 }
881 if frame.payload.len() != 40 {
882 return Err(Error::StreamFrameLen {
883 declared: 40,
884 available: frame.payload.len() as u64,
885 });
886 }
887 let mut hash = [0u8; 32];
888 hash.copy_from_slice(&frame.payload[..32]);
889 let size = read_u64(&frame.payload, 32)?;
890 if self.ended {
891 let base = self
893 .module
894 .as_ref()
895 .ok_or(Error::StreamState("no module to patch"))?;
896 let updated = apply_delta(base, &self.pending_deltas)?;
897 self.pending_deltas.clear();
898 let image = encode_with(&updated, EncodeOpts::default())?;
899 if image.len() as u64 != size {
900 return Err(Error::StreamDelta("delta changed the declared size"));
901 }
902 if hash_image(&image)? != hash {
903 return Err(Error::StreamHashMismatch);
904 }
905 let validated = Validated::new(
906 &image,
907 crate::layout::DecodeOpts {
908 limits: self.limits,
909 verify_hash: true,
910 verify_canonical: false,
911 },
912 )?;
913 let module = crate::decode::decode_validated(&validated)?;
914 self.module = Some(module);
915 self.image = Some(image);
916 return Ok(());
917 }
918 let image = self.assemble(size)?;
919 if hash_image(&image)? != hash {
920 return Err(Error::StreamHashMismatch);
921 }
922 let validated = Validated::new(
923 &image,
924 crate::layout::DecodeOpts {
925 limits: self.limits,
926 verify_hash: true,
927 verify_canonical: false,
928 },
929 )?;
930 let module = crate::decode::decode_validated(&validated)?;
931 self.module = Some(module);
932 self.image = Some(image);
933 self.ended = true;
934 }
935 }
936 Ok(())
937 }
938
939 fn assemble(&self, size: u64) -> Result<Vec<u8>> {
940 if size < HEADER_SIZE as u64 {
941 return Err(Error::StreamDelta("declared size smaller than a header"));
942 }
943 if self.seen.iter().any(|seen| !seen) {
944 return Err(Error::StreamState("generation missing a section"));
945 }
946 let mut image = vec![0u8; HEADER_SIZE];
947 let mut header = crate::layout::Header {
948 major: crate::layout::FORMAT_MAJOR,
949 minor: crate::layout::FORMAT_MINOR,
950 flags: self.flags,
951 header_size: HEADER_SIZE as u16,
952 content_hash: [0u8; 32],
953 dialect_hash: crate::layout::dialect_hash(),
954 total_size: 0,
955 sections: [crate::layout::SectionSpan::default(); SECTION_COUNT],
956 };
957 for section in Section::ALL {
958 while image.len() % 8 != 0 {
959 image.push(0);
960 }
961 header.span_mut(section).offset = image.len() as u32;
962 header.span_mut(section).count = self.counts[section as usize];
963 image.extend_from_slice(&self.sections[section as usize]);
964 }
965 while image.len() % 8 != 0 {
966 image.push(0);
967 }
968 header.total_size = image.len() as u64;
969 header.write(&mut image);
970 let hash = hash_image(&image)?;
971 image[16..48].copy_from_slice(&hash);
972 if image.len() as u64 != size {
973 return Err(Error::StreamDelta(
974 "reassembled size disagrees with the end frame",
975 ));
976 }
977 Ok(image)
978 }
979
980 pub fn module(&self) -> Option<&InfinityModule> {
982 self.module.as_ref()
983 }
984
985 pub fn image(&self) -> Option<&[u8]> {
987 self.image.as_deref()
988 }
989
990 pub fn finish(self) -> Result<StreamedModule> {
993 if !self.partial.is_empty() {
994 let needed = if self.partial.len() >= 20 {
998 let declared = u32::from_le_bytes([
999 self.partial[16],
1000 self.partial[17],
1001 self.partial[18],
1002 self.partial[19],
1003 ]) as u64;
1004 FRAME_HEADER_SIZE as u64 + declared
1005 } else {
1006 FRAME_HEADER_SIZE as u64
1007 };
1008 return Err(Error::Truncated {
1009 needed,
1010 found: self.partial.len() as u64,
1011 });
1012 }
1013 let module = self.module.ok_or(Error::StreamUnterminated)?;
1014 let image = self.image.ok_or(Error::StreamUnterminated)?;
1015 let hash = hash_image(&image)?;
1016 Ok(StreamedModule {
1017 module,
1018 image,
1019 content_hash: hash,
1020 })
1021 }
1022
1023 pub fn frames(&self) -> u64 {
1025 self.frames
1026 }
1027}
1028
1029#[derive(Debug)]
1031pub struct StreamEncoder {
1032 frames: VecDeque<Frame>,
1033 sequence: u64,
1034}
1035
1036impl StreamEncoder {
1037 pub fn new(module: &InfinityModule, chunk: usize) -> Result<Self> {
1039 let frames = encode_stream(module, chunk)?;
1040 let sequence = frames.len() as u64;
1041 Ok(Self {
1042 frames: frames.into(),
1043 sequence,
1044 })
1045 }
1046
1047 pub fn push_delta(&mut self, base: &InfinityModule, ops: &[DeltaOp]) -> Result<()> {
1050 let updated = apply_delta(base, ops)?;
1051 let image = encode_with(&updated, EncodeOpts::default())?;
1052 let mut payload = Vec::new();
1053 payload.extend_from_slice(&(ops.len() as u32).to_le_bytes());
1054 for op in ops {
1055 op.encode(&mut payload)?;
1056 }
1057 self.frames
1058 .push_back(Frame::new(FrameKind::Delta, self.sequence, payload));
1059 self.sequence += 1;
1060 let mut end = Vec::new();
1061 end.extend_from_slice(&hash_image(&image)?);
1062 end.extend_from_slice(&(image.len() as u64).to_le_bytes());
1063 self.frames
1064 .push_back(Frame::new(FrameKind::End, self.sequence, end));
1065 self.sequence += 1;
1066 Ok(())
1067 }
1068
1069 pub fn frames(&self) -> Vec<Frame> {
1071 self.frames.iter().cloned().collect()
1072 }
1073
1074 pub fn to_bytes(&self) -> Vec<u8> {
1076 let mut out = Vec::new();
1077 for frame in &self.frames {
1078 out.extend_from_slice(&frame.to_bytes());
1079 }
1080 out
1081 }
1082}
1083
1084fn write_string(out: &mut Vec<u8>, text: &str) -> Result<()> {
1085 let bytes = text.as_bytes();
1086 if bytes.len() > u16::MAX as usize {
1087 return Err(Error::TooLarge);
1088 }
1089 out.extend_from_slice(&(bytes.len() as u16).to_le_bytes());
1090 out.extend_from_slice(bytes);
1091 Ok(())
1092}
1093
1094fn write_opt_string(out: &mut Vec<u8>, text: Option<&str>) -> Result<()> {
1095 match text {
1096 None => {
1097 out.push(0);
1098 Ok(())
1099 }
1100 Some(text) => {
1101 out.push(1);
1102 write_string(out, text)
1103 }
1104 }
1105}
1106
1107fn take_u8(buf: &[u8], cursor: &mut usize) -> Result<u8> {
1108 if *cursor >= buf.len() {
1109 return Err(Error::StreamDelta("truncated delta payload"));
1110 }
1111 let value = buf[*cursor];
1112 *cursor += 1;
1113 Ok(value)
1114}
1115
1116fn take_u32(buf: &[u8], cursor: &mut usize) -> Result<u32> {
1117 if *cursor + 4 > buf.len() {
1118 return Err(Error::StreamDelta("truncated delta payload"));
1119 }
1120 let value = read_u32(buf, *cursor)?;
1121 *cursor += 4;
1122 Ok(value)
1123}
1124
1125fn take_u64(buf: &[u8], cursor: &mut usize) -> Result<u64> {
1126 if *cursor + 8 > buf.len() {
1127 return Err(Error::StreamDelta("truncated delta payload"));
1128 }
1129 let value = read_u64(buf, *cursor)?;
1130 *cursor += 8;
1131 Ok(value)
1132}
1133
1134fn take_i64(buf: &[u8], cursor: &mut usize) -> Result<i64> {
1135 take_u64(buf, cursor).map(|value| value as i64)
1136}
1137
1138fn take_hash(buf: &[u8], cursor: &mut usize) -> Result<[u8; 32]> {
1139 if *cursor + 32 > buf.len() {
1140 return Err(Error::StreamDelta("truncated delta payload"));
1141 }
1142 let mut hash = [0u8; 32];
1143 hash.copy_from_slice(&buf[*cursor..*cursor + 32]);
1144 *cursor += 32;
1145 Ok(hash)
1146}
1147
1148fn take_string(buf: &[u8], cursor: &mut usize) -> Result<String> {
1149 if *cursor + 2 > buf.len() {
1150 return Err(Error::StreamDelta("truncated delta payload"));
1151 }
1152 let len = u16::from_le_bytes([buf[*cursor], buf[*cursor + 1]]) as usize;
1153 *cursor += 2;
1154 if *cursor + len > buf.len() {
1155 return Err(Error::StreamDelta("truncated delta payload"));
1156 }
1157 let text = std::str::from_utf8(&buf[*cursor..*cursor + len])
1158 .map_err(|_| Error::Utf8 {
1159 offset: *cursor as u64,
1160 })?
1161 .to_string();
1162 *cursor += len;
1163 Ok(text)
1164}
1165
1166fn take_opt_string(buf: &[u8], cursor: &mut usize) -> Result<Option<String>> {
1167 match take_u8(buf, cursor)? {
1168 0 => Ok(None),
1169 1 => Ok(Some(take_string(buf, cursor)?)),
1170 _ => Err(Error::StreamDelta("invalid optional string tag")),
1171 }
1172}
1173
1174fn decode_fragment(
1175 buf: &[u8],
1176 cursor: &mut usize,
1177 limits: crate::layout::Limits,
1178) -> Result<InfinityModule> {
1179 let len = take_u32(buf, cursor)? as usize;
1180 if *cursor + len > buf.len() {
1181 return Err(Error::StreamDelta("truncated fragment"));
1182 }
1183 let bytes = &buf[*cursor..*cursor + len];
1184 *cursor += len;
1185 crate::decode::decode_with(
1189 bytes,
1190 crate::layout::DecodeOpts {
1191 limits,
1192 verify_hash: true,
1193 verify_canonical: false,
1194 },
1195 )
1196}
1197
1198pub fn delta_hash(module: &InfinityModule, ops: &[DeltaOp]) -> Result<[u8; 32]> {
1200 let updated = apply_delta(module, ops)?;
1201 let image = encode_with(&updated, EncodeOpts::default())?;
1202 hash_image(&image)
1203}
1204
1205pub fn decode_stream(bytes: &[u8]) -> Result<StreamedModule> {
1207 let mut decoder = StreamDecoder::new();
1208 decoder.push_bytes(bytes)?;
1209 decoder.finish()
1210}
1211
1212pub fn aligned_section_len(len: usize) -> usize {
1214 align8(len as u64) as usize
1215}
1216
1217pub fn delta_annotations(ops: &[DeltaOp]) -> Vec<Annotation> {
1219 ops.iter()
1220 .filter_map(|op| match op {
1221 DeltaOp::SetAnnotation { node, kind } => Some(Annotation {
1222 node: *node,
1223 kind: kind.clone(),
1224 }),
1225 _ => None,
1226 })
1227 .collect()
1228}