use std::io;
use std::io::{Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use blake3::Hasher as Blake3;
use serde::{Deserialize, Serialize};
use crate::error::{FileError, Result};
use crate::record::{crc32, read_u32_le, Reader, Writer, HEADER_LEN, REC_HDR};
const TAIL_MAGIC: &[u8; 8] = b"MFFTAIL1";
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct SegmentType(pub String);
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SegmentEntry {
pub kind: SegmentType,
pub offset: u64,
pub length: u64,
pub crc32: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Directory {
pub entries: Vec<SegmentEntry>,
pub flags: u32,
pub file_hash: [u8; 32],
}
pub struct DocumentWriter {
w: Writer,
segments: Vec<SegmentEntry>,
path: PathBuf,
}
impl DocumentWriter {
#[cfg_attr(feature = "dev-tracing", tracing::instrument(skip(path), fields(
crate_name = "file",
file_path = %path.as_ref().display()
)))]
pub fn begin<P: AsRef<Path>>(path: P) -> Result<Self> {
let p = path.as_ref().to_path_buf();
Ok(Self { w: Writer::create(&p, 0)?, segments: Vec::new(), path: p })
}
pub fn add_segment(
&mut self,
kind: SegmentType,
payload: &[u8],
) -> Result<()> {
let off = self.w.len();
let _ = self.w.append(payload)?;
let crc = crc32(payload);
self.segments.push(SegmentEntry {
kind,
offset: off,
length: (REC_HDR as u64) + payload.len() as u64,
crc32: crc,
});
Ok(())
}
#[cfg_attr(feature = "dev-tracing", tracing::instrument(skip(self), fields(
crate_name = "file",
segment_count = self.segments.len(),
file_path = %self.path.display()
)))]
pub fn finalize(mut self) -> Result<()> {
self.w.flush()?;
let mut hasher = Blake3::new();
let r = Reader::open(&self.path)?;
for bytes in r.iter() {
hasher.update(bytes);
}
let hash = *hasher.finalize().as_bytes();
let dir =
Directory { entries: self.segments, flags: 0, file_hash: hash };
let bytes =
bincode::serde::encode_to_vec(&dir, bincode::config::standard())
.map_err(io::Error::other)
.map_err(FileError::Io)?;
let dir_off = self.w.append(&bytes)?;
self.w.flush()?;
{
let file = &mut self.w.file;
file.seek(SeekFrom::Start(self.w.logical_end))?;
file.write_all(TAIL_MAGIC)?;
file.write_all(&dir_off.to_le_bytes())?;
file.sync_data()?;
}
Ok(())
}
}
pub struct DocumentReader {
r: Reader,
dir: Directory,
}
impl DocumentReader {
#[cfg_attr(feature = "dev-tracing", tracing::instrument(skip(path), fields(
crate_name = "file",
file_path = %path.as_ref().display()
)))]
pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
let r = Reader::open(path)?;
let mut last_off = HEADER_LEN as u64;
let phys_len = r.mmap.len();
if phys_len >= 16 {
let tail = &r.mmap[phys_len - 16..phys_len];
if &tail[..8] == TAIL_MAGIC {
let mut off_bytes = [0u8; 8];
off_bytes.copy_from_slice(&tail[8..16]);
let off = u64::from_le_bytes(off_bytes);
if (off as usize) + REC_HDR <= r.logical_end as usize {
let len =
read_u32_le(&r.mmap[off as usize..off as usize + 4])
as usize;
let s = off as usize + REC_HDR;
let e = s + len;
if e <= r.logical_end as usize {
let stored_crc = read_u32_le(
&r.mmap[off as usize + 4..off as usize + 8],
);
if crc32(&r.mmap[s..e]) == stored_crc {
last_off = off;
}
}
}
}
}
if last_off == (HEADER_LEN as u64) {
let mut p = HEADER_LEN;
let end = r.logical_end as usize;
let mut fallback_last = HEADER_LEN as u64;
while p + REC_HDR <= end {
let len = read_u32_le(&r.mmap[p..p + 4]) as usize;
if len == 0 {
break;
}
let s = p + REC_HDR;
let e = s + len;
if e > end {
break;
}
let stored_crc = read_u32_le(&r.mmap[p + 4..p + 8]);
if crc32(&r.mmap[s..e]) != stored_crc {
break;
}
fallback_last = p as u64;
p = e;
}
last_off = fallback_last;
}
let dir_bytes = r.get_at(last_off)?;
let (dir, _) = bincode::serde::decode_from_slice::<Directory, _>(
dir_bytes,
bincode::config::standard(),
)
.map_err(io::Error::other)
.map_err(FileError::Io)?;
let mut hasher = Blake3::new();
let mut q = HEADER_LEN;
let end2 = last_off as usize;
while q + REC_HDR <= end2 {
let len = read_u32_le(&r.mmap[q..q + 4]) as usize;
if len == 0 {
break;
}
let s = q + REC_HDR;
let e = s + len;
if e > end2 {
break;
}
let stored_crc = read_u32_le(&r.mmap[q + 4..q + 8]);
if crc32(&r.mmap[s..e]) != stored_crc {
break;
}
hasher.update(&r.mmap[s..e]);
q = e;
}
let calc = *hasher.finalize().as_bytes();
if calc != dir.file_hash {
return Err(FileError::BadHeader);
}
Ok(Self { r, dir })
}
#[cfg_attr(feature = "dev-tracing", tracing::instrument(skip(self, callback), fields(
crate_name = "file",
segment_type = ?kind,
total_segments = self.dir.segments.len()
)))]
pub fn read_segments<F>(
&self,
kind: SegmentType,
mut callback: F,
) -> Result<()>
where
F: FnMut(usize, &[u8]) -> Result<()>,
{
for (index, entry) in self.dir.entries.iter().enumerate() {
if entry.kind == kind {
let bytes = self.r.get_at(entry.offset)?;
if crc32(bytes) != entry.crc32 {
return Err(FileError::CrcMismatch(entry.offset));
}
callback(index, bytes)?;
}
}
Ok(())
}
}