nordisk/heimdall
guests/verify-bundle/tests/common/zwriter.rs
// SPDX-License-Identifier: AGPL-3.0-only
//! A test-only znippy v0.7 archive writer that can also write what znippy never
//! would: duplicate rows, hostile names, lying sizes, wrong checksums, a
//! compressed flag over stored bytes, odd sections. `writer_output_reads_back_
//! through_znippys_own_reader` (in `archive.rs`) proves its honest output is a
//! real znippy archive by opening it with znippy-common itself.
use std::collections::BTreeMap;
use std::sync::Arc;
use arrow_array::{
ArrayRef, BooleanArray, Int8Array, RecordBatch, StringArray, UInt32Array, UInt64Array,
};
use arrow_ipc::writer::StreamWriter;
use arrow_schema::{DataType, Field, Schema};
pub const MAGIC: &[u8; 8] = b"ZNPYMIDX";
/// One OpenZL frame of `bytes`, written exactly as znippy writes a blob.
pub fn openzl_frame(bytes: &[u8]) -> Vec<u8> {
znippy_common::codec::CompressCtx::new(3)
.unwrap()
.compress(bytes)
.unwrap()
}
#[derive(Debug, Clone)]
pub struct Row {
pub path: String,
pub chunk_seq: u32,
pub fdata_offset: u64,
pub compressed: bool,
pub uncompressed_size: u64,
pub blob_offset: u64,
pub blob_size: u64,
pub checksum: Vec<u8>,
}
#[derive(Debug, Clone)]
pub struct RawArchive {
pub blobs: Vec<u8>,
pub rows: Vec<Row>,
/// `Some(v)` stamps `znippy_format_version` into the data sub-index schema.
pub version: Option<String>,
/// Extra manifest sections: (module_name, raw bytes) written after the data index.
pub extra_sections: Vec<(String, Vec<u8>)>,
/// Width of the checksum column (32 is the real one).
pub checksum_width: i32,
/// Leave the checksum column out entirely.
pub omit_checksum: bool,
}
impl RawArchive {
/// Every member stored (uncompressed), split into `chunk`-byte chunks.
pub fn from_members(members: &BTreeMap<String, Vec<u8>>, chunk: usize) -> Self {
let mut blobs = Vec::new();
let mut rows = Vec::new();
for (path, bytes) in members {
let pieces: Vec<&[u8]> = if bytes.is_empty() {
vec![&bytes[..]]
} else {
bytes.chunks(chunk).collect()
};
let mut fdata = 0u64;
for (seq, piece) in pieces.into_iter().enumerate() {
rows.push(Row {
path: path.clone(),
chunk_seq: seq as u32,
fdata_offset: fdata,
compressed: false,
uncompressed_size: piece.len() as u64,
blob_offset: blobs.len() as u64,
blob_size: piece.len() as u64,
checksum: blake3::hash(piece).as_bytes().to_vec(),
});
fdata += piece.len() as u64;
blobs.extend_from_slice(piece);
}
}
RawArchive {
blobs,
rows,
version: Some("3".into()),
extra_sections: Vec::new(),
checksum_width: 32,
omit_checksum: false,
}
}
/// Every member through REAL OpenZL frames (znippy's own codec, its pinned
/// write format), split into `chunk`-byte chunks; empty members stay stored
/// as znippy writes them. Checksums cover the uncompressed bytes.
pub fn from_members_openzl(members: &BTreeMap<String, Vec<u8>>, chunk: usize) -> Self {
let mut a = Self::from_members(members, chunk);
let plain = std::mem::take(&mut a.blobs);
for r in &mut a.rows {
let piece = &plain[r.blob_offset as usize..(r.blob_offset + r.blob_size) as usize];
if piece.is_empty() {
r.blob_offset = a.blobs.len() as u64;
continue;
}
let frame = openzl_frame(piece);
r.compressed = true;
r.blob_offset = a.blobs.len() as u64;
r.blob_size = frame.len() as u64;
a.blobs.extend_from_slice(&frame);
}
a
}
/// Append `frame` as one more compressed chunk of `path` that declares 0
/// uncompressed bytes (the member's size and digests stay honest).
pub fn append_compressed_tail(&mut self, path: &str, frame: &[u8]) {
let last = self
.rows
.iter()
.filter(|r| r.path == path)
.max_by_key(|r| r.chunk_seq)
.unwrap()
.clone();
let off = self.blobs.len() as u64;
self.blobs.extend_from_slice(frame);
self.rows.push(Row {
path: path.into(),
chunk_seq: last.chunk_seq + 1,
fdata_offset: last.fdata_offset + last.uncompressed_size,
compressed: true,
uncompressed_size: 0,
blob_offset: off,
blob_size: frame.len() as u64,
checksum: blake3::hash(b"").as_bytes().to_vec(),
});
}
/// The blob bytes of `path`'s first chunk.
pub fn blob_mut(&mut self, path: &str) -> &mut [u8] {
let r = self.row_mut(path).clone();
&mut self.blobs[r.blob_offset as usize..(r.blob_offset + r.blob_size) as usize]
}
pub fn row_mut(&mut self, path: &str) -> &mut Row {
self.rows
.iter_mut()
.find(|r| r.path == path)
.unwrap_or_else(|| panic!("no row for {path}"))
}
pub fn encode(&self) -> Vec<u8> {
let mut out = self.blobs.clone();
let mut sections: Vec<(String, u64, u64, u64)> = Vec::new();
let data = self.index_stream();
sections.push((
"data".into(),
out.len() as u64,
data.len() as u64,
self.rows.len() as u64,
));
out.extend_from_slice(&data);
for (name, bytes) in &self.extra_sections {
sections.push((name.clone(), out.len() as u64, bytes.len() as u64, 0));
out.extend_from_slice(bytes);
}
let manifest_offset = out.len() as u64;
out.extend_from_slice(&manifest_stream(§ions));
out.extend_from_slice(MAGIC);
out.extend_from_slice(&manifest_offset.to_le_bytes());
out
}
fn index_stream(&self) -> Vec<u8> {
let mut fields = vec![
Field::new("relative_path", DataType::Utf8, false),
Field::new("chunk_seq", DataType::UInt32, false),
Field::new("fdata_offset", DataType::UInt64, false),
Field::new("compressed", DataType::Boolean, false),
Field::new("uncompressed_size", DataType::UInt64, false),
Field::new("blob_offset", DataType::UInt64, false),
Field::new("blob_size", DataType::UInt64, false),
];
if !self.omit_checksum {
fields.push(Field::new(
"checksum",
DataType::FixedSizeBinary(self.checksum_width),
false,
));
}
let mut meta = std::collections::HashMap::new();
if let Some(v) = &self.version {
meta.insert("znippy_format_version".to_string(), v.clone());
}
let schema = Arc::new(Schema::new_with_metadata(fields, meta));
let r = &self.rows;
let mut cols: Vec<ArrayRef> = vec![
Arc::new(StringArray::from_iter_values(
r.iter().map(|x| x.path.as_str()),
)),
Arc::new(UInt32Array::from_iter_values(r.iter().map(|x| x.chunk_seq))),
Arc::new(UInt64Array::from_iter_values(
r.iter().map(|x| x.fdata_offset),
)),
Arc::new(BooleanArray::from(
r.iter().map(|x| x.compressed).collect::<Vec<_>>(),
)),
Arc::new(UInt64Array::from_iter_values(
r.iter().map(|x| x.uncompressed_size),
)),
Arc::new(UInt64Array::from_iter_values(
r.iter().map(|x| x.blob_offset),
)),
Arc::new(UInt64Array::from_iter_values(r.iter().map(|x| x.blob_size))),
];
if !self.omit_checksum {
let w = self.checksum_width as usize;
let mut b = arrow_array::builder::FixedSizeBinaryBuilder::with_capacity(
r.len(),
self.checksum_width,
);
for x in r {
let mut c = x.checksum.clone();
c.resize(w, 0);
b.append_value(&c).unwrap();
}
cols.push(Arc::new(b.finish()));
}
let batch = RecordBatch::try_new(schema.clone(), cols).unwrap();
let mut buf = Vec::new();
{
let mut w = StreamWriter::try_new(&mut buf, &schema).unwrap();
w.write(&batch).unwrap();
w.finish().unwrap();
}
buf
}
}
pub fn manifest_stream(sections: &[(String, u64, u64, u64)]) -> Vec<u8> {
let schema = Arc::new(Schema::new(vec![
Field::new("pkg_type", DataType::Int8, false),
Field::new("repo", DataType::Utf8, false),
Field::new("module_name", DataType::Utf8, false),
Field::new("index_offset", DataType::UInt64, false),
Field::new("index_len", DataType::UInt64, false),
Field::new("row_count", DataType::UInt64, false),
]));
let n = sections.len();
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int8Array::from(vec![0i8; n])),
Arc::new(StringArray::from_iter_values(std::iter::repeat_n("", n))),
Arc::new(StringArray::from_iter_values(
sections.iter().map(|s| s.0.as_str()),
)),
Arc::new(UInt64Array::from_iter_values(sections.iter().map(|s| s.1))),
Arc::new(UInt64Array::from_iter_values(sections.iter().map(|s| s.2))),
Arc::new(UInt64Array::from_iter_values(sections.iter().map(|s| s.3))),
],
)
.unwrap();
let mut buf = Vec::new();
{
let mut w = StreamWriter::try_new(&mut buf, &schema).unwrap();
w.write(&batch).unwrap();
w.finish().unwrap();
}
buf
}