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, 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(),
permissions: 0o100644,
}
}
}
enum ZipArchiveState {
Start,
FileOpen(ZipFile),
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(())
}
}