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