use std::{
collections::{HashMap, HashSet, VecDeque},
fs::File,
io::{Read, Seek, SeekFrom},
path::Path,
};
use sha2::{Digest, Sha256};
use crate::{
Error, Result,
format::{
BLOCK_SIZE, ENTRIES_PER_OFFSET_RECORD, EntryData, FILE_ENTRY_SIZE, FOOTER_SIZE, FileEntry,
Footer, MAGIC, OFFSET_RECORD_SIZE, OffsetRecord, ROOT_NAME_OFFSET, VERSION_1, eq_name,
path_components,
},
};
pub type NodeHandle = u32;
pub const ROOT_NODE: NodeHandle = 0;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EntryKind {
File,
Directory,
}
#[derive(Clone, Copy, Debug)]
pub struct DirEntry<'a> {
pub name: &'a [u8],
pub kind: EntryKind,
pub size: u64,
pub handle: NodeHandle,
}
impl DirEntry<'_> {
#[must_use]
pub fn is_file(&self) -> bool {
self.kind == EntryKind::File
}
#[must_use]
pub fn is_directory(&self) -> bool {
self.kind == EntryKind::Directory
}
}
pub struct ArchiveReader<R> {
source: R,
file_size: u64,
expected_hash: [u8; 32],
compressed_data_offset: u64,
compressed_data_size: u64,
offset_records: Vec<OffsetRecord>,
names: Vec<u8>,
entries: Vec<FileEntry>,
cache: BlockCache,
}
impl ArchiveReader<File> {
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
Self::new(File::open(path)?)
}
}
impl<R: Read + Seek> ArchiveReader<R> {
pub fn new(mut source: R) -> Result<Self> {
let file_size = source.seek(SeekFrom::End(0))?;
if file_size < FOOTER_SIZE as u64 {
return Err(Error::InvalidArchive(
"file is too small to contain a footer",
));
}
let footer_offset = file_size - FOOTER_SIZE as u64;
let mut footer_bytes = [0; FOOTER_SIZE];
read_exact_at(&mut source, footer_offset, &mut footer_bytes)?;
let footer = Footer::decode(&footer_bytes);
validate_footer(&footer, file_size, footer_offset)?;
let offset_section = footer.sections[1];
let offset_len = usize::try_from(offset_section.size)
.map_err(|_| Error::InvalidArchive("offset table is too large"))?;
if offset_len % OFFSET_RECORD_SIZE != 0 {
return Err(Error::InvalidArchive("offset table has a partial record"));
}
let mut bytes = vec![0; offset_len];
read_exact_at(&mut source, offset_section.offset, &mut bytes)?;
let offset_records = bytes
.chunks_exact(OFFSET_RECORD_SIZE)
.map(OffsetRecord::decode)
.collect::<Vec<_>>();
let names_section = footer.sections[2];
let names_len = usize::try_from(names_section.size)
.map_err(|_| Error::InvalidArchive("name table is too large"))?;
let mut names = vec![0; names_len];
read_exact_at(&mut source, names_section.offset, &mut names)?;
let tree_section = footer.sections[3];
let tree_len = usize::try_from(tree_section.size)
.map_err(|_| Error::InvalidArchive("file tree is too large"))?;
if tree_len == 0 || tree_len % FILE_ENTRY_SIZE != 0 {
return Err(Error::InvalidArchive("file tree has an invalid size"));
}
let mut bytes = vec![0; tree_len];
read_exact_at(&mut source, tree_section.offset, &mut bytes)?;
let entries = bytes
.chunks_exact(FILE_ENTRY_SIZE)
.map(FileEntry::decode)
.collect::<Vec<_>>();
validate_tree(&entries, &names, &offset_records)?;
Ok(Self {
source,
file_size,
expected_hash: footer.integrity_hash,
compressed_data_offset: footer.sections[0].offset,
compressed_data_size: footer.sections[0].size,
offset_records,
names,
entries,
cache: BlockCache::new(64),
})
}
pub fn lookup(&self, path: impl AsRef<[u8]>) -> Option<NodeHandle> {
self.lookup_kind(path, None)
}
pub fn lookup_kind(
&self,
path: impl AsRef<[u8]>,
expected: Option<EntryKind>,
) -> Option<NodeHandle> {
let mut current = ROOT_NODE;
for component in path_components(path.as_ref()) {
let entry = self.entries.get(current as usize)?;
let EntryData::Directory { start, count } = entry.data else {
return None;
};
let end = start.checked_add(count)?;
current = (start..end).find(|index| {
self.entries
.get(*index as usize)
.and_then(|entry| self.name(entry.name_offset).ok())
.is_some_and(|name| eq_name(name, component))
})?;
}
expected
.is_none_or(|kind| self.entry_kind(current) == Some(kind))
.then_some(current)
}
#[must_use]
pub fn entry_kind(&self, handle: NodeHandle) -> Option<EntryKind> {
self.entries
.get(handle as usize)
.map(|entry| match entry.data {
EntryData::File { .. } => EntryKind::File,
EntryData::Directory { .. } => EntryKind::Directory,
})
}
pub fn directory_len(&self, handle: NodeHandle) -> Result<u32> {
match self.entry(handle)?.data {
EntryData::Directory { count, .. } => Ok(count),
EntryData::File { .. } => Err(Error::NotADirectory),
}
}
pub fn directory_entry(&self, handle: NodeHandle, index: u32) -> Result<DirEntry<'_>> {
let EntryData::Directory { start, count } = self.entry(handle)?.data else {
return Err(Error::NotADirectory);
};
if index >= count {
return Err(Error::InvalidPath("directory entry index is out of range"));
}
let child_handle = start
.checked_add(index)
.ok_or(Error::InvalidArchive("directory range overflows"))?;
let child = self.entry(child_handle)?;
let (kind, size) = match child.data {
EntryData::File { size, .. } => (EntryKind::File, size),
EntryData::Directory { .. } => (EntryKind::Directory, 0),
};
Ok(DirEntry {
name: self.name(child.name_offset)?,
kind,
size,
handle: child_handle,
})
}
pub fn directory_entries(&self, handle: NodeHandle) -> Result<Vec<DirEntry<'_>>> {
(0..self.directory_len(handle)?)
.map(|index| self.directory_entry(handle, index))
.collect()
}
pub fn file_size(&self, handle: NodeHandle) -> Result<u64> {
match self.entry(handle)?.data {
EntryData::File { size, .. } => Ok(size),
EntryData::Directory { .. } => Err(Error::NotAFile),
}
}
pub fn read_file(
&mut self,
handle: NodeHandle,
offset: u64,
buffer: &mut [u8],
) -> Result<usize> {
let EntryData::File {
offset: file_offset,
size,
} = self.entry(handle)?.data
else {
return Err(Error::NotAFile);
};
if offset >= size || buffer.is_empty() {
return Ok(0);
}
let available = usize::try_from((size - offset).min(buffer.len() as u64)).unwrap();
let mut raw_offset = file_offset
.checked_add(offset)
.ok_or(Error::InvalidArchive("file offset overflows"))?;
let mut written = 0;
while written < available {
let block_index = raw_offset / BLOCK_SIZE as u64;
let block_offset = raw_offset as usize % BLOCK_SIZE;
let step = (available - written).min(BLOCK_SIZE - block_offset);
let block = self.load_block(block_index)?;
buffer[written..written + step]
.copy_from_slice(&block[block_offset..block_offset + step]);
written += step;
raw_offset += step as u64;
}
Ok(written)
}
pub fn read_file_to_end(&mut self, handle: NodeHandle) -> Result<Vec<u8>> {
let size = usize::try_from(self.file_size(handle)?).map_err(|_| Error::ArchiveTooLarge)?;
let mut bytes = vec![0; size];
let read = self.read_file(handle, 0, &mut bytes)?;
if read != size {
return Err(Error::InvalidArchive("file data ended unexpectedly"));
}
Ok(bytes)
}
pub fn verify_integrity(&mut self) -> Result<bool> {
const HASH_OFFSET_IN_FOOTER: u64 = 6 * 16;
let hash_start = self.file_size - FOOTER_SIZE as u64 + HASH_OFFSET_IN_FOOTER;
let hash_end = hash_start + 32;
self.source.seek(SeekFrom::Start(0))?;
let mut hasher = Sha256::new();
let mut absolute = 0_u64;
let mut buffer = [0_u8; 64 * 1024];
while absolute < self.file_size {
let wanted =
usize::try_from((self.file_size - absolute).min(buffer.len() as u64)).unwrap();
self.source.read_exact(&mut buffer[..wanted])?;
let overlap_start = absolute.max(hash_start);
let overlap_end = (absolute + wanted as u64).min(hash_end);
if overlap_start < overlap_end {
buffer[(overlap_start - absolute) as usize..(overlap_end - absolute) as usize]
.fill(0);
}
hasher.update(&buffer[..wanted]);
absolute += wanted as u64;
}
Ok(hasher.finalize().as_slice() == self.expected_hash)
}
pub fn into_inner(self) -> R {
self.source
}
fn entry(&self, handle: NodeHandle) -> Result<&FileEntry> {
self.entries
.get(handle as usize)
.ok_or(Error::InvalidPath("node handle is out of range"))
}
fn name(&self, offset: u32) -> Result<&[u8]> {
if offset == ROOT_NAME_OFFSET {
return Ok(&[]);
}
let offset = offset as usize;
let first = *self
.names
.get(offset)
.ok_or(Error::InvalidArchive("name offset is out of range"))?;
let (length, header_len) = if first & 0x80 == 0 {
(usize::from(first), 1)
} else {
let second = *self
.names
.get(offset + 1)
.ok_or(Error::InvalidArchive("extended name header is truncated"))?;
(usize::from(first & 0x7f) | (usize::from(second) << 7), 2)
};
let start = offset + header_len;
let end = start
.checked_add(length)
.ok_or(Error::InvalidArchive("name length overflows"))?;
self.names
.get(start..end)
.ok_or(Error::InvalidArchive("name is truncated"))
}
fn load_block(&mut self, block_index: u64) -> Result<&[u8]> {
if self.cache.contains(block_index) {
return Ok(self.cache.get(block_index).unwrap());
}
let record_index = usize::try_from(block_index / ENTRIES_PER_OFFSET_RECORD as u64)
.map_err(|_| Error::InvalidArchive("block index is too large"))?;
let sub_index = block_index as usize % ENTRIES_PER_OFFSET_RECORD;
let record = self
.offset_records
.get(record_index)
.ok_or(Error::InvalidArchive("file references a missing block"))?;
let relative_offset = record.sizes[..sub_index]
.iter()
.try_fold(record.base_offset, |offset, size| {
offset.checked_add(u64::from(*size) + 1)
})
.ok_or(Error::InvalidArchive("compressed block offset overflows"))?;
let compressed_size = usize::from(record.sizes[sub_index]) + 1;
let relative_end = relative_offset
.checked_add(compressed_size as u64)
.ok_or(Error::InvalidArchive("compressed block range overflows"))?;
if relative_end > self.compressed_data_size {
return Err(Error::InvalidArchive(
"compressed block is outside its section",
));
}
let mut compressed = vec![0; compressed_size];
read_exact_at(
&mut self.source,
self.compressed_data_offset + relative_offset,
&mut compressed,
)?;
let block = if compressed_size == BLOCK_SIZE {
compressed
} else {
zstd::bulk::decompress(&compressed, BLOCK_SIZE).map_err(Error::Io)?
};
if block.len() != BLOCK_SIZE {
return Err(Error::InvalidArchive(
"decompressed block has the wrong size",
));
}
self.cache.insert(block_index, block);
Ok(self.cache.get(block_index).unwrap())
}
}
fn read_exact_at(source: &mut (impl Read + Seek), offset: u64, bytes: &mut [u8]) -> Result<()> {
if bytes.is_empty() {
return Ok(());
}
source.seek(SeekFrom::Start(offset))?;
source.read_exact(bytes)?;
Ok(())
}
fn validate_footer(footer: &Footer, file_size: u64, footer_offset: u64) -> Result<()> {
if footer.magic != MAGIC {
return Err(Error::InvalidArchive("footer magic does not match"));
}
if footer.version != VERSION_1 {
return Err(Error::InvalidArchive("archive version is not supported"));
}
if footer.total_size != file_size {
return Err(Error::InvalidArchive("footer size does not match the file"));
}
if footer
.sections
.iter()
.any(|section| !section.is_within(footer_offset))
{
return Err(Error::InvalidArchive(
"a section is outside the archive body",
));
}
Ok(())
}
fn validate_tree(entries: &[FileEntry], names: &[u8], offsets: &[OffsetRecord]) -> Result<()> {
let root = entries
.first()
.ok_or(Error::InvalidArchive("file tree is empty"))?;
if !matches!(root.data, EntryData::Directory { .. }) || root.name_offset != ROOT_NAME_OFFSET {
return Err(Error::InvalidArchive(
"first file-tree entry is not the root",
));
}
let max_uncompressed = (offsets.len() as u64)
.checked_mul(ENTRIES_PER_OFFSET_RECORD as u64)
.and_then(|blocks| blocks.checked_mul(BLOCK_SIZE as u64))
.ok_or(Error::InvalidArchive("block table size overflows"))?;
for (index, entry) in entries.iter().enumerate() {
if index != 0 {
validate_name(names, entry.name_offset)?;
}
match entry.data {
EntryData::Directory { start, count } => {
let end = start
.checked_add(count)
.ok_or(Error::InvalidArchive("directory range overflows"))?;
if end as usize > entries.len() {
return Err(Error::InvalidArchive(
"directory range is outside the file tree",
));
}
}
EntryData::File { offset, size } => {
if offset
.checked_add(size)
.is_none_or(|end| end > max_uncompressed)
{
return Err(Error::InvalidArchive(
"file range is outside the block table",
));
}
}
}
}
let mut visited = vec![false; entries.len()];
let mut pending = vec![0_usize];
visited[0] = true;
while let Some(index) = pending.pop() {
let EntryData::Directory { start, count } = entries[index].data else {
continue;
};
let mut child_names = HashSet::with_capacity(count as usize);
for child in start..start + count {
let child = child as usize;
if visited[child] {
return Err(Error::InvalidArchive(
"file tree contains a cycle or shared child",
));
}
visited[child] = true;
let name = decoded_name(names, entries[child].name_offset)?;
let folded = name.iter().map(u8::to_ascii_lowercase).collect::<Vec<_>>();
if !child_names.insert(folded) {
return Err(Error::InvalidArchive(
"directory contains duplicate case-insensitive names",
));
}
pending.push(child);
}
}
if visited.contains(&false) {
return Err(Error::InvalidArchive(
"file tree contains an unreachable entry",
));
}
Ok(())
}
fn validate_name(names: &[u8], offset: u32) -> Result<()> {
decoded_name(names, offset).map(|_| ())
}
fn decoded_name(names: &[u8], offset: u32) -> Result<&[u8]> {
if offset == ROOT_NAME_OFFSET {
return Err(Error::InvalidArchive(
"non-root entry uses the root name marker",
));
}
let offset = offset as usize;
let first = *names
.get(offset)
.ok_or(Error::InvalidArchive("name offset is out of range"))?;
let (length, header) = if first & 0x80 == 0 {
(usize::from(first), 1)
} else {
let second = *names
.get(offset + 1)
.ok_or(Error::InvalidArchive("extended name header is truncated"))?;
(usize::from(first & 0x7f) | (usize::from(second) << 7), 2)
};
if length == 0 || offset + header + length > names.len() {
return Err(Error::InvalidArchive("entry name is empty or truncated"));
}
Ok(&names[offset + header..offset + header + length])
}
struct BlockCache {
capacity: usize,
blocks: HashMap<u64, Vec<u8>>,
order: VecDeque<u64>,
}
impl BlockCache {
fn new(capacity: usize) -> Self {
Self {
capacity,
blocks: HashMap::with_capacity(capacity),
order: VecDeque::with_capacity(capacity),
}
}
fn contains(&self, index: u64) -> bool {
self.blocks.contains_key(&index)
}
fn get(&mut self, index: u64) -> Option<&[u8]> {
if self.blocks.contains_key(&index) {
if let Some(position) = self.order.iter().position(|cached| *cached == index) {
self.order.remove(position);
}
self.order.push_back(index);
}
self.blocks.get(&index).map(Vec::as_slice)
}
fn insert(&mut self, index: u64, block: Vec<u8>) {
if self.blocks.len() == self.capacity
&& let Some(oldest) = self.order.pop_front()
{
self.blocks.remove(&oldest);
}
self.blocks.insert(index, block);
self.order.push_back(index);
}
}