zesven 2.0.0

//! Writer initialization and finalization.
//!
//! This module provides methods for creating writers and finishing archive writing,
//! including signature header writing.

use std::fs::File;
use std::io::{BufWriter, Seek, SeekFrom, Write};
use std::path::Path;

use crate::format::{SIGNATURE, SIGNATURE_HEADER_SIZE};
use crate::volume::{MultiVolumeWriter, VolumeConfig};
use crate::{Error, Result};

use super::options::{WriteOptions, WriteResult};
use super::{StreamInfo, Writer, WriterState};

impl Writer<BufWriter<File>> {
    /// Creates a new archive file at the given path.
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the archive file to create
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be created.
    pub fn create_path(path: impl AsRef<Path>) -> Result<Self> {
        let file = File::create(path.as_ref()).map_err(Error::Io)?;
        let writer = BufWriter::new(file);
        Self::create(writer)
    }

    /// Finishes writing the archive.
    ///
    /// # Returns
    ///
    /// A WriteResult with statistics about the written archive.
    ///
    /// # Errors
    ///
    /// Returns an error if header writing fails.
    pub fn finish(self) -> Result<WriteResult> {
        let (result, _sink) = self.finish_into_inner()?;
        Ok(result)
    }
}

impl Writer<std::io::Cursor<Vec<u8>>> {
    /// Finishes writing the archive to an owned cursor.
    pub fn finish(self) -> Result<WriteResult> {
        let (result, _sink) = self.finish_into_inner()?;
        Ok(result)
    }
}

impl Writer<std::io::Cursor<&mut Vec<u8>>> {
    /// Finishes writing the archive to a borrowed cursor.
    pub fn finish(self) -> Result<WriteResult> {
        let (result, _sink) = self.finish_into_inner()?;
        Ok(result)
    }
}

impl Writer<MultiVolumeWriter> {
    /// Creates a new multi-volume archive writer.
    ///
    /// The archive will be split across multiple files when each volume
    /// reaches the configured size limit.
    ///
    /// # Arguments
    ///
    /// * `config` - Volume configuration specifying size and base path
    ///
    /// # Errors
    ///
    /// Returns an error if the first volume file cannot be created.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use zesven::{Writer, VolumeConfig, ArchivePath};
    ///
    /// let config = VolumeConfig::new("archive.7z", 50 * 1024 * 1024); // 50 MB volumes
    /// let mut writer = Writer::create_multivolume(config)?;
    /// writer.add_bytes(ArchivePath::new("data.bin")?, &large_data)?;
    /// let result = writer.finish()?;
    /// println!("Created {} volumes", result.volume_count);
    /// ```
    pub fn create_multivolume(config: VolumeConfig) -> Result<Self> {
        let writer = MultiVolumeWriter::create(config)?;
        Self::create(writer)
    }

    /// Finishes writing the multi-volume archive.
    ///
    /// This finalizes all volumes and returns a WriteResult with volume information.
    ///
    /// # Returns
    ///
    /// A WriteResult with statistics including volume_count and volume_sizes.
    ///
    /// # Errors
    ///
    /// Returns an error if header writing or volume finalization fails.
    pub fn finish(self) -> Result<WriteResult> {
        let (mut result, mv_writer) = self.finish_into_inner()?;

        // Finalize the multi-volume writer and get volume sizes
        let volume_sizes = mv_writer.finish()?;
        result.volume_count = volume_sizes.len() as u32;
        result.volume_sizes = volume_sizes;

        Ok(result)
    }
}

impl<W: Write + Seek> Writer<W> {
    /// Creates a new archive writer.
    ///
    /// # Arguments
    ///
    /// * `sink` - The writer to output archive data to
    ///
    /// # Errors
    ///
    /// Returns an error if the initial seek fails.
    pub fn create(mut sink: W) -> Result<Self> {
        // Where the archive begins, which is wherever the sink happens to be
        // rather than nought: a caller may be writing after a prefix of their
        // own. Everything measured or seeked to is relative to this.
        let start_pos = sink.stream_position().map_err(Error::Io)?;

        // Reserve space for signature header (32 bytes) by writing zeros
        let placeholder = [0u8; SIGNATURE_HEADER_SIZE as usize];
        sink.write_all(&placeholder).map_err(Error::Io)?;

        Ok(Self {
            sink,
            start_pos,
            options: WriteOptions::default(),
            state: WriterState::AcceptingEntries,
            entries: Vec::new(),
            stream_info: StreamInfo::default(),
            compressed_bytes: 0,
            solid_buffer: Vec::new(),
            solid_buffer_size: 0,
            pending_batch: Vec::new(),
            pending_batch_size: 0,
            active_options: std::sync::Arc::new(WriteOptions::default()),
            last_path: None,
            #[cfg(feature = "aes")]
            archive_salt: None,
            #[cfg(feature = "aes")]
            archive_password: None,
        })
    }

    /// Sets the write options.
    ///
    /// Entries already accepted keep the options they were accepted under.
    /// Anything still waiting is written out with those before the next entry
    /// is taken, so a change here never reaches back over work already done.
    pub fn options(mut self, options: WriteOptions) -> Self {
        self.options = options.clone();
        self.active_options = std::sync::Arc::new(options);
        self
    }

    /// Refuses a password that differs from the one this archive is keyed on.
    ///
    /// Checked before an entry is accepted, so a caller learns while the entry
    /// is still theirs to reconsider. It only reports; the password is fixed
    /// when one is actually used to derive a key, so an entry that is rejected,
    /// or one that never gets encrypted such as a directory, leaves the archive
    /// free to be keyed on something else.
    #[cfg(feature = "aes")]
    pub(crate) fn check_password(&self, options: &WriteOptions) -> Result<()> {
        let (Some(held), Some(wanted)) = (&self.archive_password, &options.password) else {
            return Ok(());
        };
        if !options.is_encrypted() || held.as_utf16_le() == wanted.as_utf16_le() {
            return Ok(());
        }

        Err(Error::InvalidFormat(
            "the password cannot be changed once an entry has been \
             encrypted: one archive is opened with one password, and \
             entries written under different ones cannot all be read"
                .into(),
        ))
    }

    /// Records the password a key is about to be derived from.
    ///
    /// Called from the places that actually encrypt, so what is remembered is
    /// what the archive is really keyed on.
    #[cfg(feature = "aes")]
    pub(crate) fn hold_password(&mut self, password: &crate::crypto::Password) -> Result<()> {
        match &self.archive_password {
            Some(held) if held.as_utf16_le() != password.as_utf16_le() => {
                Err(Error::InvalidFormat(
                    "the password cannot be changed once an entry has been \
                     encrypted: one archive is opened with one password, and \
                     entries written under different ones cannot all be read"
                        .into(),
                ))
            }
            Some(_) => Ok(()),
            None => {
                self.archive_password = Some(password.clone());
                Ok(())
            }
        }
    }

    /// Writes out anything waiting under options that are no longer current.
    ///
    /// Called before an entry is accepted, so the buffers never hold work from
    /// two different sets of options.
    pub(crate) fn settle_stale_buffers(&mut self) -> Result<()> {
        let stale = self
            .pending_batch
            .first()
            .or_else(|| self.solid_buffer.first())
            .is_some_and(|entry| !std::sync::Arc::ptr_eq(&entry.options, &self.active_options));

        if stale {
            self.flush_buffered_entries()?;
        }
        Ok(())
    }

    /// Writes out every entry waiting to be compressed, in the order they came.
    pub(crate) fn flush_buffered_entries(&mut self) -> Result<()> {
        self.flush_pending_batch()?;
        if !self.solid_buffer.is_empty() {
            self.flush_solid_buffer()?;
        }
        Ok(())
    }

    /// Finishes writing the archive and returns the underlying sink.
    ///
    /// This is useful when you need access to the written data, such as
    /// when writing to a `Cursor<Vec<u8>>` and need to retrieve the buffer.
    ///
    /// # Returns
    ///
    /// A tuple of (WriteResult, W) where W is the underlying sink.
    ///
    /// # Errors
    ///
    /// Returns an error if header writing fails.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use zesven::write::Writer;
    /// use std::io::Cursor;
    ///
    /// let mut writer = Writer::create(Cursor::new(Vec::new()))?;
    /// writer.add_bytes("test.txt".try_into()?, b"Hello")?;
    /// let (result, cursor) = writer.finish_into_inner()?;
    /// let archive_bytes = cursor.into_inner();
    /// ```
    pub fn finish_into_inner(mut self) -> Result<(WriteResult, W)> {
        // Also rejects a writer poisoned by a partial write, so a failure
        // cannot be turned into an archive by ignoring its error.
        self.ensure_accepting_entries()?;
        // The header is encrypted with the archive's password, which must be
        // the one its entries were encrypted with.
        #[cfg(feature = "aes")]
        {
            let options = self.options.clone();
            self.check_password(&options)?;
        }
        self.state = WriterState::Building;

        // Everything still waiting, written with the options it was accepted
        // under and in the order it arrived.
        self.flush_buffered_entries()?;

        let header_data = self.encode_header()?;

        // An encrypted header is stored as a packed stream in the data area,
        // followed by the small structure that describes it. The structure is the
        // archive's next header, so it must be written last and its position is
        // what the signature header points at.
        #[cfg(feature = "aes")]
        if self.options.is_header_encrypted() {
            let payload_pos = self.sink.stream_position().map_err(Error::Io)?;
            let nonce = self.nonce_for_stream()?;
            let (payload, structure) = self.encode_encrypted_header(
                &header_data,
                payload_pos - self.start_pos - SIGNATURE_HEADER_SIZE,
                nonce,
            )?;

            self.sink.write_all(&payload).map_err(Error::Io)?;
            let header_pos = self.sink.stream_position().map_err(Error::Io)?;
            self.sink.write_all(&structure).map_err(Error::Io)?;
            let archive_len = self.sink.stream_position().map_err(Error::Io)?;
            self.write_signature_header(header_pos, &structure)?;

            return self.finish_state(archive_len);
        }

        // A plain header is written compressed when that is smaller, in the same
        // shape as the encrypted one: payload in the data area, structure last.
        #[cfg(feature = "lzma2")]
        {
            let payload_pos = self.sink.stream_position().map_err(Error::Io)?;
            if let Some((payload, structure)) = self.encode_compressed_header(
                &header_data,
                payload_pos - self.start_pos - SIGNATURE_HEADER_SIZE,
            )? {
                self.sink.write_all(&payload).map_err(Error::Io)?;
                let header_pos = self.sink.stream_position().map_err(Error::Io)?;
                self.sink.write_all(&structure).map_err(Error::Io)?;
                let archive_len = self.sink.stream_position().map_err(Error::Io)?;
                self.write_signature_header(header_pos, &structure)?;

                return self.finish_state(archive_len);
            }
        }

        let header_pos = self.sink.stream_position().map_err(Error::Io)?;
        self.sink.write_all(&header_data).map_err(Error::Io)?;
        let archive_len = self.sink.stream_position().map_err(Error::Io)?;

        // Write signature header at start
        self.write_signature_header(header_pos, &header_data)?;

        self.finish_state(archive_len)
    }

    /// Marks the writer finished and builds the write result.
    ///
    /// `archive_len` is where writing ended, taken before the signature header
    /// was written: that seeks back to the start, so asking the sink afterwards
    /// reported the 32 bytes of the signature as the size of the archive.
    fn finish_state(mut self, archive_len: u64) -> Result<(WriteResult, W)> {
        // The signature header is written last, over the start of the archive,
        // and a buffered sink may still be holding it. Flushing here rather
        // than leaving it to `Drop` is what turns a failed write into an error
        // the caller sees: `BufWriter::drop` discards the result.
        self.sink.flush().map_err(Error::Io)?;

        self.state = WriterState::Finished;

        // Build result
        let result = WriteResult {
            entries_written: self
                .entries
                .iter()
                .filter(|e| !e.meta.is_directory && !e.meta.is_anti)
                .count(),
            directories_written: self.entries.iter().filter(|e| e.meta.is_directory).count(),
            total_size: self.entries.iter().map(|e| e.uncompressed_size).sum(),
            compressed_size: self.compressed_bytes,
            volume_count: 1,
            volume_sizes: vec![archive_len - self.start_pos],
        };

        Ok((result, self.sink))
    }

    /// Writes the signature header at the start of the file.
    pub(crate) fn write_signature_header(
        &mut self,
        header_pos: u64,
        header_data: &[u8],
    ) -> Result<()> {
        // Calculate values
        let next_header_offset = header_pos - self.start_pos - SIGNATURE_HEADER_SIZE;
        let next_header_size = header_data.len() as u64;
        let next_header_crc = crc32fast::hash(header_data);

        // Build start header (20 bytes)
        let mut start_header = Vec::with_capacity(20);
        start_header.extend_from_slice(&next_header_offset.to_le_bytes());
        start_header.extend_from_slice(&next_header_size.to_le_bytes());
        start_header.extend_from_slice(&next_header_crc.to_le_bytes());

        let start_header_crc = crc32fast::hash(&start_header);

        // Seek to where the archive starts, which is not necessarily the start
        // of the sink.
        let start = self.start_pos;
        self.sink.seek(SeekFrom::Start(start)).map_err(Error::Io)?;

        // Write signature (6 bytes)
        self.sink.write_all(SIGNATURE).map_err(Error::Io)?;

        // Write version (2 bytes)
        self.sink.write_all(&[0x00, 0x04]).map_err(Error::Io)?;

        // Write start header CRC (4 bytes)
        self.sink
            .write_all(&start_header_crc.to_le_bytes())
            .map_err(Error::Io)?;

        // Write start header (20 bytes)
        self.sink.write_all(&start_header).map_err(Error::Io)?;

        Ok(())
    }

    /// The header's view of what this writer has produced.
    pub(crate) fn header_model(&self) -> super::header_encode::HeaderModel<'_> {
        super::header_encode::HeaderModel {
            stream_info: &self.stream_info,
            entries: &self.entries,
            options: &self.options,
        }
    }

    /// Encodes the archive header from that view.
    pub(crate) fn encode_header(&self) -> Result<Vec<u8>> {
        self.header_model().encode_header()
    }

    /// Writes to the sink, marking the writer unusable if it fails.
    ///
    /// Once any of an entry's bytes are in the sink, a failure has left data
    /// that belongs to no folder, and every folder written after it would be
    /// found at the wrong offset. Every write of entry data goes through here
    /// so that no path can quietly leave the writer usable.
    pub(crate) fn write_entry_bytes(&mut self, data: &[u8]) -> Result<()> {
        match self.sink.write_all(data) {
            Ok(()) => Ok(()),
            Err(e) => self.fail(Error::Io(e)),
        }
    }

    /// Marks the writer unusable and returns the error that caused it.
    ///
    /// For failures that happen once bytes are already in the sink: whatever
    /// was written belongs to no folder, and every folder after it would be
    /// found at the wrong offset.
    pub(crate) fn fail<T>(&mut self, error: Error) -> Result<T> {
        self.state = WriterState::Failed;
        Err(error)
    }

    /// Ensures the writer is in the AcceptingEntries state.
    pub(crate) fn ensure_accepting_entries(&self) -> Result<()> {
        if self.state == WriterState::Failed {
            return Err(Error::InvalidFormat(
                "an earlier entry failed partway through writing; \
                 this archive cannot be completed"
                    .into(),
            ));
        }
        if self.state != WriterState::AcceptingEntries {
            return Err(Error::InvalidFormat(
                "Writer is not accepting entries".into(),
            ));
        }

        self.options.validate()?;

        Ok(())
    }
}