sit/

archive.rs

use std::{
    fmt::Debug,
    fs,
    io::{self},
    path::{self},
    rc::Rc,
};

use binrw::BinReaderExt;
use macintosh_utils::Fork;

use crate::{
    Entry, Error, VerifyingEntryReader,
    structs::{
        Algorithm, ArchiveHeader, File, Version, v1,
        v5::{self, EntryBinReadArgs},
    },
    verify::VerifyingIterator,
};

pub struct EntryReader<'a, T: io::Read + io::Seek>(crate::algos::EntryReader<'a, T>, u16);

impl<'a, T: io::Read + io::Seek> EntryReader<'a, T> {
    #[inline]
    pub(crate) fn try_from(
        reader: &'a mut Rc<T>,
        algo: Algorithm,
        uncompressed_size: u64,
        compressed_size: usize,
        offset: u64,
        checksum: u16,
    ) -> Result<Self, Error> {
        Ok(Self(
            crate::algos::EntryReader::try_from(
                reader,
                algo,
                uncompressed_size,
                compressed_size,
                offset,
            )?,
            checksum,
        ))
    }

    pub fn verifying(self) -> VerifyingEntryReader<'a, T> {
        self.0.verifying(self.1)
    }

    pub fn verify(self) -> Result<(), Error> {
        // Arsenic compression uses 32-bit checksums interleaved with the compressed blocks
        // so it get's special treatment
        let is_arsenic = matches!(self.0, crate::algos::EntryReader::Arsenic { .. });

        VerifyingEntryReader::new(self.0, self.1, is_arsenic).slurp()
    }
}

impl<'a, T: io::Read + io::Seek> io::Read for EntryReader<'a, T> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl<'a, T: io::Read + io::Seek> io::Seek for EntryReader<'a, T> {
    #[inline]
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        self.0.seek(pos)
    }

    #[inline]
    fn stream_len(&mut self) -> io::Result<u64> {
        self.0.stream_len()
    }

    #[inline]
    fn stream_position(&mut self) -> io::Result<u64> {
        self.0.stream_position()
    }
}

#[derive(Debug)]
pub struct Archive<R> {
    header_location: u64,
    header: ArchiveHeader,
    inner: Rc<R>,
}

impl Archive<fs::File> {
    pub fn open_path<P: AsRef<path::Path>>(p: P) -> Result<Self, Error> {
        Archive::try_from(fs::File::open(p.as_ref())?)
    }
}

impl<R: io::Seek + io::Read> Archive<R> {
    pub fn try_from(mut inner: R) -> Result<Self, Error> {
        let header_location = inner.stream_position()?;

        match inner.read_be() {
            Ok(header) => Ok(Self {
                inner: Rc::new(inner),
                header_location,
                header,
            }),
            Err(binrw::Error::AssertFail { .. }) => Err(Error::InvalidFile(
                crate::error::InvalidFileReason::InvalidHeader,
            )),
            Err(e) => Err(e.into()),
        }
    }

    pub fn version(&self) -> Version {
        self.header.version()
    }

    pub fn header(&self) -> &ArchiveHeader {
        &self.header
    }

    /// Verify the archive's integrity by validating all checksums
    pub fn verify(&mut self) -> Result<(), Error> {
        if !self.header.checksum_valid() {
            return Err(Error::ChecksumMismatch(
                crate::error::ChecksumLocation::ArchiveHeader,
            ));
        }

        self.iter().try_for_each(|e| self.verify_entry(e))
    }

    /// Open an archive entry for reading
    pub fn open_fork<'a, E: ReadableEntry>(
        &'a mut self,
        entry: &E,
        fork: Fork,
    ) -> Result<EntryReader<'a, R>, Error> {
        let algo = entry.algorithm(fork);
        let compressed_size = entry.compressed_size(fork);
        let uncompressed_size = entry.uncompressed_size(fork);
        let offset = entry.offset(fork);
        let checksum = entry.checksum(fork);

        if entry.encrypted(fork) {
            return Err(Error::UnsupportedFeature(
                crate::error::UnsupportedFeature::Encryption,
            ));
        }

        EntryReader::try_from(
            &mut self.inner,
            algo,
            uncompressed_size as u64,
            compressed_size,
            offset,
            checksum,
        )
    }

    /// Iterate through all archive entries
    pub fn iter(&self) -> EntryIterator<R> {
        // SAFETY: Archives can only be iterated once, and opening an entry locks the archive until
        // the handle is released. That ensures we don't modify the underlying reader concurrently
        self.ensure_not_iterating();

        let catalog_offset = match &self.header {
            ArchiveHeader::V1(hdr) => self.header_location + hdr.first_entry_offset(),
            ArchiveHeader::V5(hdr) => self.header_location + hdr.first_entry_offset(),
        };

        EntryIterator::new(
            self.inner.clone(),
            self.header.entry_count(),
            catalog_offset,
            matches!(self.header, ArchiveHeader::V1(_)),
        )
    }

    /// Returns the underlying reader
    pub fn into_inner(self) -> R {
        self.ensure_not_iterating();

        // SAFETY: Since no one is iterating, and any open entry would borrow the archive mutably
        // this is safe
        let Archive { inner, .. } = self;
        unsafe { Rc::try_unwrap(inner).unwrap_unchecked() }
    }

    pub fn reset(&mut self) -> Result<(), Error> {
        self.ensure_not_iterating();

        // SAFETY: Since no one is iterating, and any open entry would borrow the archive mutably,
        // this is safe
        unsafe {
            Rc::get_mut_unchecked(&mut self.inner)
                .seek(io::SeekFrom::Start(self.header_location))?;
        }

        Ok(())
    }

    fn ensure_not_iterating(&self) {
        if Rc::strong_count(&self.inner) != 1 || Rc::weak_count(&self.inner) != 0 {
            panic!("Can not modify archive while an iterator is runnning")
        }
    }

    /// Verify integrity of an entry
    pub fn verify_entry(&mut self, e: Entry) -> Result<(), Error> {
        if let Entry::Directory(crate::structs::Directory::V5(dir)) = &e
            && dir.marks_end()
        {
            return Ok(());
        }

        if !e.is_file() {
            return Ok(());
        }

        if e.has(Fork::Resource) {
            self.open_fork(&e, Fork::Resource)?.verifying().slurp()?;
        }

        if e.has(Fork::Data) {
            self.open_fork(&e, Fork::Data)?.verifying().slurp()?;
        }

        Ok(())
    }
}

/// Trait implemented by archive entries to create a decompression stream
pub trait ReadableEntry {
    /// Algorithm used to compress the entry's data or resource fork
    fn algorithm(&self, fork: Fork) -> Algorithm;
    /// Amount of bytes occupied by the specified fork
    fn compressed_size(&self, fork: Fork) -> usize;
    /// Size of the fork in bytes after decompression
    fn uncompressed_size(&self, fork: Fork) -> usize;
    fn encrypted(&self, fork: Fork) -> bool;
    /// Offset from start of the archive header to beginning of the compressed data
    fn offset(&self, fork: Fork) -> u64;
    /// Checksum of the uncompressed data
    fn checksum(&self, fork: Fork) -> u16;
}

impl ReadableEntry for Entry {
    fn encrypted(&self, fork: Fork) -> bool {
        match self {
            Entry::File(e) => e.encrypted(fork),
            Entry::Directory(_) => false,
            Entry::DirectoryEnd(_) => false,
        }
    }

    fn algorithm(&self, fork: Fork) -> Algorithm {
        match self {
            Entry::File(e) => e.algorithm(fork),
            Entry::Directory(e) => e.algorithm(fork),
            Entry::DirectoryEnd(_) => Algorithm::None,
        }
    }

    fn compressed_size(&self, fork: Fork) -> usize {
        match self {
            Entry::File(e) => e.compressed_size(fork),
            Entry::Directory(e) => e.compressed_size(fork),
            Entry::DirectoryEnd(_) => 0,
        }
    }

    fn uncompressed_size(&self, fork: Fork) -> usize {
        match self {
            Entry::File(e) => e.uncompressed_size(fork),
            Entry::Directory(e) => e.uncompressed_size(fork),
            Entry::DirectoryEnd(_) => 0,
        }
    }

    fn offset(&self, fork: Fork) -> u64 {
        match self {
            Entry::File(e) => e.offset(fork),
            Entry::Directory(e) => e.offset(fork),
            Entry::DirectoryEnd(off) => *off,
        }
    }

    fn checksum(&self, fork: Fork) -> u16 {
        match self {
            Entry::File(file) => file.checksum(fork),
            Entry::Directory(_) => 0,
            Entry::DirectoryEnd(_) => 0,
        }
    }
}

impl ReadableEntry for File {
    #[inline]
    fn algorithm(&self, fork: Fork) -> Algorithm {
        File::compression_method(self, fork)
    }

    #[inline]
    fn compressed_size(&self, fork: Fork) -> usize {
        File::compressed_size(self, fork)
    }

    #[inline]
    fn uncompressed_size(&self, fork: Fork) -> usize {
        File::uncompressed_size(self, fork)
    }

    #[inline]
    fn encrypted(&self, fork: Fork) -> bool {
        File::encrypted(self, fork)
    }

    #[inline]
    fn offset(&self, fork: Fork) -> u64 {
        File::offset(self, fork)
    }

    #[inline]
    fn checksum(&self, fork: Fork) -> u16 {
        File::checksum(self, fork)
    }
}

impl ReadableEntry for v5::File {
    #[inline]
    fn algorithm(&self, fork: Fork) -> Algorithm {
        v5::File::compression_method(self, fork)
    }

    #[inline]
    fn compressed_size(&self, fork: Fork) -> usize {
        v5::File::compressed_size(self, fork)
    }

    #[inline]
    fn uncompressed_size(&self, fork: Fork) -> usize {
        v5::File::uncompressed_size(self, fork)
    }

    #[inline]
    fn encrypted(&self, fork: Fork) -> bool {
        v5::File::encrypted(self, fork)
    }

    #[inline]
    fn offset(&self, fork: Fork) -> u64 {
        v5::File::offset(self, fork)
    }

    #[inline]
    fn checksum(&self, fork: Fork) -> u16 {
        v5::File::checksum(self, fork)
    }
}

impl ReadableEntry for v1::File {
    #[inline]
    fn algorithm(&self, fork: Fork) -> Algorithm {
        v1::File::compression_method(self, fork)
    }

    #[inline]
    fn compressed_size(&self, fork: Fork) -> usize {
        v1::File::compressed_size(self, fork)
    }

    #[inline]
    fn uncompressed_size(&self, fork: Fork) -> usize {
        v1::File::uncompressed_size(self, fork)
    }

    #[inline]
    fn encrypted(&self, fork: Fork) -> bool {
        v1::File::encrypted(self, fork)
    }

    #[inline]
    fn offset(&self, fork: Fork) -> u64 {
        v1::File::offset(self, fork)
    }

    #[inline]
    fn checksum(&self, fork: Fork) -> u16 {
        v1::File::checksum(self, fork)
    }
}

pub struct EntryIterator<R: io::Read + io::Seek> {
    next_offset: u64,
    next_file_index: usize,
    stack: Vec<u32>,
    v1: bool,

    pub(crate) reader: Rc<R>,
}

impl<R: io::Read + io::Seek> EntryIterator<R> {
    fn new(reader: Rc<R>, entry_count: usize, offset: u64, v1: bool) -> Self {
        Self {
            next_offset: offset,
            next_file_index: 0,
            stack: vec![entry_count as u32],
            reader,
            v1,
        }
    }

    pub fn verifying(self) -> VerifyingIterator<R> {
        let Self {
            next_offset,
            stack,
            reader,
            v1,
            next_file_index,
        } = self;

        VerifyingIterator {
            next_offset,
            stack,
            reader,
            v1,
            _next_file_index: next_file_index,
        }
    }
}

impl<R: io::Read + io::Seek> Iterator for EntryIterator<R> {
    type Item = Entry;

    fn next(&mut self) -> Option<Self::Item> {
        let reader = unsafe { Rc::get_mut_unchecked(&mut self.reader) };

        match self.stack.last_mut() {
            None => return None,
            Some(0) => {
                self.stack.pop();

                // Don't send directory end marker for the root catalog
                return if self.stack.is_empty() {
                    None
                } else {
                    Some(Entry::DirectoryEnd(self.next_offset))
                };
            }
            Some(d) => *d -= 1,
        }

        let Ok(entry_offset) = reader.seek(io::SeekFrom::Start(self.next_offset)) else {
            log::warn!("Failed seeking to next archive entry");
            return None;
        };

        if self.v1 {
            let Ok(entry) = reader.read_be::<v1::Entry>() else {
                return None;
            };

            let Ok(payload_offset) = reader.stream_position() else {
                return None;
            };

            match entry {
                v1::Entry::Directory(dir) => {
                    self.next_offset = payload_offset;
                    self.stack.push(u32::MAX);

                    Some(Entry::Directory(dir.into()))
                }
                v1::Entry::DirectoryEnd => {
                    self.next_offset = payload_offset;
                    self.stack.pop();

                    Some(Entry::DirectoryEnd(payload_offset))
                }
                v1::Entry::File(mut file) => {
                    self.next_offset = payload_offset
                        + file.data_compressed_size as u64
                        + file.rsrc_compressed_size as u64;
                    log::debug!(
                        "Entry: {}, {:?} {} 0x{:04x}",
                        file.file_name,
                        file.data_compression,
                        file.data_compressed_size,
                        file.checksum(Fork::Data)
                    );

                    file.index = self.next_file_index;
                    self.next_file_index += 1;

                    Some(Entry::File(file.into()))
                }
            }
        } else {
            let Ok(entry) = reader.read_be_args::<v5::Entry>(
                EntryBinReadArgs::builder().offset(entry_offset).finalize(),
            ) else {
                return None;
            };

            let Ok(payload_offset) = reader.stream_position() else {
                return None;
            };

            match entry {
                v5::Entry::Directory(dir) => {
                    if dir.marks_end() {
                        self.next_offset = payload_offset;
                        return self.next();
                    }

                    self.stack.push(dir.child_count);
                    self.next_offset = dir.first_child_offset;

                    Some(Entry::Directory(dir.into()))
                }
                v5::Entry::File(mut file) => {
                    self.next_offset = file.next_entry_offset as u64;
                    file.payload_offset = payload_offset;

                    file.index = self.next_file_index;
                    self.next_file_index += 1;

                    Some(Entry::File(file.into()))
                }
            }
        }
    }
}