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 > self.limits.max_total_bytes {
694 return Err(Error::Limit {
695 kind: "total bytes",
696 limit: self.limits.max_total_bytes,
697 });
698 }
699 self.total_size = Some(size);
700 self.flags = flags;
701 self.begun = true;
702 }
703 FrameKind::Batch => {
704 if !self.begun || self.ended {
705 return Err(Error::StreamState("batch outside a generation"));
706 }
707 if frame.payload.len() < 9 {
708 return Err(Error::StreamFrameLen {
709 declared: 9,
710 available: frame.payload.len() as u64,
711 });
712 }
713 let section = frame.payload[0] as usize;
714 if section >= SECTION_COUNT {
715 return Err(Error::StreamState("unknown section id"));
716 }
717 let count = read_u32(&frame.payload, 1)?;
718 let len = read_u32(&frame.payload, 5)? as usize;
719 if frame.payload.len() != 9 + len {
720 return Err(Error::StreamFrameLen {
721 declared: len as u32,
722 available: (frame.payload.len() - 9) as u64,
723 });
724 }
725 if section != self.cursor && section != self.cursor + 1 {
726 return Err(Error::StreamState("sections out of order"));
727 }
728 if section == self.cursor + 1 {
729 self.cursor = section;
730 }
731 if self.seen[section] && self.counts[section] != count {
732 return Err(Error::StreamState("section count changed mid-stream"));
733 }
734 self.seen[section] = true;
735 self.counts[section] = count;
736 self.sections[section].extend_from_slice(&frame.payload[9..]);
737 let accumulated: usize = self.sections.iter().map(|section| section.len()).sum();
741 let declared = self
742 .total_size
743 .map(|size| size.min(usize::MAX as u64) as usize)
744 .unwrap_or(usize::MAX);
745 if accumulated > declared {
746 return Err(Error::Limit {
747 kind: "stream bytes",
748 limit: declared as u64,
749 });
750 }
751 if accumulated as u64 > self.limits.max_total_bytes {
752 return Err(Error::Limit {
753 kind: "stream bytes",
754 limit: self.limits.max_total_bytes,
755 });
756 }
757 }
758 FrameKind::Delta => {
759 if !self.ended {
760 return Err(Error::StreamState("delta before end"));
761 }
762 let mut cursor = 0usize;
763 let count = take_u32(&frame.payload, &mut cursor)?;
764 if count > crate::layout::Limits::default().max_aux {
765 return Err(Error::StreamDelta("too many delta operations"));
766 }
767 if self.pending_deltas.len().saturating_add(count as usize)
768 > self.limits.max_aux as usize
769 {
770 return Err(Error::Limit {
771 kind: "delta operations",
772 limit: self.limits.max_aux as u64,
773 });
774 }
775 for _ in 0..count {
776 self.pending_deltas.push(DeltaOp::decode(
777 &frame.payload,
778 &mut cursor,
779 self.limits,
780 )?);
781 }
782 if cursor != frame.payload.len() {
783 return Err(Error::StreamDelta("trailing bytes in delta payload"));
784 }
785 }
786 FrameKind::End => {
787 if !self.begun {
788 return Err(Error::StreamState("end before begin"));
789 }
790 if frame.payload.len() != 40 {
791 return Err(Error::StreamFrameLen {
792 declared: 40,
793 available: frame.payload.len() as u64,
794 });
795 }
796 let mut hash = [0u8; 32];
797 hash.copy_from_slice(&frame.payload[..32]);
798 let size = read_u64(&frame.payload, 32)?;
799 if self.ended {
800 let base = self
802 .module
803 .as_ref()
804 .ok_or(Error::StreamState("no module to patch"))?;
805 let updated = apply_delta(base, &self.pending_deltas)?;
806 self.pending_deltas.clear();
807 let image = encode_with(&updated, EncodeOpts::default())?;
808 if image.len() as u64 != size {
809 return Err(Error::StreamDelta("delta changed the declared size"));
810 }
811 if hash_image(&image)? != hash {
812 return Err(Error::StreamHashMismatch);
813 }
814 let validated = Validated::new(
815 &image,
816 crate::layout::DecodeOpts {
817 limits: self.limits,
818 verify_hash: true,
819 verify_canonical: false,
820 },
821 )?;
822 let module = crate::decode::decode_validated(&validated)?;
823 self.module = Some(module);
824 self.image = Some(image);
825 return Ok(());
826 }
827 let image = self.assemble(size)?;
828 if hash_image(&image)? != hash {
829 return Err(Error::StreamHashMismatch);
830 }
831 let validated = Validated::new(
832 &image,
833 crate::layout::DecodeOpts {
834 limits: self.limits,
835 verify_hash: true,
836 verify_canonical: false,
837 },
838 )?;
839 let module = crate::decode::decode_validated(&validated)?;
840 self.module = Some(module);
841 self.image = Some(image);
842 self.ended = true;
843 }
844 }
845 Ok(())
846 }
847
848 fn assemble(&self, size: u64) -> Result<Vec<u8>> {
849 if size < HEADER_SIZE as u64 {
850 return Err(Error::StreamDelta("declared size smaller than a header"));
851 }
852 if self.seen.iter().any(|seen| !seen) {
853 return Err(Error::StreamState("generation missing a section"));
854 }
855 let mut image = vec![0u8; HEADER_SIZE];
856 let mut header = crate::layout::Header {
857 major: crate::layout::FORMAT_MAJOR,
858 minor: crate::layout::FORMAT_MINOR,
859 flags: self.flags,
860 header_size: HEADER_SIZE as u16,
861 content_hash: [0u8; 32],
862 dialect_hash: crate::layout::dialect_hash(),
863 total_size: 0,
864 sections: [crate::layout::SectionSpan::default(); SECTION_COUNT],
865 };
866 for section in Section::ALL {
867 while image.len() % 8 != 0 {
868 image.push(0);
869 }
870 header.span_mut(section).offset = image.len() as u32;
871 header.span_mut(section).count = self.counts[section as usize];
872 image.extend_from_slice(&self.sections[section as usize]);
873 }
874 while image.len() % 8 != 0 {
875 image.push(0);
876 }
877 header.total_size = image.len() as u64;
878 header.write(&mut image);
879 let hash = hash_image(&image)?;
880 image[16..48].copy_from_slice(&hash);
881 if image.len() as u64 != size {
882 return Err(Error::StreamDelta(
883 "reassembled size disagrees with the end frame",
884 ));
885 }
886 Ok(image)
887 }
888
889 pub fn module(&self) -> Option<&InfinityModule> {
891 self.module.as_ref()
892 }
893
894 pub fn image(&self) -> Option<&[u8]> {
896 self.image.as_deref()
897 }
898
899 pub fn finish(self) -> Result<StreamedModule> {
902 if !self.partial.is_empty() {
903 let needed = if self.partial.len() >= 20 {
907 let declared = u32::from_le_bytes([
908 self.partial[16],
909 self.partial[17],
910 self.partial[18],
911 self.partial[19],
912 ]) as u64;
913 FRAME_HEADER_SIZE as u64 + declared
914 } else {
915 FRAME_HEADER_SIZE as u64
916 };
917 return Err(Error::Truncated {
918 needed,
919 found: self.partial.len() as u64,
920 });
921 }
922 let module = self.module.ok_or(Error::StreamUnterminated)?;
923 let image = self.image.ok_or(Error::StreamUnterminated)?;
924 let hash = hash_image(&image)?;
925 Ok(StreamedModule {
926 module,
927 image,
928 content_hash: hash,
929 })
930 }
931
932 pub fn frames(&self) -> u64 {
934 self.frames
935 }
936}
937
938#[derive(Debug)]
940pub struct StreamEncoder {
941 frames: VecDeque<Frame>,
942 sequence: u64,
943}
944
945impl StreamEncoder {
946 pub fn new(module: &InfinityModule, chunk: usize) -> Result<Self> {
948 let frames = encode_stream(module, chunk)?;
949 let sequence = frames.len() as u64;
950 Ok(Self {
951 frames: frames.into(),
952 sequence,
953 })
954 }
955
956 pub fn push_delta(&mut self, base: &InfinityModule, ops: &[DeltaOp]) -> Result<()> {
959 let updated = apply_delta(base, ops)?;
960 let image = encode_with(&updated, EncodeOpts::default())?;
961 let mut payload = Vec::new();
962 payload.extend_from_slice(&(ops.len() as u32).to_le_bytes());
963 for op in ops {
964 op.encode(&mut payload)?;
965 }
966 self.frames
967 .push_back(Frame::new(FrameKind::Delta, self.sequence, payload));
968 self.sequence += 1;
969 let mut end = Vec::new();
970 end.extend_from_slice(&hash_image(&image)?);
971 end.extend_from_slice(&(image.len() as u64).to_le_bytes());
972 self.frames
973 .push_back(Frame::new(FrameKind::End, self.sequence, end));
974 self.sequence += 1;
975 Ok(())
976 }
977
978 pub fn frames(&self) -> Vec<Frame> {
980 self.frames.iter().cloned().collect()
981 }
982
983 pub fn to_bytes(&self) -> Vec<u8> {
985 let mut out = Vec::new();
986 for frame in &self.frames {
987 out.extend_from_slice(&frame.to_bytes());
988 }
989 out
990 }
991}
992
993fn write_string(out: &mut Vec<u8>, text: &str) -> Result<()> {
994 let bytes = text.as_bytes();
995 if bytes.len() > u16::MAX as usize {
996 return Err(Error::TooLarge);
997 }
998 out.extend_from_slice(&(bytes.len() as u16).to_le_bytes());
999 out.extend_from_slice(bytes);
1000 Ok(())
1001}
1002
1003fn write_opt_string(out: &mut Vec<u8>, text: Option<&str>) -> Result<()> {
1004 match text {
1005 None => {
1006 out.push(0);
1007 Ok(())
1008 }
1009 Some(text) => {
1010 out.push(1);
1011 write_string(out, text)
1012 }
1013 }
1014}
1015
1016fn take_u8(buf: &[u8], cursor: &mut usize) -> Result<u8> {
1017 if *cursor >= buf.len() {
1018 return Err(Error::StreamDelta("truncated delta payload"));
1019 }
1020 let value = buf[*cursor];
1021 *cursor += 1;
1022 Ok(value)
1023}
1024
1025fn take_u32(buf: &[u8], cursor: &mut usize) -> Result<u32> {
1026 if *cursor + 4 > buf.len() {
1027 return Err(Error::StreamDelta("truncated delta payload"));
1028 }
1029 let value = read_u32(buf, *cursor)?;
1030 *cursor += 4;
1031 Ok(value)
1032}
1033
1034fn take_u64(buf: &[u8], cursor: &mut usize) -> Result<u64> {
1035 if *cursor + 8 > buf.len() {
1036 return Err(Error::StreamDelta("truncated delta payload"));
1037 }
1038 let value = read_u64(buf, *cursor)?;
1039 *cursor += 8;
1040 Ok(value)
1041}
1042
1043fn take_i64(buf: &[u8], cursor: &mut usize) -> Result<i64> {
1044 take_u64(buf, cursor).map(|value| value as i64)
1045}
1046
1047fn take_hash(buf: &[u8], cursor: &mut usize) -> Result<[u8; 32]> {
1048 if *cursor + 32 > buf.len() {
1049 return Err(Error::StreamDelta("truncated delta payload"));
1050 }
1051 let mut hash = [0u8; 32];
1052 hash.copy_from_slice(&buf[*cursor..*cursor + 32]);
1053 *cursor += 32;
1054 Ok(hash)
1055}
1056
1057fn take_string(buf: &[u8], cursor: &mut usize) -> Result<String> {
1058 if *cursor + 2 > buf.len() {
1059 return Err(Error::StreamDelta("truncated delta payload"));
1060 }
1061 let len = u16::from_le_bytes([buf[*cursor], buf[*cursor + 1]]) as usize;
1062 *cursor += 2;
1063 if *cursor + len > buf.len() {
1064 return Err(Error::StreamDelta("truncated delta payload"));
1065 }
1066 let text = std::str::from_utf8(&buf[*cursor..*cursor + len])
1067 .map_err(|_| Error::Utf8 {
1068 offset: *cursor as u64,
1069 })?
1070 .to_string();
1071 *cursor += len;
1072 Ok(text)
1073}
1074
1075fn take_opt_string(buf: &[u8], cursor: &mut usize) -> Result<Option<String>> {
1076 match take_u8(buf, cursor)? {
1077 0 => Ok(None),
1078 1 => Ok(Some(take_string(buf, cursor)?)),
1079 _ => Err(Error::StreamDelta("invalid optional string tag")),
1080 }
1081}
1082
1083fn decode_fragment(
1084 buf: &[u8],
1085 cursor: &mut usize,
1086 limits: crate::layout::Limits,
1087) -> Result<InfinityModule> {
1088 let len = take_u32(buf, cursor)? as usize;
1089 if *cursor + len > buf.len() {
1090 return Err(Error::StreamDelta("truncated fragment"));
1091 }
1092 let bytes = &buf[*cursor..*cursor + len];
1093 *cursor += len;
1094 crate::decode::decode_with(
1098 bytes,
1099 crate::layout::DecodeOpts {
1100 limits,
1101 verify_hash: true,
1102 verify_canonical: false,
1103 },
1104 )
1105}
1106
1107pub fn delta_hash(module: &InfinityModule, ops: &[DeltaOp]) -> Result<[u8; 32]> {
1109 let updated = apply_delta(module, ops)?;
1110 let image = encode_with(&updated, EncodeOpts::default())?;
1111 hash_image(&image)
1112}
1113
1114pub fn decode_stream(bytes: &[u8]) -> Result<StreamedModule> {
1116 let mut decoder = StreamDecoder::new();
1117 decoder.push_bytes(bytes)?;
1118 decoder.finish()
1119}
1120
1121pub fn aligned_section_len(len: usize) -> usize {
1123 align8(len as u64) as usize
1124}
1125
1126pub fn delta_annotations(ops: &[DeltaOp]) -> Vec<Annotation> {
1128 ops.iter()
1129 .filter_map(|op| match op {
1130 DeltaOp::SetAnnotation { node, kind } => Some(Annotation {
1131 node: *node,
1132 kind: kind.clone(),
1133 }),
1134 _ => None,
1135 })
1136 .collect()
1137}