Skip to content

Commit

Permalink
use indexmap
Browse files Browse the repository at this point in the history
  • Loading branch information
cosmicexplorer committed May 2, 2024
1 parent 01a8ff4 commit 7a5348b
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 48 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ chrono = { version = "0.4.38", optional = true }
constant_time_eq = { version = "0.3.0", optional = true }
crc32fast = "1.4.0"
flate2 = { version = "1.0.28", default-features = false, optional = true }
indexmap = "2"
hmac = { version = "0.12.1", optional = true, features = ["reset"] }
pbkdf2 = { version = "0.12.2", optional = true }
sha1 = { version = "0.10.6", optional = true }
Expand Down
49 changes: 22 additions & 27 deletions src/read.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,8 @@ use crate::spec;
use crate::types::{AesMode, AesVendorVersion, DateTime, System, ZipFileData};
use crate::zipcrypto::{ZipCryptoReader, ZipCryptoReaderValid, ZipCryptoValidator};
use byteorder::{LittleEndian, ReadBytesExt};
use std::borrow::{Borrow, Cow};
use std::collections::HashMap;
use indexmap::IndexMap;
use std::borrow::Cow;
use std::io::{self, prelude::*};
use std::ops::Deref;
use std::path::{Path, PathBuf};
Expand Down Expand Up @@ -47,8 +47,7 @@ pub(crate) mod zip_archive {
/// Extract immutable data from `ZipArchive` to make it cheap to clone
#[derive(Debug)]
pub(crate) struct Shared {
pub(crate) files: Box<[super::ZipFileData]>,
pub(crate) names_map: super::HashMap<Box<str>, usize>,
pub(crate) files: super::IndexMap<Box<str>, super::ZipFileData>,
pub(super) offset: u64,
pub(super) dir_start: u64,
pub(super) dir_end: u64,
Expand Down Expand Up @@ -487,21 +486,18 @@ impl<R: Read + Seek> ZipArchive<R> {
} else {
dir_info.number_of_files
};
let mut files = Vec::with_capacity(file_capacity);
let mut names_map = HashMap::with_capacity(file_capacity);
let mut files = IndexMap::with_capacity(file_capacity);
reader.seek(io::SeekFrom::Start(dir_info.directory_start))?;
for _ in 0..dir_info.number_of_files {
let file = central_header_to_zip_file(reader, dir_info.archive_offset)?;
names_map.insert(file.file_name.clone(), files.len());
files.push(file);
files.insert(file.file_name.clone(), file);
}
let dir_end = reader.seek(io::SeekFrom::Start(dir_info.directory_start))?;
if dir_info.disk_number != dir_info.disk_with_central_directory {
unsupported_zip_error("Support for multi-disk files is not implemented")
} else {
Ok(Shared {
files: files.into(),
names_map,
files,
offset: dir_info.archive_offset,
dir_start: dir_info.directory_start,
dir_end,
Expand Down Expand Up @@ -606,7 +602,7 @@ impl<R: Read + Seek> ZipArchive<R> {

/// Returns an iterator over all the file and directory names in this archive.
pub fn file_names(&self) -> impl Iterator<Item = &str> {
self.shared.names_map.keys().map(Box::borrow)
self.shared.files.keys().map(|s| s.as_ref())
}

/// Search for a file entry by name, decrypt with given password
Expand Down Expand Up @@ -634,7 +630,7 @@ impl<R: Read + Seek> ZipArchive<R> {
/// Get the index of a file entry by name, if it's present.
#[inline(always)]
pub fn index_for_name(&self, name: &str) -> Option<usize> {
self.shared.names_map.get(name).copied()
self.shared.files.get_index_of(&*name)
}

/// Get the index of a file entry by path, if it's present.
Expand All @@ -648,16 +644,16 @@ impl<R: Read + Seek> ZipArchive<R> {
pub fn name_for_index(&self, index: usize) -> Option<&str> {
self.shared
.files
.get(index)
.map(|file_data| &*file_data.file_name)
.get_index(index)
.map(|(name, _)| name.as_ref())
}

fn by_name_with_optional_password<'a>(
&'a mut self,
name: &str,
password: Option<&[u8]>,
) -> ZipResult<ZipFile<'a>> {
let Some(index) = self.index_for_name(name) else {
let Some(index) = self.shared.files.get_index_of(name) else {
return Err(ZipError::FileNotFound);
};
self.by_index_with_optional_password(index, password)
Expand Down Expand Up @@ -692,28 +688,27 @@ impl<R: Read + Seek> ZipArchive<R> {
/// Get a contained file by index without decompressing it
pub fn by_index_raw(&mut self, file_number: usize) -> ZipResult<ZipFile<'_>> {
let reader = &mut self.reader;
self.shared
let (_, data) = self
.shared
.files
.get(file_number)
.ok_or(ZipError::FileNotFound)
.and_then(move |data| {
Ok(ZipFile {
crypto_reader: None,
reader: ZipFileReader::Raw(find_content(data, reader)?),
data: Cow::Borrowed(data),
})
})
.get_index(file_number)
.ok_or(ZipError::FileNotFound)?;
Ok(ZipFile {
crypto_reader: None,
reader: ZipFileReader::Raw(find_content(data, reader)?),
data: Cow::Borrowed(data),
})
}

fn by_index_with_optional_password(
&mut self,
file_number: usize,
mut password: Option<&[u8]>,
) -> ZipResult<ZipFile<'_>> {
let data = self
let (_, data) = self
.shared
.files
.get(file_number)
.get_index(file_number)
.ok_or(ZipError::FileNotFound)?;

match (password, data.encrypted) {
Expand Down
38 changes: 17 additions & 21 deletions src/write.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use byteorder::{LittleEndian, WriteBytesExt};
#[cfg(any(feature = "_deflate-any", feature = "bzip2", feature = "zstd",))]
use core::num::NonZeroU64;
use crc32fast::Hasher;
use std::collections::HashMap;
use indexmap::IndexMap;
use std::default::Default;
use std::io;
use std::io::prelude::*;
Expand Down Expand Up @@ -111,8 +111,7 @@ pub(crate) mod zip_writer {
/// ```
pub struct ZipWriter<W: Write + Seek> {
pub(super) inner: GenericZipWriter<W>,
pub(super) files: Vec<ZipFileData>,
pub(super) files_by_name: HashMap<Box<str>, usize>,
pub(super) files: IndexMap<Box<str>, ZipFileData>,
pub(super) stats: ZipWriterStats,
pub(super) writing_to_file: bool,
pub(super) writing_raw: bool,
Expand Down Expand Up @@ -435,7 +434,7 @@ impl<W: Write + Seek> Write for ZipWriter<W> {
if let Ok(count) = write_result {
self.stats.update(&buf[0..count]);
if self.stats.bytes_written > spec::ZIP64_BYTES_THR
&& !self.files.last_mut().unwrap().large_file
&& !self.files.last_mut().unwrap().1.large_file
{
self.abort_file().unwrap();
return Err(io::Error::new(
Expand Down Expand Up @@ -479,8 +478,7 @@ impl<A: Read + Write + Seek> ZipWriter<A> {

Ok(ZipWriter {
inner: Storer(MaybeEncrypted::Unencrypted(readwriter)),
files: metadata.files.into(),
files_by_name: metadata.names_map,
files: metadata.files,
stats: Default::default(),
writing_to_file: false,
comment: footer.zip_file_comment,
Expand Down Expand Up @@ -608,8 +606,7 @@ impl<W: Write + Seek> ZipWriter<W> {
pub fn new(inner: W) -> ZipWriter<W> {
ZipWriter {
inner: Storer(MaybeEncrypted::Unencrypted(inner)),
files: Vec::new(),
files_by_name: HashMap::new(),
files: IndexMap::new(),
stats: Default::default(),
writing_to_file: false,
writing_raw: false,
Expand Down Expand Up @@ -808,15 +805,12 @@ impl<W: Write + Seek> ZipWriter<W> {
}

fn insert_file_data(&mut self, file: ZipFileData) -> ZipResult<usize> {
let name = &file.file_name;
if self.files_by_name.contains_key(name) {
if self.files.contains_key(&file.file_name) {
return Err(InvalidArchive("Duplicate filename"));
}
let name = name.to_owned();
self.files.push(file);
let index = self.files.len() - 1;
self.files_by_name.insert(name, index);
Ok(index)
let name = file.file_name.to_owned();
self.files.insert(name.clone(), file);
Ok(self.files.get_index_of(&name).unwrap())
}

fn finish_file(&mut self) -> ZipResult<()> {
Expand All @@ -837,7 +831,7 @@ impl<W: Write + Seek> ZipWriter<W> {
if !self.writing_raw {
let file = match self.files.last_mut() {
None => return Ok(()),
Some(f) => f,
Some((_, f)) => f,
};
file.crc32 = self.stats.hasher.clone().finalize();
file.uncompressed_size = self.stats.bytes_written;
Expand Down Expand Up @@ -877,8 +871,7 @@ impl<W: Write + Seek> ZipWriter<W> {
/// Removes the file currently being written from the archive if there is one, or else removes
/// the file most recently written.
pub fn abort_file(&mut self) -> ZipResult<()> {
let last_file = self.files.pop().ok_or(ZipError::FileNotFound)?;
self.files_by_name.remove(&last_file.file_name);
let (_, last_file) = self.files.pop().ok_or(ZipError::FileNotFound)?;
let make_plain_writer = self.inner.prepare_next_writer(
Stored,
None,
Expand All @@ -891,7 +884,7 @@ impl<W: Write + Seek> ZipWriter<W> {
// overwrite a valid file and corrupt the archive
let rewind_safe: bool = match last_file.data_start.get() {
None => self.files.is_empty(),
Some(last_file_start) => self.files.iter().all(|file| {
Some(last_file_start) => self.files.values().all(|file| {
file.data_start
.get()
.is_some_and(|start| start < last_file_start)
Expand Down Expand Up @@ -1184,7 +1177,7 @@ impl<W: Write + Seek> ZipWriter<W> {
let writer = self.inner.get_plain();

let central_start = writer.stream_position()?;
for file in self.files.iter() {
for file in self.files.values() {
write_central_directory_header(writer, file)?;
}
let central_size = writer.stream_position()? - central_start;
Expand Down Expand Up @@ -1230,7 +1223,10 @@ impl<W: Write + Seek> ZipWriter<W> {
}

fn index_by_name(&self, name: &str) -> ZipResult<usize> {
Ok(*self.files_by_name.get(name).ok_or(ZipError::FileNotFound)?)
Ok(self
.files
.get_index_of(name)
.ok_or(ZipError::FileNotFound)?)
}

/// Adds another entry to the central directory referring to the same content as an existing
Expand Down

0 comments on commit 7a5348b

Please sign in to comment.