use std::collections::{BTreeMap, BTreeSet};
use prikk_error::{PrikkError, Result};
use prikk_object::{
BlockPayload, ObjectEnvelope, ObjectId, ObjectType, RefStatePayload, Signature, SignerRole,
};
use crate::author_key_index::{
AuthorKeyEntry, check_author_key_conflict, lookup_author_key_entries,
record_author_key_material,
};
use crate::byte_cursor::ByteCursor;
use crate::file_codec::{decode_envelope_file, encode_envelope_file, push_bytes_u64, push_u64};
use crate::fsutil::len_to_u64;
use crate::layout::{LockableContainer, RepositoryLayout};
use crate::lock::{ActiveLock, acquire_container_locks};
use crate::object_store::{ObjectReadSnapshot, ObjectReader, ObjectWriteSession, ObjectWriter};
use crate::patch_replay::decode::{
DecodedDeletePreimage, DecodedOperationKind, decode_patch_operations,
};
use crate::refs::{RefStore, ensure_ref_target_valid};
const BUNDLE_MAGIC: &[u8; 8] = b"PBNDL002";
const RETIRED_BUNDLE_MAGIC_V1: &[u8; 8] = b"PBNDL001";
pub const DEFAULT_BUNDLE_MAX_OBJECT_COUNT: usize = 100_000;
pub const DEFAULT_BUNDLE_MAX_TOTAL_BYTES: usize = 256 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BundleImportOptions {
pub max_object_count: usize,
pub max_total_bytes: usize,
}
impl BundleImportOptions {
#[must_use]
pub const fn default_limits() -> Self {
Self {
max_object_count: DEFAULT_BUNDLE_MAX_OBJECT_COUNT,
max_total_bytes: DEFAULT_BUNDLE_MAX_TOTAL_BYTES,
}
}
#[must_use]
pub const fn with_max_object_count(mut self, max_object_count: usize) -> Self {
self.max_object_count = max_object_count;
self
}
#[must_use]
pub const fn with_max_total_bytes(mut self, max_total_bytes: usize) -> Self {
self.max_total_bytes = max_total_bytes;
self
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BundleExportReport {
pub ref_name: String,
pub tip_block_id: ObjectId,
pub object_count: usize,
pub author_key_count: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BundleImportReport {
pub ref_name: String,
pub ref_state_id: ObjectId,
pub object_count: usize,
pub written_object_count: usize,
pub recorded_author_key_count: usize,
}
pub fn export_bundle(
layout: &RepositoryLayout,
ref_name: &str,
) -> Result<(BundleExportReport, Vec<u8>)> {
let ref_store = RefStore::new(layout.clone());
let object_store = ObjectReadSnapshot::open(layout)?;
let Some(ref_state_id) = ref_store.read_current_ref_state_id(ref_name)? else {
return Err(PrikkError::Integrity(format!(
"ref {ref_name} does not exist, nothing to export"
)));
};
let ref_state_envelope = object_store
.read_typed(ref_state_id, ObjectType::RefState)?
.ok_or_else(|| PrikkError::Integrity(format!("missing RefState object: {ref_state_id}")))?;
let ref_state_payload = RefStatePayload::decode_canonical(
&ref_state_envelope.canonical_payload,
ref_state_envelope.schema_version,
)?;
let mut tag_envelopes: Vec<ObjectEnvelope> = Vec::new();
let tip_block_id =
resolve_ref_target_block(&object_store, &ref_state_payload, &mut tag_envelopes)?;
let mut ref_state_chain: Vec<ObjectEnvelope> = vec![ref_state_envelope];
let mut required_attestation_ids: BTreeSet<ObjectId> = ref_state_payload
.required_attestation_ids
.iter()
.copied()
.collect();
let mut ancestors = crate::merge_evidence::ancestors_inclusive(&object_store, tip_block_id)?;
let mut previous = ref_state_payload.previous_ref_state_id;
let mut seen_ref_states: BTreeSet<ObjectId> = BTreeSet::from([ref_state_id]);
while let Some(previous_id) = previous {
if !seen_ref_states.insert(previous_id) {
return Err(PrikkError::Integrity(format!(
"RefState chain for {ref_name} contains a cycle at {previous_id}"
)));
}
let envelope = read_required(&object_store, previous_id, ObjectType::RefState)?;
let payload = RefStatePayload::decode_canonical(
&envelope.canonical_payload,
envelope.schema_version,
)?;
required_attestation_ids.extend(payload.required_attestation_ids.iter().copied());
let target_block_id =
resolve_ref_target_block(&object_store, &payload, &mut tag_envelopes)?;
ancestors.extend(crate::merge_evidence::ancestors_inclusive(
&object_store,
target_block_id,
)?);
previous = payload.previous_ref_state_id;
ref_state_chain.push(envelope);
}
let mut patch_ids: BTreeSet<ObjectId> = BTreeSet::new();
let mut blob_ids: BTreeSet<ObjectId> = BTreeSet::new();
for payload in ancestors.values() {
patch_ids.extend(payload.patch_ids.iter().copied());
if let Some(blob_id) = payload.snapshot_blob_ref {
blob_ids.insert(blob_id);
}
}
let mut objects: Vec<ObjectEnvelope> = ref_state_chain;
objects.append(&mut tag_envelopes);
for block_id in ancestors.keys() {
objects.push(read_required(&object_store, *block_id, ObjectType::Block)?);
}
let mut patch_envelopes: Vec<ObjectEnvelope> = Vec::with_capacity(patch_ids.len());
for patch_id in &patch_ids {
let envelope = read_required(&object_store, *patch_id, ObjectType::Patch)?;
for operation in
crate::patch_replay::decode::decode_patch_operations(&envelope.canonical_payload)?
{
match operation.kind {
crate::patch_replay::decode::DecodedOperationKind::CreateFile {
blob_id, ..
} => {
blob_ids.insert(blob_id);
}
crate::patch_replay::decode::DecodedOperationKind::ReplaceBinary {
old_blob_id,
new_blob_id,
..
} => {
blob_ids.insert(old_blob_id);
blob_ids.insert(new_blob_id);
}
crate::patch_replay::decode::DecodedOperationKind::DeleteNode {
preimage:
crate::patch_replay::decode::DecodedDeletePreimage::File { old_blob_id, .. },
..
} => {
blob_ids.insert(old_blob_id);
}
_ => {}
}
}
patch_envelopes.push(envelope);
}
objects.extend(patch_envelopes);
for blob_id in &blob_ids {
objects.push(read_required(&object_store, *blob_id, ObjectType::Blob)?);
}
for attestation_id in &required_attestation_ids {
objects.push(read_required(
&object_store,
*attestation_id,
ObjectType::Attestation,
)?);
}
let mut author_key_ids: BTreeSet<String> = BTreeSet::new();
for envelope in &objects {
if envelope.object_type != ObjectType::Patch {
continue;
}
if let Some(signature) = envelope
.signatures
.iter()
.find(|signature| signature.signer_role == SignerRole::Author)
{
author_key_ids.insert(signature.key_id.clone());
}
}
let mut author_keys: Vec<AuthorKeyEntry> = Vec::with_capacity(author_key_ids.len());
for key_id in &author_key_ids {
let entries = lookup_author_key_entries(layout, key_id)?;
let mut distinct = entries.iter().map(|entry| entry.public_key);
let Some(first) = distinct.next() else {
continue;
};
if distinct.any(|public_key| public_key != first) {
return Err(PrikkError::Integrity(format!(
"author key_id {key_id} has more than one distinct recorded public key locally; \
refusing to export a provenance claim this repository's own material does not \
agree on -- run doctor, though no repair exists for this container"
)));
}
author_keys.push(AuthorKeyEntry {
key_id: key_id.clone(),
public_key: first,
});
}
let object_count = objects.len();
let author_key_count = author_keys.len();
let bytes = encode_bundle(ref_name, &objects, &author_keys)?;
Ok((
BundleExportReport {
ref_name: ref_name.to_string(),
tip_block_id,
object_count,
author_key_count,
},
bytes,
))
}
pub fn import_bundle(
layout: &RepositoryLayout,
bytes: &[u8],
options: &BundleImportOptions,
) -> Result<BundleImportReport> {
if bytes.len() > options.max_total_bytes {
return Err(PrikkError::MalformedData(format!(
"bundle is {} bytes, over the configured limit of {} bytes",
bytes.len(),
options.max_total_bytes
)));
}
let (origin_ref_name, objects, author_keys) = decode_bundle(bytes, options.max_object_count)?;
let Some(ref_state_envelope) = objects.first() else {
return Err(PrikkError::MalformedData(
"bundle contains no objects".to_string(),
));
};
if ref_state_envelope.object_type != ObjectType::RefState {
return Err(PrikkError::MalformedData(
"bundle's first object must be the exported ref's RefState".to_string(),
));
}
let ref_state_id = ref_state_envelope.object_id();
let mut bundle_key_ids: BTreeMap<&str, [u8; 32]> = BTreeMap::new();
for entry in &author_keys {
match bundle_key_ids.get(entry.key_id.as_str()) {
Some(existing) if *existing != entry.public_key => {
return Err(PrikkError::MalformedData(format!(
"bundle's author-key section carries two different public keys for key_id {} \
-- refusing the whole import",
entry.key_id
)));
}
Some(_) => {}
None => {
bundle_key_ids.insert(&entry.key_id, entry.public_key);
}
}
}
let read_snapshot = ObjectReadSnapshot::open(layout)?;
let bundle_objects_by_id: BTreeMap<ObjectId, ObjectEnvelope> = objects
.iter()
.map(|envelope| (envelope.object_id(), envelope.clone()))
.collect();
let ref_state_payload = RefStatePayload::decode_canonical(
&ref_state_envelope.canonical_payload,
ref_state_envelope.schema_version,
)?;
let combined_reader = BundleAndLocalReader {
bundle_objects: &bundle_objects_by_id,
local: &read_snapshot,
};
ensure_ref_target_valid(
&combined_reader,
ref_state_payload.kind,
ref_state_payload.target_object_id,
ref_state_id,
)?;
for envelope in &objects {
if envelope.object_type != ObjectType::Patch {
continue;
}
for operation in decode_patch_operations(&envelope.canonical_payload)? {
for blob_id in bundle_referenced_blob_ids(&operation.kind) {
if !bundle_objects_by_id.contains_key(&blob_id)
&& !read_snapshot.contains_object(ObjectType::Blob, blob_id)
{
return Err(PrikkError::Integrity(format!(
"patch {} references blob {blob_id}, which is neither carried by this \
bundle nor already present in this repository -- refusing the whole \
import, no partial write",
envelope.object_id()
)));
}
}
}
}
for envelope in &objects {
if envelope.object_type != ObjectType::Block {
continue;
}
let block_id = envelope.object_id();
let block_payload = BlockPayload::decode_canonical(&envelope.canonical_payload)?;
for patch_id in &block_payload.patch_ids {
if !bundle_objects_by_id.contains_key(patch_id)
&& !read_snapshot.contains_object(ObjectType::Patch, *patch_id)
{
return Err(PrikkError::Integrity(format!(
"block {block_id} names patch {patch_id}, which is neither carried by this \
bundle nor already present in this repository -- refusing the whole import, \
no partial write"
)));
}
}
for parent_block_id in &block_payload.parent_block_ids {
if !bundle_objects_by_id.contains_key(parent_block_id)
&& !read_snapshot.contains_object(ObjectType::Block, *parent_block_id)
{
return Err(PrikkError::Integrity(format!(
"block {block_id} names parent {parent_block_id}, which is neither carried \
by this bundle nor already present in this repository -- refusing the whole \
import, no partial write"
)));
}
}
if let Some(snapshot_blob_id) = block_payload.snapshot_blob_ref {
if !bundle_objects_by_id.contains_key(&snapshot_blob_id)
&& !read_snapshot.contains_object(ObjectType::Blob, snapshot_blob_id)
{
return Err(PrikkError::Integrity(format!(
"block {block_id} names snapshot blob {snapshot_blob_id}, which is neither \
carried by this bundle nor already present in this repository -- refusing \
the whole import, no partial write"
)));
}
}
}
let mut object_store = ObjectWriteSession::open(layout)?;
let mut written_object_count = 0_usize;
for envelope in &objects {
let id = envelope.object_id();
if !object_store.contains_object(envelope.object_type, id)? {
written_object_count = written_object_count.checked_add(1).ok_or_else(|| {
PrikkError::Integrity("bundle import written-object count overflow".to_string())
})?;
}
object_store.write_object(envelope)?;
}
let mut recorded_author_key_count = 0_usize;
{
let active_lock = ActiveLock::acquire(layout)?;
for (&key_id, &public_key) in &bundle_key_ids {
check_author_key_conflict(layout, key_id, public_key)?;
}
for entry in &author_keys {
record_author_key_material(layout, &entry.key_id, entry.public_key, &active_lock)?;
recorded_author_key_count =
recorded_author_key_count.checked_add(1).ok_or_else(|| {
PrikkError::Integrity(
"bundle import recorded-author-key count overflow".to_string(),
)
})?;
}
}
let received_ref_name = format!("remotes/{origin_ref_name}");
let _received_index_lock =
acquire_container_locks(layout, &[LockableContainer::ReceivedIndex])?;
crate::received::write_received_pointer(layout, &received_ref_name, ref_state_id)?;
Ok(BundleImportReport {
ref_name: received_ref_name,
ref_state_id,
object_count: objects.len(),
written_object_count,
recorded_author_key_count,
})
}
fn read_required(
object_store: &impl ObjectReader,
id: ObjectId,
object_type: ObjectType,
) -> Result<ObjectEnvelope> {
object_store
.read_typed(id, object_type)?
.ok_or_else(|| PrikkError::Integrity(format!("missing {object_type} object: {id}")))
}
struct BundleAndLocalReader<'a> {
bundle_objects: &'a BTreeMap<ObjectId, ObjectEnvelope>,
local: &'a ObjectReadSnapshot,
}
impl ObjectReader for BundleAndLocalReader<'_> {
fn read_object(&self, id: ObjectId) -> Result<Option<ObjectEnvelope>> {
if let Some(envelope) = self.bundle_objects.get(&id) {
return Ok(Some(envelope.clone()));
}
self.local.read_object(id)
}
}
fn bundle_referenced_blob_ids(kind: &DecodedOperationKind) -> Vec<ObjectId> {
match kind {
DecodedOperationKind::CreateFile { blob_id, .. } => vec![*blob_id],
DecodedOperationKind::ReplaceBinary {
old_blob_id,
new_blob_id,
..
} => vec![*old_blob_id, *new_blob_id],
DecodedOperationKind::DeleteNode {
preimage: DecodedDeletePreimage::File { old_blob_id, .. },
..
} => vec![*old_blob_id],
_ => Vec::new(),
}
}
fn resolve_ref_target_block(
object_store: &impl ObjectReader,
ref_state_payload: &RefStatePayload,
tag_envelopes: &mut Vec<ObjectEnvelope>,
) -> Result<ObjectId> {
let (target_block_id, tag_envelope) =
crate::refs::resolve_ref_tip_block(object_store, ref_state_payload)?;
if let Some(tag_envelope) = tag_envelope {
tag_envelopes.push(tag_envelope);
}
Ok(target_block_id)
}
fn encode_bundle(
ref_name: &str,
objects: &[ObjectEnvelope],
author_keys: &[AuthorKeyEntry],
) -> Result<Vec<u8>> {
let mut out = Vec::new();
out.extend_from_slice(BUNDLE_MAGIC);
push_bytes_u64(&mut out, ref_name.as_bytes())?;
push_u64(&mut out, len_to_u64(objects.len())?);
for envelope in objects {
push_bytes_u64(&mut out, &encode_envelope_file(envelope)?)?;
}
push_u64(&mut out, len_to_u64(author_keys.len())?);
for entry in author_keys {
push_bytes_u64(&mut out, entry.key_id.as_bytes())?;
out.extend_from_slice(&entry.public_key);
}
Ok(out)
}
#[cfg(all(test, target_os = "linux"))]
fn encode_bundle_v1_for_test(ref_name: &str, objects: &[ObjectEnvelope]) -> Result<Vec<u8>> {
let mut out = Vec::new();
out.extend_from_slice(RETIRED_BUNDLE_MAGIC_V1);
push_bytes_u64(&mut out, ref_name.as_bytes())?;
push_u64(&mut out, len_to_u64(objects.len())?);
for envelope in objects {
push_bytes_u64(&mut out, &encode_envelope_file(envelope)?)?;
}
Ok(out)
}
fn decode_bundle(
bytes: &[u8],
max_object_count: usize,
) -> Result<(String, Vec<ObjectEnvelope>, Vec<AuthorKeyEntry>)> {
let mut cursor = ByteCursor::new(bytes);
let magic = cursor.read_array::<8>()?;
let has_author_key_section = if &magic == BUNDLE_MAGIC {
true
} else if &magic == RETIRED_BUNDLE_MAGIC_V1 {
false
} else {
return Err(PrikkError::MalformedData(
"invalid bundle magic".to_string(),
));
};
let ref_name_bytes = cursor.read_bytes_u64()?;
let ref_name = String::from_utf8(ref_name_bytes).map_err(|err| {
PrikkError::MalformedData(format!("invalid bundle ref name utf-8: {err}"))
})?;
let count = cursor.read_u64()?;
if count > len_to_u64(max_object_count)? {
return Err(PrikkError::MalformedData(format!(
"bundle declares {count} objects, over the configured limit of {max_object_count}"
)));
}
let mut objects = Vec::new();
for _ in 0..count {
let encoded = cursor.read_bytes_u64()?;
objects.push(decode_envelope_file(&encoded)?);
}
let mut author_keys = Vec::new();
if has_author_key_section {
let author_key_count = cursor.read_u64()?;
if author_key_count > len_to_u64(max_object_count)? {
return Err(PrikkError::MalformedData(format!(
"bundle declares {author_key_count} author key entries, over the configured limit \
of {max_object_count}"
)));
}
for _ in 0..author_key_count {
let key_id_bytes = cursor.read_bytes_u64()?;
let key_id = String::from_utf8(key_id_bytes).map_err(|err| {
PrikkError::MalformedData(format!("invalid bundle author key_id utf-8: {err}"))
})?;
Signature::validate_key_id(&key_id)?;
let public_key = cursor.read_array::<32>()?;
author_keys.push(AuthorKeyEntry { key_id, public_key });
}
}
if !cursor.is_finished() {
return Err(PrikkError::MalformedData(
"trailing bytes in bundle".to_string(),
));
}
Ok((ref_name, objects, author_keys))
}
#[cfg(all(test, target_os = "linux"))]
mod tests;