use prikk_error::{PrikkError, Result};
use prikk_object::{ObjectEnvelope, ObjectId, ObjectType};
use crate::byte_cursor::ByteCursor;
use crate::container::{self, ContainerRecordStatus, container_magic};
use crate::file_codec::push_u16;
use crate::frame_resync::resync_to_next_magic;
use crate::fsutil::{append_file_required, len_to_u64, read_file_if_exists};
use crate::layout::{ContainerSlot, RepositoryLayout, persisted_object_types};
use prikk_hash::sha256;
const INDEX_MAGIC: &[u8; 8] = b"PIDXENT1";
const INDEX_VERSION: u16 = 1;
const INDEX_HEADER_LEN: usize = 8 + 2 + 8 + 32;
const INDEX_BODY_LEN: usize = 32 + 2 + 1 + 8 + 8 + 32;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct IndexEntry {
pub(crate) object_id: ObjectId,
pub(crate) object_type: ObjectType,
pub(crate) slot: ContainerSlot,
pub(crate) offset: u64,
pub(crate) length: u64,
pub(crate) container_checksum: [u8; 32],
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum IndexRecordStatus {
Evaluated,
Failed { message: String },
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct IndexRecordOutcome {
pub(crate) offset: usize,
pub(crate) status: IndexRecordStatus,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct IndexReplay {
pub(crate) entries: Vec<IndexEntry>,
pub(crate) trailing_partial_bytes: usize,
pub(crate) record_outcomes: Vec<IndexRecordOutcome>,
}
impl IndexReplay {
#[must_use]
pub(crate) fn has_item_failure(&self) -> bool {
self.record_outcomes
.iter()
.any(|outcome| matches!(outcome.status, IndexRecordStatus::Failed { .. }))
}
}
fn slot_code(slot: ContainerSlot) -> u8 {
match slot {
ContainerSlot::A => 0,
ContainerSlot::B => 1,
}
}
fn slot_from_code(code: u8) -> Result<ContainerSlot> {
match code {
0 => Ok(ContainerSlot::A),
1 => Ok(ContainerSlot::B),
other => Err(PrikkError::MalformedData(format!(
"unrecognized container slot code {other}"
))),
}
}
fn encode_entry_body(entry: &IndexEntry) -> Vec<u8> {
let mut body = Vec::with_capacity(INDEX_BODY_LEN);
body.extend_from_slice(entry.object_id.as_bytes());
push_u16(&mut body, entry.object_type.code());
body.push(slot_code(entry.slot));
body.extend_from_slice(&entry.offset.to_be_bytes());
body.extend_from_slice(&entry.length.to_be_bytes());
body.extend_from_slice(&entry.container_checksum);
body
}
fn decode_entry_body(body: &[u8]) -> Result<IndexEntry> {
let mut cursor = ByteCursor::new(body);
let object_id = ObjectId::from_bytes(cursor.read_array::<32>()?);
let object_type = ObjectType::from_code(cursor.read_u16()?)?;
let slot = slot_from_code(cursor.read_array::<1>()?[0])?;
let offset = cursor.read_u64()?;
let length = cursor.read_u64()?;
let container_checksum = cursor.read_array::<32>()?;
if !cursor.is_finished() {
return Err(PrikkError::MalformedData(
"trailing bytes in index entry body".to_string(),
));
}
Ok(IndexEntry {
object_id,
object_type,
slot,
offset,
length,
container_checksum,
})
}
pub(crate) fn encode_index_record(entry: &IndexEntry) -> Result<Vec<u8>> {
let body = encode_entry_body(entry);
let body_len = len_to_u64(body.len())?;
let checksum = index_record_checksum(body_len, &body);
let mut out = Vec::with_capacity(INDEX_HEADER_LEN + body.len());
out.extend_from_slice(INDEX_MAGIC);
push_u16(&mut out, INDEX_VERSION);
out.extend_from_slice(&body_len.to_be_bytes());
out.extend_from_slice(&checksum);
out.extend_from_slice(&body);
Ok(out)
}
fn index_record_checksum(body_len: u64, body: &[u8]) -> [u8; 32] {
let mut preimage = Vec::new();
preimage.extend_from_slice(INDEX_MAGIC);
preimage.extend_from_slice(&INDEX_VERSION.to_be_bytes());
preimage.extend_from_slice(&body_len.to_be_bytes());
preimage.extend_from_slice(body);
sha256(&preimage)
}
struct IndexHeader {
body_len: u64,
checksum: [u8; 32],
}
fn parse_index_header(header: &[u8]) -> Result<IndexHeader> {
let mut cursor = ByteCursor::new(header);
let magic = cursor.read_array::<8>()?;
if &magic != INDEX_MAGIC {
return Err(PrikkError::MalformedData(
"invalid index record magic".to_string(),
));
}
let version = cursor.read_u16()?;
if version != INDEX_VERSION {
return Err(PrikkError::UnsupportedFormatVersion(u32::from(version)));
}
let body_len = cursor.read_u64()?;
let checksum = cursor.read_array::<32>()?;
if !cursor.is_finished() {
return Err(PrikkError::MalformedData(
"trailing bytes in index header".to_string(),
));
}
Ok(IndexHeader { body_len, checksum })
}
enum FrameAttempt {
Record {
entry: IndexEntry,
next_offset: usize,
},
TrailingPartial {
remaining: usize,
},
Invalid {
message: String,
},
}
fn parse_frame_at(bytes: &[u8], offset: usize) -> FrameAttempt {
let remaining = bytes.len().saturating_sub(offset);
if remaining < INDEX_HEADER_LEN {
return FrameAttempt::TrailingPartial { remaining };
}
let header_end = offset + INDEX_HEADER_LEN;
let Some(header) = bytes.get(offset..header_end) else {
return FrameAttempt::TrailingPartial { remaining };
};
let header_values = match parse_index_header(header) {
Ok(values) => values,
Err(err) => {
return FrameAttempt::Invalid {
message: err.to_string(),
};
}
};
let Ok(body_len) = usize::try_from(header_values.body_len) else {
return FrameAttempt::Invalid {
message: "index body length does not fit usize".to_string(),
};
};
let Some(body_end) = header_end.checked_add(body_len) else {
return FrameAttempt::Invalid {
message: "index body end overflow".to_string(),
};
};
let Some(body) = bytes.get(header_end..body_end) else {
return FrameAttempt::TrailingPartial { remaining };
};
let expected = index_record_checksum(header_values.body_len, body);
if expected != header_values.checksum {
return FrameAttempt::Invalid {
message: format!("index checksum mismatch at byte offset {offset}"),
};
}
match decode_entry_body(body) {
Ok(entry) => FrameAttempt::Record {
entry,
next_offset: body_end,
},
Err(err) => FrameAttempt::Invalid {
message: err.to_string(),
},
}
}
pub(crate) fn decode_index_records(bytes: &[u8], start_offset: usize) -> Result<IndexReplay> {
let mut entries = Vec::new();
let mut record_outcomes = Vec::new();
let mut offset = start_offset;
loop {
match parse_frame_at(bytes, offset) {
FrameAttempt::Record { entry, next_offset } => {
record_outcomes.push(IndexRecordOutcome {
offset,
status: IndexRecordStatus::Evaluated,
});
entries.push(entry);
offset = next_offset;
}
FrameAttempt::TrailingPartial { remaining } => {
return Ok(IndexReplay {
entries,
trailing_partial_bytes: remaining,
record_outcomes,
});
}
FrameAttempt::Invalid { message } => {
record_outcomes.push(IndexRecordOutcome {
offset,
status: IndexRecordStatus::Failed { message },
});
match resync_to_next_magic(bytes, offset + 1, INDEX_MAGIC.as_slice()) {
Some(next) => offset = next,
None => {
return Ok(IndexReplay {
entries,
trailing_partial_bytes: 0,
record_outcomes,
});
}
}
}
}
}
}
pub(crate) fn replay_index(layout: &RepositoryLayout) -> Result<IndexReplay> {
Ok(replay_index_with_extent(layout)?.0)
}
pub(crate) fn replay_index_with_extent(layout: &RepositoryLayout) -> Result<(IndexReplay, u64)> {
let relative = layout.repository_relative(&layout.container_index_path())?;
let Some(bytes) = read_file_if_exists(layout.repository_mutation_root(), &relative)? else {
return Ok((
IndexReplay {
entries: Vec::new(),
trailing_partial_bytes: 0,
record_outcomes: Vec::new(),
},
0,
));
};
#[cfg(test)]
record_replay_index_decode_for_test();
let replay = decode_index_records(&bytes, 0)?;
let extent = len_to_u64(bytes.len().saturating_sub(replay.trailing_partial_bytes))?;
Ok((replay, extent))
}
pub(crate) fn replay_index_tail_with_extent(
layout: &RepositoryLayout,
start_offset: u64,
) -> Result<(IndexReplay, u64)> {
let relative = layout.repository_relative(&layout.container_index_path())?;
let Some(bytes) = read_file_if_exists(layout.repository_mutation_root(), &relative)? else {
return Ok((
IndexReplay {
entries: Vec::new(),
trailing_partial_bytes: 0,
record_outcomes: Vec::new(),
},
0,
));
};
#[cfg(test)]
record_replay_index_decode_for_test();
let start = usize::try_from(start_offset)
.map_err(|_| PrikkError::Integrity("index start offset exceeds usize".to_string()))?;
let replay = decode_index_records(&bytes, start)?;
let extent = len_to_u64(bytes.len().saturating_sub(replay.trailing_partial_bytes))?;
Ok((replay, extent))
}
#[cfg(test)]
std::thread_local! {
static REPLAY_INDEX_DECODE_COUNT: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}
#[cfg(test)]
fn record_replay_index_decode_for_test() {
REPLAY_INDEX_DECODE_COUNT.with(|count| count.set(count.get() + 1));
}
#[cfg(test)]
pub(crate) fn reset_replay_index_decode_count_for_test() {
REPLAY_INDEX_DECODE_COUNT.with(|count| count.set(0));
}
#[cfg(test)]
pub(crate) fn replay_index_decode_count_for_test() -> usize {
REPLAY_INDEX_DECODE_COUNT.with(|count| count.get())
}
pub(crate) fn lookup_object_location(
layout: &RepositoryLayout,
object_id: ObjectId,
) -> Result<Option<IndexEntry>> {
let replay = replay_index(layout)?;
if replay.has_item_failure() {
return Err(PrikkError::Integrity(
"object index has a damaged entry; run doctor before reading".to_string(),
));
}
Ok(replay
.entries
.into_iter()
.rev()
.find(|entry| entry.object_id == object_id))
}
pub(crate) fn read_object_envelope_at(
layout: &RepositoryLayout,
entry: &IndexEntry,
) -> Result<ObjectEnvelope> {
let container_relative =
layout.repository_relative(&layout.container_slot_path(entry.object_type, entry.slot))?;
let Some(bytes) = read_file_if_exists(layout.repository_mutation_root(), &container_relative)?
else {
return Err(PrikkError::Integrity(format!(
"index names container {:?} slot {:?}, which does not exist",
entry.object_type, entry.slot
)));
};
let offset = usize::try_from(entry.offset)
.map_err(|_| PrikkError::Integrity("index entry offset exceeds usize".to_string()))?;
let record = container::decode_container_record_at(entry.object_type, &bytes, offset)?
.ok_or_else(|| {
PrikkError::Integrity(format!(
"index entry for {} names an offset past its container's end",
entry.object_id
))
})?;
Ok(record.envelope)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum WriteDecision {
New,
AlreadyPresent(ObjectId),
}
pub(crate) fn decide_write_outcome(
layout: &RepositoryLayout,
object_type: ObjectType,
envelope: &ObjectEnvelope,
existing: Option<&IndexEntry>,
) -> Result<WriteDecision> {
let object_id = envelope.object_id();
let Some(existing) = existing else {
return Ok(WriteDecision::New);
};
if existing.object_type != object_type {
return Err(PrikkError::Integrity(format!(
"existing index entry for {object_id} has type {}, expected {object_type}",
existing.object_type
)));
}
let existing_envelope = read_object_envelope_at(layout, existing)?;
if existing_envelope != *envelope {
return Err(PrikkError::Integrity(format!(
"existing container record for {object_id} differs from candidate"
)));
}
Ok(WriteDecision::AlreadyPresent(object_id))
}
pub(crate) fn append_object_to_container(
layout: &RepositoryLayout,
object_type: ObjectType,
envelope: &ObjectEnvelope,
) -> Result<IndexEntry> {
let object_id = envelope.object_id();
let record_bytes = container::encode_container_record(object_type, envelope)?;
let container_relative =
layout.repository_relative(&layout.container_slot_path(object_type, ContainerSlot::A))?;
let existing_len = read_file_if_exists(layout.repository_mutation_root(), &container_relative)?
.map_or(0, |bytes| bytes.len());
let offset = len_to_u64(existing_len)?;
let length = len_to_u64(record_bytes.len())?;
let container_checksum = frame_checksum(object_type, &record_bytes)?;
append_file_required(
layout.repository_mutation_root(),
&container_relative,
&record_bytes,
)?;
let entry = IndexEntry {
object_id,
object_type,
slot: ContainerSlot::A,
offset,
length,
container_checksum,
};
let index_bytes = encode_index_record(&entry)?;
let index_relative = layout.repository_relative(&layout.container_index_path())?;
append_file_required(
layout.repository_mutation_root(),
&index_relative,
&index_bytes,
)?;
Ok(entry)
}
fn frame_checksum(object_type: ObjectType, record_bytes: &[u8]) -> Result<[u8; 32]> {
let magic = container_magic(object_type)?;
let checksum_start = magic.len() + 2 + 8;
let checksum_end = checksum_start + 32;
let checksum_bytes = record_bytes
.get(checksum_start..checksum_end)
.ok_or_else(|| {
PrikkError::Integrity("just-encoded container record is too short".to_string())
})?;
let mut checksum = [0_u8; 32];
checksum.copy_from_slice(checksum_bytes);
Ok(checksum)
}
#[allow(dead_code)]
pub(crate) fn rebuild_index_from_containers(layout: &RepositoryLayout) -> Result<Vec<IndexEntry>> {
let mut entries = Vec::new();
for object_type in persisted_object_types() {
let relative = layout
.repository_relative(&layout.container_slot_path(object_type, ContainerSlot::A))?;
let Some(bytes) = read_file_if_exists(layout.repository_mutation_root(), &relative)? else {
continue;
};
let replay = container::decode_container_records(object_type, &bytes)?;
let mut records = replay.records.iter();
for outcome in &replay.record_outcomes {
let ContainerRecordStatus::Evaluated {
frame_len,
checksum,
} = &outcome.status
else {
continue;
};
let Some(record) = records.next() else {
return Err(PrikkError::Integrity(
"container replay outcome/record count mismatch".to_string(),
));
};
entries.push(IndexEntry {
object_id: record.envelope.object_id(),
object_type,
slot: ContainerSlot::A,
offset: len_to_u64(outcome.offset)?,
length: len_to_u64(*frame_len)?,
container_checksum: *checksum,
});
}
}
Ok(entries)
}
#[cfg(test)]
pub(crate) fn remove_index_entry_for_test(
layout: &RepositoryLayout,
object_id: ObjectId,
) -> Result<()> {
let path = layout.container_index_path();
let bytes = std::fs::read(&path)?;
let replay = decode_index_records(&bytes, 0)?;
let mut entries = replay.entries.iter();
for outcome in &replay.record_outcomes {
let IndexRecordStatus::Evaluated = &outcome.status else {
continue;
};
let Some(entry) = entries.next() else {
return Err(PrikkError::Integrity(
"index replay outcome/entry count mismatch".to_string(),
));
};
if entry.object_id != object_id {
continue;
}
let frame_len = INDEX_HEADER_LEN + INDEX_BODY_LEN;
let end = outcome.offset + frame_len;
let mut retained = bytes.get(..outcome.offset).unwrap_or_default().to_vec();
retained.extend_from_slice(bytes.get(end..).unwrap_or_default());
std::fs::write(&path, retained)?;
return Ok(());
}
Err(PrikkError::Integrity(format!(
"no index entry for {object_id} to remove"
)))
}
#[cfg(test)]
mod tests;