sett 0.3.0

Rust port of sett (data compression, encryption and transfer tool).
Documentation
use std::{
    borrow::Cow,
    io::{self, Write},
};

use anyhow::Result;
use chrono::{DateTime, Utc};

mod reader;
mod spec;

pub(crate) use reader::ZipReader;

#[derive(Debug, Default, Clone)]
struct ZipFile {
    name: String,
    offset: u64,
    size: u64,
    hasher: crc32fast::Hasher,
    crc32: u32,
    flags: u16,
    external_attributes: u32,
    timestamp: DateTime<Utc>,
    compression_method: u16,
}

impl ZipFile {
    fn new(name: String, offset: u64, opts: &ZipFileOpts) -> Self {
        Self {
            name,
            offset,
            timestamp: opts.timestamp,
            flags: 1u16 << 3 | 1u16 << 11, // Streaming + UTF-8
            external_attributes: opts.permissions << 16,
            ..Default::default()
        }
    }
    fn to_central_directory_header(&self) -> spec::CentralDirectoryHeader {
        spec::CentralDirectoryHeader {
            compression_method: self.compression_method,
            crc32: self.crc32,
            external_attributes: self.external_attributes,
            flags: self.flags,
            modified: self.timestamp,
            name: Cow::Borrowed(&self.name),
            offset: self.offset,
            size: self.size,
        }
    }
    fn to_local_file_header(&self) -> spec::LocalFileHeader {
        spec::LocalFileHeader {
            flags: self.flags,
            modified: self.timestamp,
            name: Cow::Borrowed(&self.name),
            ..Default::default()
        }
    }
    fn to_data_descriptor(&self) -> spec::DataDescriptor {
        spec::DataDescriptor {
            crc32: self.crc32,
            size: self.size,
        }
    }
    fn finish(mut self) -> Self {
        self.crc32 = self.hasher.clone().finalize();
        self
    }
}

pub(super) struct ZipFileOpts {
    timestamp: DateTime<Utc>,
    permissions: u32,
}

impl Default for ZipFileOpts {
    fn default() -> Self {
        Self {
            timestamp: Utc::now(),
            // Default permissions are rw-r--r--
            permissions: 0o100644,
        }
    }
}

enum ZipArchiveState {
    Start,
    /// Only write file data
    FileOpen(ZipFile),
    /// Only write metadata, not files
    FileClosed,
}

impl ZipArchiveState {
    fn unwrap(self) -> ZipFile {
        if let ZipArchiveState::FileOpen(f) = self {
            f
        } else {
            panic!("Only `FileOpen` can be unwrapped");
        }
    }
}

pub(super) struct ZipArchive<W: Write> {
    inner: W,
    files: Vec<ZipFile>,
    state: ZipArchiveState,
    size: u64,
}

impl<W: io::Write> ZipArchive<W> {
    pub(super) fn new(writer: W) -> Self {
        Self {
            inner: writer,
            files: Vec::new(),
            state: ZipArchiveState::Start,
            size: 0,
        }
    }

    fn finish_file(&mut self) -> Result<()> {
        if let ZipArchiveState::FileOpen(_) = &self.state {
            let closed = std::mem::replace(&mut self.state, ZipArchiveState::FileClosed)
                .unwrap()
                .finish();
            self.write_all(&closed.to_data_descriptor().build())?;
            self.files.push(closed);
        }
        Ok(())
    }

    pub(super) fn add(&mut self, file: &str, opts: &ZipFileOpts) -> Result<()> {
        self.finish_file()?;
        let f = ZipFile::new(file.into(), self.size, opts);
        self.write_all(&f.to_local_file_header().build())?;
        self.state = ZipArchiveState::FileOpen(f);
        Ok(())
    }

    pub(super) fn finish(&mut self) -> Result<u64> {
        self.finish_file()?;
        let central_directory_offset = self.size;
        let headers = self
            .files
            .iter()
            .map(|f| f.to_central_directory_header().build())
            .collect::<Vec<_>>();
        for header in &headers {
            self.write_all(header)?;
        }

        let central_directory_size = self.size - central_directory_offset;
        let zip64_central_directory_offset = self.size;

        if self.size > spec::ZIP64_THRESHOLD_BYTES
            || self.files.len() as u64 > spec::ZIP64_THRESHOLD_FILES
        {
            let zip64_end = spec::Zip64CentralDirectoryEnd {
                disk_number_of_records: self.files.len() as u64,
                total_number_of_records: self.files.len() as u64,
                central_directory_size,
                central_directory_offset,
                ..Default::default()
            };
            self.write_all(&zip64_end.build())?;
            let zip64_end_locator = spec::Zip64CentralDirectoryEndLocator {
                disk_with_zip64_central_directory_end: 0,
                zip64_central_directory_end_offset: zip64_central_directory_offset,
                total_number_of_disks: 1,
            };
            self.write_all(&zip64_end_locator.build())?;
        }

        let end = spec::CentralDirectoryEnd {
            disk_number: 0,
            disk_with_central_directory: 0,
            disk_number_of_records: self.files.len() as u16,
            total_number_of_records: self.files.len() as u16,
            central_directory_size: central_directory_size.min(spec::ZIP64_THRESHOLD_BYTES) as u32,
            central_directory_offset: central_directory_offset.min(spec::ZIP64_THRESHOLD_BYTES)
                as u32,
            comment: Vec::with_capacity(0),
        };
        self.write_all(&end.build())?;
        self.flush()?;
        Ok(self.size)
    }
}

impl<W: Write> Write for ZipArchive<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let written = self.inner.write(buf)?;
        if let ZipArchiveState::FileOpen(f) = &mut self.state {
            f.hasher.update(&buf[..written]);
            f.size += written as u64;
        }
        self.size += written as u64;
        Ok(written)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.inner.flush()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_zip() -> anyhow::Result<()> {
        use std::io::Cursor;
        let mut zip = ZipArchive::new(Cursor::new(Vec::new()));
        let opts = Default::default();
        zip.add("test.txt", &opts)?;
        zip.write_all(b"First file")?;
        zip.add("test2.txt", &opts)?;
        zip.write_all(b"Second file")?;
        zip.finish()?;
        Ok(())
    }
}