#![allow(missing_docs)]
use std::collections::VecDeque;
use crate::canonical::remap_node;
use crate::encode::{encode_with, hash_image};
use crate::error::{Error, Result};
use crate::ir::{Annotation, AnnotationKind, InfinityModule, NodeId, Provenance};
use crate::layout::{
align8, read_u32, read_u64, EncodeOpts, Section, HEADER_SIZE, SECTION_COUNT, STREAM_MAGIC,
};
use crate::validate::Validated;
pub const STREAM_VERSION: u8 = 1;
pub const FRAME_HEADER_SIZE: usize = 32;
pub const DEFAULT_CHUNK: usize = 64 << 10;
pub const MAX_FRAME_PAYLOAD: u32 = 64 << 20;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FrameKind {
Begin,
Batch,
Delta,
End,
}
impl FrameKind {
pub const fn code(self) -> u8 {
match self {
FrameKind::Begin => 1,
FrameKind::Batch => 2,
FrameKind::Delta => 3,
FrameKind::End => 4,
}
}
pub const fn from_code(code: u8) -> Option<Self> {
match code {
1 => Some(FrameKind::Begin),
2 => Some(FrameKind::Batch),
3 => Some(FrameKind::Delta),
4 => Some(FrameKind::End),
_ => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Frame {
pub kind: FrameKind,
pub sequence: u64,
pub payload: Vec<u8>,
}
impl Frame {
pub fn new(kind: FrameKind, sequence: u64, payload: Vec<u8>) -> Self {
Frame {
kind,
sequence,
payload,
}
}
pub fn to_bytes(&self) -> Vec<u8> {
let mut out = Vec::with_capacity(FRAME_HEADER_SIZE + self.payload.len());
out.extend_from_slice(&STREAM_MAGIC);
out.push(STREAM_VERSION);
out.push(self.kind.code());
out.extend_from_slice(&0u16.to_le_bytes());
out.extend_from_slice(&self.sequence.to_le_bytes());
out.extend_from_slice(&(self.payload.len() as u32).to_le_bytes());
out.extend_from_slice(&0u32.to_le_bytes());
let checksum = frame_checksum(&self.payload);
out.extend_from_slice(&checksum[..8]);
out.extend_from_slice(&self.payload);
out
}
pub fn parse(buf: &[u8]) -> Result<(Frame, usize)> {
if buf.len() < 4 || buf[0..4] != STREAM_MAGIC {
return Err(Error::StreamMagic);
}
if buf.len() < FRAME_HEADER_SIZE {
return Err(Error::Truncated {
needed: FRAME_HEADER_SIZE as u64,
found: buf.len() as u64,
});
}
let version = buf[4];
if version != STREAM_VERSION {
return Err(Error::StreamVersion(version));
}
let kind = FrameKind::from_code(buf[5]).ok_or(Error::StreamFrameKind(buf[5]))?;
let flags = u16::from_le_bytes([buf[6], buf[7]]);
if flags != 0 {
return Err(Error::UnknownFlags(flags));
}
let sequence = u64::from_le_bytes(buf[8..16].try_into().expect("8 bytes"));
let len = u32::from_le_bytes(buf[16..20].try_into().expect("4 bytes"));
if len > MAX_FRAME_PAYLOAD {
return Err(Error::StreamFrameLen {
declared: len,
available: buf.len() as u64,
});
}
let reserved = u32::from_le_bytes(buf[20..24].try_into().expect("4 bytes"));
if reserved != 0 {
return Err(Error::ReservedNonZero {
section: "stream-frame",
offset: 20,
});
}
let payload_start = FRAME_HEADER_SIZE;
let payload_end = payload_start + len as usize;
if payload_end > buf.len() {
return Err(Error::StreamFrameLen {
declared: len,
available: (buf.len() - payload_start) as u64,
});
}
let payload = &buf[payload_start..payload_end];
let expected = frame_checksum(payload);
if buf[24..32] != expected[..8] {
return Err(Error::HashMismatch);
}
Ok((
Frame {
kind,
sequence,
payload: payload.to_vec(),
},
payload_end,
))
}
}
pub fn frame_checksum(payload: &[u8]) -> [u8; 32] {
*blake3::hash(payload).as_bytes()
}
pub fn encode_stream(module: &InfinityModule, chunk: usize) -> Result<Vec<Frame>> {
encode_stream_with(module, chunk, EncodeOpts::default())
}
pub fn encode_stream_with(
module: &InfinityModule,
chunk: usize,
opts: EncodeOpts,
) -> Result<Vec<Frame>> {
if chunk == 0 {
return Err(Error::InvalidArgument("chunk size must be non-zero"));
}
let image = encode_with(module, opts)?;
let header = crate::layout::Header::parse(&image)?;
let mut frames = Vec::new();
let mut sequence = 0u64;
let mut begin = Vec::with_capacity(24);
begin.extend_from_slice(&header.major.to_le_bytes());
begin.extend_from_slice(&header.minor.to_le_bytes());
begin.extend_from_slice(&header.flags.to_le_bytes());
begin.extend_from_slice(&0u16.to_le_bytes());
begin.extend_from_slice(&header.dialect_hash.to_le_bytes());
begin.extend_from_slice(&header.total_size.to_le_bytes());
frames.push(Frame::new(FrameKind::Begin, sequence, begin));
sequence += 1;
for section in Section::ALL {
let span = header.span(section);
let start = span.offset as usize;
let len = section_body_len(&image, &header, section);
let body = &image[start..start + len];
if body.is_empty() {
let mut payload = Vec::with_capacity(9);
payload.push(section as u8);
payload.extend_from_slice(&span.count.to_le_bytes());
payload.extend_from_slice(&0u32.to_le_bytes());
frames.push(Frame::new(FrameKind::Batch, sequence, payload));
sequence += 1;
continue;
}
for piece in body.chunks(chunk) {
let mut payload = Vec::with_capacity(9 + piece.len());
payload.push(section as u8);
payload.extend_from_slice(&span.count.to_le_bytes());
payload.extend_from_slice(&(piece.len() as u32).to_le_bytes());
payload.extend_from_slice(piece);
frames.push(Frame::new(FrameKind::Batch, sequence, payload));
sequence += 1;
}
}
let mut end = Vec::with_capacity(40);
let hash = hash_image(&image)?;
end.extend_from_slice(&hash);
end.extend_from_slice(&(image.len() as u64).to_le_bytes());
frames.push(Frame::new(FrameKind::End, sequence, end));
Ok(frames)
}
fn section_body_len(image: &[u8], header: &crate::layout::Header, section: Section) -> usize {
let span = header.span(section);
let start = span.offset as usize;
let end = match section {
Section::Strings => {
if span.count == 0 {
start
} else {
let heap = read_u32(image, start + span.count as usize * 4).unwrap_or(0) as usize;
start + (span.count as usize + 1) * 4 + heap
}
}
Section::Blobs => {
let mut heap = 0usize;
for index in 0..span.count {
let len = read_u32(image, start + index as usize * 40 + 36).unwrap_or(0) as usize;
heap += len;
}
start + span.count as usize * 40 + heap
}
Section::Provenance => start + span.count as usize,
Section::Secrets => {
let mut heap = 0usize;
for index in 0..span.count {
let len = read_u32(image, start + index as usize * 8 + 4).unwrap_or(0) as usize;
heap += len;
}
start + span.count as usize * 8 + heap
}
Section::Extensions => {
let mut end = start + span.count as usize * 16;
for index in 0..span.count {
let dir = start + index as usize * 16;
let offset = read_u32(image, dir + 8).unwrap_or(0) as usize;
let len = read_u32(image, dir + 12).unwrap_or(0) as usize;
end = end.max(offset + len);
}
end
}
Section::Nodes => start + span.count as usize * 32,
Section::Children => start + span.count as usize * 4,
Section::Fields => start + span.count as usize * 16,
Section::Parts => start + span.count as usize * 8,
Section::Edges => start + span.count as usize * 8,
Section::Annotations => start + span.count as usize * 16,
Section::Units => start + span.count as usize * 8,
Section::Types => start + span.count as usize * 16,
Section::Aux => start + span.count as usize * 8,
Section::Numbers => start + span.count as usize * 8,
Section::Hashes => start + span.count as usize * 32,
Section::Bindings => start + span.count as usize * 8,
};
end.saturating_sub(start)
}
#[derive(Debug, Clone, PartialEq)]
pub enum DeltaOp {
SetRoot(Option<NodeId>),
PutNode {
target: NodeId,
fragment: Box<InfinityModule>,
},
SetBinding { name: String, node: NodeId },
RemoveBinding { name: String },
SetAnnotation { node: NodeId, kind: AnnotationKind },
ClearAnnotations,
SetProvenance(Option<Provenance>),
}
impl DeltaOp {
fn encode(&self, out: &mut Vec<u8>) -> Result<()> {
match self {
DeltaOp::SetRoot(root) => {
out.push(0);
out.extend_from_slice(&root.map(|id| id.0).unwrap_or(u32::MAX).to_le_bytes());
}
DeltaOp::PutNode { target, fragment } => {
out.push(1);
out.extend_from_slice(&target.0.to_le_bytes());
let image = encode_with(fragment, EncodeOpts::default())?;
out.extend_from_slice(&(image.len() as u32).to_le_bytes());
out.extend_from_slice(&image);
}
DeltaOp::SetBinding { name, node } => {
out.push(3);
let bytes = name.as_bytes();
out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
out.extend_from_slice(bytes);
out.extend_from_slice(&node.0.to_le_bytes());
}
DeltaOp::RemoveBinding { name } => {
out.push(4);
let bytes = name.as_bytes();
out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
out.extend_from_slice(bytes);
}
DeltaOp::SetAnnotation { node, kind } => {
out.push(5);
out.extend_from_slice(&node.0.to_le_bytes());
match kind {
AnnotationKind::Intent(text) => {
out.push(0);
write_string(out, text)?;
}
AnnotationKind::Constraint { source, expr } => {
out.push(1);
write_string(out, source)?;
out.extend_from_slice(&expr.0.to_le_bytes());
}
AnnotationKind::AiImmutable(flag) => {
out.push(2);
out.push(u8::from(*flag));
}
}
}
DeltaOp::ClearAnnotations => out.push(6),
DeltaOp::SetProvenance(prov) => {
out.push(7);
match prov {
None => out.push(0),
Some(prov) => {
out.push(1);
write_opt_string(out, prov.producer.as_deref())?;
write_opt_string(out, prov.model.as_deref())?;
write_opt_string(out, prov.tool.as_deref())?;
out.extend_from_slice(&prov.timestamp_ms.to_le_bytes());
match prov.parent {
None => out.push(0),
Some(hash) => {
out.push(1);
out.extend_from_slice(&hash);
}
}
out.extend_from_slice(&(prov.entries.len() as u32).to_le_bytes());
for entry in &prov.entries {
write_string(out, &entry.key)?;
match &entry.value {
crate::ir::ProvValue::Text(text) => {
out.push(0);
write_string(out, text)?;
}
crate::ir::ProvValue::Int(value) => {
out.push(1);
out.extend_from_slice(&value.to_le_bytes());
}
crate::ir::ProvValue::Flag(flag) => {
out.push(2);
out.push(u8::from(*flag));
}
crate::ir::ProvValue::Hash(hash) => {
out.push(3);
out.extend_from_slice(hash);
}
}
}
}
}
}
}
Ok(())
}
fn decode(buf: &[u8], cursor: &mut usize, limits: crate::layout::Limits) -> Result<DeltaOp> {
let tag = take_u8(buf, cursor)?;
let op = match tag {
0 => {
let id = take_u32(buf, cursor)?;
DeltaOp::SetRoot(if id == u32::MAX {
None
} else {
Some(NodeId(id))
})
}
1 => {
let target = NodeId(take_u32(buf, cursor)?);
let fragment = Box::new(decode_fragment(buf, cursor, limits)?);
DeltaOp::PutNode { target, fragment }
}
3 => {
let name = take_string(buf, cursor)?;
let node = NodeId(take_u32(buf, cursor)?);
DeltaOp::SetBinding { name, node }
}
4 => {
let name = take_string(buf, cursor)?;
DeltaOp::RemoveBinding { name }
}
5 => {
let node = NodeId(take_u32(buf, cursor)?);
let kind_tag = take_u8(buf, cursor)?;
let kind = match kind_tag {
0 => AnnotationKind::Intent(take_string(buf, cursor)?),
1 => {
let source = take_string(buf, cursor)?;
let expr = NodeId(take_u32(buf, cursor)?);
AnnotationKind::Constraint { source, expr }
}
2 => AnnotationKind::AiImmutable(take_u8(buf, cursor)? != 0),
_ => return Err(Error::StreamDelta("unknown annotation kind")),
};
DeltaOp::SetAnnotation { node, kind }
}
6 => DeltaOp::ClearAnnotations,
7 => {
if take_u8(buf, cursor)? == 0 {
DeltaOp::SetProvenance(None)
} else {
let producer = take_opt_string(buf, cursor)?;
let model = take_opt_string(buf, cursor)?;
let tool = take_opt_string(buf, cursor)?;
let timestamp_ms = take_u64(buf, cursor)?;
let parent = if take_u8(buf, cursor)? == 0 {
None
} else {
Some(take_hash(buf, cursor)?)
};
let count = take_u32(buf, cursor)?;
let mut entries = Vec::with_capacity(count.min(4096) as usize);
for _ in 0..count {
let key = take_string(buf, cursor)?;
let value = match take_u8(buf, cursor)? {
0 => crate::ir::ProvValue::Text(take_string(buf, cursor)?),
1 => crate::ir::ProvValue::Int(take_i64(buf, cursor)?),
2 => crate::ir::ProvValue::Flag(take_u8(buf, cursor)? != 0),
3 => crate::ir::ProvValue::Hash(take_hash(buf, cursor)?),
_ => return Err(Error::StreamDelta("unknown provenance value kind")),
};
entries.push(crate::ir::ProvEntry { key, value });
}
DeltaOp::SetProvenance(Some(Provenance {
producer,
timestamp_ms,
model,
tool,
parent,
entries,
}))
}
}
_ => return Err(Error::StreamDelta("unknown delta operation")),
};
Ok(op)
}
}
pub fn apply_delta(module: &InfinityModule, ops: &[DeltaOp]) -> Result<InfinityModule> {
let mut out = module.clone();
for op in ops {
match op {
DeltaOp::SetRoot(root) => {
if let Some(root) = root {
if out.get(*root).is_none() {
return Err(Error::StreamDelta("root node does not exist"));
}
}
out.set_root(*root);
}
DeltaOp::PutNode { target, fragment } => {
if out.get(*target).is_none() {
return Err(Error::StreamDelta("target node does not exist"));
}
if fragment.node_count() == 0 {
return Err(Error::StreamDelta("fragment is empty"));
}
let replacement = merge_fragment(&mut out, fragment)?;
let content = out
.get(replacement)
.cloned()
.ok_or(Error::StreamDelta("fragment root missing"))?;
*out.node_mut(*target)? = content;
}
DeltaOp::SetBinding { name, node } => {
if out.get(*node).is_none() {
return Err(Error::StreamDelta("binding target does not exist"));
}
out.bind(name.clone(), *node);
}
DeltaOp::RemoveBinding { name } => {
out.remove_binding(name);
}
DeltaOp::SetAnnotation { node, kind } => {
if out.get(*node).is_none() {
return Err(Error::StreamDelta("annotation subject does not exist"));
}
out.annotate(*node, kind.clone());
}
DeltaOp::ClearAnnotations => out.clear_annotations(),
DeltaOp::SetProvenance(prov) => out.set_provenance(prov.clone()),
}
}
Ok(out)
}
fn merge_fragment(base: &mut InfinityModule, fragment: &InfinityModule) -> Result<NodeId> {
let base_count = base.node_count();
let map: Vec<NodeId> = (0..fragment.node_count())
.map(|index| NodeId(base_count as u32 + index as u32))
.collect();
for node in fragment.nodes() {
base.push(remap_node(node, &map)?);
}
Ok(map[fragment.root().unwrap_or(NodeId(0)).index()])
}
#[derive(Debug, Clone, PartialEq)]
pub struct StreamedModule {
pub module: InfinityModule,
pub image: Vec<u8>,
pub content_hash: [u8; 32],
}
#[derive(Debug, Default)]
pub struct StreamDecoder {
limits: crate::layout::Limits,
sections: [Vec<u8>; SECTION_COUNT],
counts: [u32; SECTION_COUNT],
seen: [bool; SECTION_COUNT],
begun: bool,
ended: bool,
expected: u64,
flags: u16,
total_size: Option<u64>,
module: Option<InfinityModule>,
image: Option<Vec<u8>>,
pending_deltas: Vec<DeltaOp>,
frames: u64,
partial: Vec<u8>,
cursor: usize,
}
impl StreamDecoder {
pub fn new() -> Self {
Self {
limits: crate::layout::Limits::default(),
sections: Default::default(),
counts: [0; SECTION_COUNT],
seen: [false; SECTION_COUNT],
begun: false,
ended: false,
expected: 0,
flags: 0,
total_size: None,
module: None,
image: None,
pending_deltas: Vec::new(),
frames: 0,
partial: Vec::new(),
cursor: 0,
}
}
pub fn with_limits(mut self, limits: crate::layout::Limits) -> Self {
self.limits = limits;
self
}
pub fn push_bytes(&mut self, buf: &[u8]) -> Result<()> {
self.partial.extend_from_slice(buf);
loop {
if self.partial.len() < FRAME_HEADER_SIZE {
return Ok(());
}
match Frame::parse(&self.partial) {
Ok((frame, consumed)) => {
self.push(&frame)?;
self.partial.drain(..consumed);
}
Err(Error::Truncated { .. }) => return Ok(()),
Err(Error::StreamFrameLen { declared, .. }) if declared <= MAX_FRAME_PAYLOAD => {
return Ok(())
}
Err(error) => return Err(error),
}
}
}
pub fn push(&mut self, frame: &Frame) -> Result<()> {
if frame.sequence != self.expected {
return Err(Error::StreamSequence {
expected: self.expected,
found: frame.sequence,
});
}
self.expected += 1;
self.frames += 1;
match frame.kind {
FrameKind::Begin => {
if self.begun {
return Err(Error::StreamState("duplicate begin frame"));
}
if self.ended {
return Err(Error::StreamState("begin after end"));
}
if frame.payload.len() != 24 {
return Err(Error::StreamFrameLen {
declared: 24,
available: frame.payload.len() as u64,
});
}
let major = u16::from_le_bytes([frame.payload[0], frame.payload[1]]);
let minor = u16::from_le_bytes([frame.payload[2], frame.payload[3]]);
if major != crate::layout::FORMAT_MAJOR || minor > crate::layout::FORMAT_MINOR {
return Err(Error::VersionMajor(major));
}
let flags = u16::from_le_bytes([frame.payload[4], frame.payload[5]]);
if flags & !crate::layout::FLAG_KNOWN_MASK != 0 {
return Err(Error::UnknownFlags(flags & !crate::layout::FLAG_KNOWN_MASK));
}
if frame.payload[6..8] != [0, 0] {
return Err(Error::ReservedNonZero {
section: "stream-begin",
offset: 6,
});
}
let dialect = read_u64(&frame.payload, 8)?;
if dialect != crate::layout::dialect_hash() {
return Err(Error::DialectMismatch(dialect));
}
let size = read_u64(&frame.payload, 16)?;
if size > self.limits.max_total_bytes {
return Err(Error::Limit {
kind: "total bytes",
limit: self.limits.max_total_bytes,
});
}
self.total_size = Some(size);
self.flags = flags;
self.begun = true;
}
FrameKind::Batch => {
if !self.begun || self.ended {
return Err(Error::StreamState("batch outside a generation"));
}
if frame.payload.len() < 9 {
return Err(Error::StreamFrameLen {
declared: 9,
available: frame.payload.len() as u64,
});
}
let section = frame.payload[0] as usize;
if section >= SECTION_COUNT {
return Err(Error::StreamState("unknown section id"));
}
let count = read_u32(&frame.payload, 1)?;
let len = read_u32(&frame.payload, 5)? as usize;
if frame.payload.len() != 9 + len {
return Err(Error::StreamFrameLen {
declared: len as u32,
available: (frame.payload.len() - 9) as u64,
});
}
if section != self.cursor && section != self.cursor + 1 {
return Err(Error::StreamState("sections out of order"));
}
if section == self.cursor + 1 {
self.cursor = section;
}
if self.seen[section] && self.counts[section] != count {
return Err(Error::StreamState("section count changed mid-stream"));
}
self.seen[section] = true;
self.counts[section] = count;
self.sections[section].extend_from_slice(&frame.payload[9..]);
let accumulated: usize = self.sections.iter().map(|section| section.len()).sum();
let declared = self
.total_size
.map(|size| size.min(usize::MAX as u64) as usize)
.unwrap_or(usize::MAX);
if accumulated > declared {
return Err(Error::Limit {
kind: "stream bytes",
limit: declared as u64,
});
}
if accumulated as u64 > self.limits.max_total_bytes {
return Err(Error::Limit {
kind: "stream bytes",
limit: self.limits.max_total_bytes,
});
}
}
FrameKind::Delta => {
if !self.ended {
return Err(Error::StreamState("delta before end"));
}
let mut cursor = 0usize;
let count = take_u32(&frame.payload, &mut cursor)?;
if count > crate::layout::Limits::default().max_aux {
return Err(Error::StreamDelta("too many delta operations"));
}
if self.pending_deltas.len().saturating_add(count as usize)
> self.limits.max_aux as usize
{
return Err(Error::Limit {
kind: "delta operations",
limit: self.limits.max_aux as u64,
});
}
for _ in 0..count {
self.pending_deltas.push(DeltaOp::decode(
&frame.payload,
&mut cursor,
self.limits,
)?);
}
if cursor != frame.payload.len() {
return Err(Error::StreamDelta("trailing bytes in delta payload"));
}
}
FrameKind::End => {
if !self.begun {
return Err(Error::StreamState("end before begin"));
}
if frame.payload.len() != 40 {
return Err(Error::StreamFrameLen {
declared: 40,
available: frame.payload.len() as u64,
});
}
let mut hash = [0u8; 32];
hash.copy_from_slice(&frame.payload[..32]);
let size = read_u64(&frame.payload, 32)?;
if self.ended {
let base = self
.module
.as_ref()
.ok_or(Error::StreamState("no module to patch"))?;
let updated = apply_delta(base, &self.pending_deltas)?;
self.pending_deltas.clear();
let image = encode_with(&updated, EncodeOpts::default())?;
if image.len() as u64 != size {
return Err(Error::StreamDelta("delta changed the declared size"));
}
if hash_image(&image)? != hash {
return Err(Error::StreamHashMismatch);
}
let validated = Validated::new(
&image,
crate::layout::DecodeOpts {
limits: self.limits,
verify_hash: true,
verify_canonical: false,
},
)?;
let module = crate::decode::decode_validated(&validated)?;
self.module = Some(module);
self.image = Some(image);
return Ok(());
}
let image = self.assemble(size)?;
if hash_image(&image)? != hash {
return Err(Error::StreamHashMismatch);
}
let validated = Validated::new(
&image,
crate::layout::DecodeOpts {
limits: self.limits,
verify_hash: true,
verify_canonical: false,
},
)?;
let module = crate::decode::decode_validated(&validated)?;
self.module = Some(module);
self.image = Some(image);
self.ended = true;
}
}
Ok(())
}
fn assemble(&self, size: u64) -> Result<Vec<u8>> {
if size < HEADER_SIZE as u64 {
return Err(Error::StreamDelta("declared size smaller than a header"));
}
if self.seen.iter().any(|seen| !seen) {
return Err(Error::StreamState("generation missing a section"));
}
let mut image = vec![0u8; HEADER_SIZE];
let mut header = crate::layout::Header {
major: crate::layout::FORMAT_MAJOR,
minor: crate::layout::FORMAT_MINOR,
flags: self.flags,
header_size: HEADER_SIZE as u16,
content_hash: [0u8; 32],
dialect_hash: crate::layout::dialect_hash(),
total_size: 0,
sections: [crate::layout::SectionSpan::default(); SECTION_COUNT],
};
for section in Section::ALL {
while image.len() % 8 != 0 {
image.push(0);
}
header.span_mut(section).offset = image.len() as u32;
header.span_mut(section).count = self.counts[section as usize];
image.extend_from_slice(&self.sections[section as usize]);
}
while image.len() % 8 != 0 {
image.push(0);
}
header.total_size = image.len() as u64;
header.write(&mut image);
let hash = hash_image(&image)?;
image[16..48].copy_from_slice(&hash);
if image.len() as u64 != size {
return Err(Error::StreamDelta(
"reassembled size disagrees with the end frame",
));
}
Ok(image)
}
pub fn module(&self) -> Option<&InfinityModule> {
self.module.as_ref()
}
pub fn image(&self) -> Option<&[u8]> {
self.image.as_deref()
}
pub fn finish(self) -> Result<StreamedModule> {
if !self.partial.is_empty() {
let needed = if self.partial.len() >= 20 {
let declared = u32::from_le_bytes([
self.partial[16],
self.partial[17],
self.partial[18],
self.partial[19],
]) as u64;
FRAME_HEADER_SIZE as u64 + declared
} else {
FRAME_HEADER_SIZE as u64
};
return Err(Error::Truncated {
needed,
found: self.partial.len() as u64,
});
}
let module = self.module.ok_or(Error::StreamUnterminated)?;
let image = self.image.ok_or(Error::StreamUnterminated)?;
let hash = hash_image(&image)?;
Ok(StreamedModule {
module,
image,
content_hash: hash,
})
}
pub fn frames(&self) -> u64 {
self.frames
}
}
#[derive(Debug)]
pub struct StreamEncoder {
frames: VecDeque<Frame>,
sequence: u64,
}
impl StreamEncoder {
pub fn new(module: &InfinityModule, chunk: usize) -> Result<Self> {
let frames = encode_stream(module, chunk)?;
let sequence = frames.len() as u64;
Ok(Self {
frames: frames.into(),
sequence,
})
}
pub fn push_delta(&mut self, base: &InfinityModule, ops: &[DeltaOp]) -> Result<()> {
let updated = apply_delta(base, ops)?;
let image = encode_with(&updated, EncodeOpts::default())?;
let mut payload = Vec::new();
payload.extend_from_slice(&(ops.len() as u32).to_le_bytes());
for op in ops {
op.encode(&mut payload)?;
}
self.frames
.push_back(Frame::new(FrameKind::Delta, self.sequence, payload));
self.sequence += 1;
let mut end = Vec::new();
end.extend_from_slice(&hash_image(&image)?);
end.extend_from_slice(&(image.len() as u64).to_le_bytes());
self.frames
.push_back(Frame::new(FrameKind::End, self.sequence, end));
self.sequence += 1;
Ok(())
}
pub fn frames(&self) -> Vec<Frame> {
self.frames.iter().cloned().collect()
}
pub fn to_bytes(&self) -> Vec<u8> {
let mut out = Vec::new();
for frame in &self.frames {
out.extend_from_slice(&frame.to_bytes());
}
out
}
}
fn write_string(out: &mut Vec<u8>, text: &str) -> Result<()> {
let bytes = text.as_bytes();
if bytes.len() > u16::MAX as usize {
return Err(Error::TooLarge);
}
out.extend_from_slice(&(bytes.len() as u16).to_le_bytes());
out.extend_from_slice(bytes);
Ok(())
}
fn write_opt_string(out: &mut Vec<u8>, text: Option<&str>) -> Result<()> {
match text {
None => {
out.push(0);
Ok(())
}
Some(text) => {
out.push(1);
write_string(out, text)
}
}
}
fn take_u8(buf: &[u8], cursor: &mut usize) -> Result<u8> {
if *cursor >= buf.len() {
return Err(Error::StreamDelta("truncated delta payload"));
}
let value = buf[*cursor];
*cursor += 1;
Ok(value)
}
fn take_u32(buf: &[u8], cursor: &mut usize) -> Result<u32> {
if *cursor + 4 > buf.len() {
return Err(Error::StreamDelta("truncated delta payload"));
}
let value = read_u32(buf, *cursor)?;
*cursor += 4;
Ok(value)
}
fn take_u64(buf: &[u8], cursor: &mut usize) -> Result<u64> {
if *cursor + 8 > buf.len() {
return Err(Error::StreamDelta("truncated delta payload"));
}
let value = read_u64(buf, *cursor)?;
*cursor += 8;
Ok(value)
}
fn take_i64(buf: &[u8], cursor: &mut usize) -> Result<i64> {
take_u64(buf, cursor).map(|value| value as i64)
}
fn take_hash(buf: &[u8], cursor: &mut usize) -> Result<[u8; 32]> {
if *cursor + 32 > buf.len() {
return Err(Error::StreamDelta("truncated delta payload"));
}
let mut hash = [0u8; 32];
hash.copy_from_slice(&buf[*cursor..*cursor + 32]);
*cursor += 32;
Ok(hash)
}
fn take_string(buf: &[u8], cursor: &mut usize) -> Result<String> {
if *cursor + 2 > buf.len() {
return Err(Error::StreamDelta("truncated delta payload"));
}
let len = u16::from_le_bytes([buf[*cursor], buf[*cursor + 1]]) as usize;
*cursor += 2;
if *cursor + len > buf.len() {
return Err(Error::StreamDelta("truncated delta payload"));
}
let text = std::str::from_utf8(&buf[*cursor..*cursor + len])
.map_err(|_| Error::Utf8 {
offset: *cursor as u64,
})?
.to_string();
*cursor += len;
Ok(text)
}
fn take_opt_string(buf: &[u8], cursor: &mut usize) -> Result<Option<String>> {
match take_u8(buf, cursor)? {
0 => Ok(None),
1 => Ok(Some(take_string(buf, cursor)?)),
_ => Err(Error::StreamDelta("invalid optional string tag")),
}
}
fn decode_fragment(
buf: &[u8],
cursor: &mut usize,
limits: crate::layout::Limits,
) -> Result<InfinityModule> {
let len = take_u32(buf, cursor)? as usize;
if *cursor + len > buf.len() {
return Err(Error::StreamDelta("truncated fragment"));
}
let bytes = &buf[*cursor..*cursor + len];
*cursor += len;
crate::decode::decode_with(
bytes,
crate::layout::DecodeOpts {
limits,
verify_hash: true,
verify_canonical: false,
},
)
}
pub fn delta_hash(module: &InfinityModule, ops: &[DeltaOp]) -> Result<[u8; 32]> {
let updated = apply_delta(module, ops)?;
let image = encode_with(&updated, EncodeOpts::default())?;
hash_image(&image)
}
pub fn decode_stream(bytes: &[u8]) -> Result<StreamedModule> {
let mut decoder = StreamDecoder::new();
decoder.push_bytes(bytes)?;
decoder.finish()
}
pub fn aligned_section_len(len: usize) -> usize {
align8(len as u64) as usize
}
pub fn delta_annotations(ops: &[DeltaOp]) -> Vec<Annotation> {
ops.iter()
.filter_map(|op| match op {
DeltaOp::SetAnnotation { node, kind } => Some(Annotation {
node: *node,
kind: kind.clone(),
}),
_ => None,
})
.collect()
}