use crate::detect::{sniff, Format};
use crate::error::{ArchiveError, Result};
use crate::peel::MAX_INFLATED;
use crate::plan::Access;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ArchiveEntry {
pub name: String,
pub size: u64,
pub is_dir: bool,
}
pub struct Archive {
format: Format,
entries: Vec<ArchiveEntry>,
backend: Backend,
}
#[allow(clippy::large_enum_variant)]
enum Backend {
Tar { compressed: Vec<u8>, outer: Format },
Zip {
archive: zip_core::ZipArchive<std::io::Cursor<Vec<u8>>>,
},
SevenZip {
reader: sevenz_rust2::ArchiveReader<std::io::Cursor<Vec<u8>>>,
},
}
impl Archive {
pub fn open(data: &[u8], name: Option<&str>) -> Result<Option<Archive>> {
Self::open_with_format(sniff(name, data), data)
}
pub(crate) fn open_with_format(format: Format, data: &[u8]) -> Result<Option<Archive>> {
match format {
Format::Tar | Format::TarGz | Format::TarBz2 => {
Self::open_tar(format, data.to_vec()).map(Some)
}
Format::Zip => Self::open_zip(data).map(Some),
Format::SevenZip => Self::open_7z(data).map(Some),
_ => Ok(None),
}
}
pub fn member_access(&mut self, index: usize) -> Result<Access> {
let count = self.entries.len();
if index >= count {
return Err(ArchiveError::IndexOutOfRange { index, count });
}
match &mut self.backend {
Backend::Zip { archive } => {
let f = archive.by_index(index).map_err(|e| ArchiveError::Read {
format: "zip",
detail: e.to_string(),
})?;
Ok(match f.compression() {
zip_core::CompressionMethod::Stored => Access::InPlace {
offset: f.data_start(),
len: f.compressed_size(),
},
zip_core::CompressionMethod::Deflated
| zip_core::CompressionMethod::Deflate64 => Access::Zran,
_ => Access::SpillToTemp,
})
}
Backend::Tar { .. } | Backend::SevenZip { .. } => Ok(Access::SpillToTemp),
}
}
#[must_use]
pub fn format(&self) -> Format {
self.format
}
#[must_use]
pub fn entries(&self) -> &[ArchiveEntry] {
&self.entries
}
pub fn read(&mut self, index: usize) -> Result<Vec<u8>> {
let count = self.entries.len();
if index >= count {
return Err(ArchiveError::IndexOutOfRange { index, count });
}
let (name, declared_size) = {
let e = &self.entries[index];
(e.name.clone(), e.size)
};
match &mut self.backend {
Backend::Tar { compressed, outer } => {
extract_tar_member_streaming(compressed, *outer, index, MAX_INFLATED)
}
Backend::Zip { archive } => read_zip_member(archive, index),
Backend::SevenZip { reader } => read_7z_member(reader, &name, declared_size),
}
}
pub fn stream_member(
&mut self,
index: usize,
out: &mut dyn std::io::Write,
cap: u64,
) -> Result<u64> {
let count = self.entries.len();
if index >= count {
return Err(ArchiveError::IndexOutOfRange { index, count });
}
let (name, declared_size) = {
let e = &self.entries[index];
(e.name.clone(), e.size)
};
match &mut self.backend {
Backend::Tar { compressed, outer } => {
stream_tar_member(compressed, *outer, index, out, cap)
}
Backend::Zip { archive } => stream_zip_member(archive, index, out, cap),
Backend::SevenZip { reader } => {
let bytes = read_7z_member(reader, &name, declared_size)?;
if bytes.len() as u64 > cap {
return Err(ArchiveError::TooLarge { cap });
}
out.write_all(&bytes).map_err(|e| ArchiveError::Read {
format: "7z",
detail: e.to_string(),
})?;
Ok(bytes.len() as u64)
}
}
}
fn open_tar(outer: Format, compressed: Vec<u8>) -> Result<Archive> {
let entries = {
let mut ar = tar::Archive::new(tar_stream(&compressed, outer));
let iter = ar.entries().map_err(|e| ArchiveError::Open {
format: "tar",
detail: e.to_string(),
})?;
let mut entries = Vec::new();
for entry in iter {
let entry = entry.map_err(|e| ArchiveError::Open {
format: "tar",
detail: e.to_string(),
})?;
let name = entry.path().map_or_else(
|_| String::from_utf8_lossy(&entry.path_bytes()).into_owned(),
|p| p.to_string_lossy().into_owned(),
);
let header = entry.header();
let size = header.size().map_err(|e| ArchiveError::Open {
format: "tar",
detail: e.to_string(),
})?;
entries.push(ArchiveEntry {
name,
size,
is_dir: header.entry_type().is_dir(),
});
}
entries
};
Ok(Archive {
format: outer,
entries,
backend: Backend::Tar { compressed, outer },
})
}
fn open_zip(data: &[u8]) -> Result<Archive> {
let mut archive =
zip_core::ZipArchive::new(std::io::Cursor::new(data.to_vec())).map_err(|e| {
ArchiveError::Open {
format: "zip",
detail: e.to_string(),
}
})?;
let count = archive.len();
let mut entries = Vec::with_capacity(count);
for i in 0..count {
let f = archive.by_index(i).map_err(|e| ArchiveError::Open {
format: "zip",
detail: e.to_string(),
})?;
entries.push(ArchiveEntry {
name: f.name().to_string(),
size: f.size(),
is_dir: f.is_dir(),
});
}
Ok(Archive {
format: Format::Zip,
entries,
backend: Backend::Zip { archive },
})
}
fn open_7z(data: &[u8]) -> Result<Archive> {
let reader = sevenz_rust2::ArchiveReader::new(
std::io::Cursor::new(data.to_vec()),
sevenz_rust2::Password::empty(),
)
.map_err(|e| ArchiveError::Open {
format: "7z",
detail: e.to_string(),
})?;
let entries = reader
.archive()
.files
.iter()
.map(|f| ArchiveEntry {
name: f.name.clone(),
size: f.size,
is_dir: f.is_directory,
})
.collect();
Ok(Archive {
format: Format::SevenZip,
entries,
backend: Backend::SevenZip { reader },
})
}
}
fn tar_stream(compressed: &[u8], outer: Format) -> Box<dyn std::io::Read + '_> {
let cursor = std::io::Cursor::new(compressed);
match outer {
Format::TarGz => Box::new(flate2::read::GzDecoder::new(cursor)),
Format::TarBz2 => Box::new(bzip2_rs::DecoderReader::new(cursor)),
_ => Box::new(cursor),
}
}
fn extract_tar_member_streaming(
compressed: &[u8],
outer: Format,
index: usize,
cap: u64,
) -> Result<Vec<u8>> {
use std::io::Read;
let mut ar = tar::Archive::new(tar_stream(compressed, outer));
let mut iter = ar.entries().map_err(|e| ArchiveError::Read {
format: "tar",
detail: e.to_string(),
})?;
let entry = iter
.nth(index)
.ok_or(ArchiveError::IndexOutOfRange {
index,
count: index,
})?
.map_err(|e| ArchiveError::Read {
format: "tar",
detail: e.to_string(),
})?;
let mut out = Vec::new();
let mut limited = entry.take(cap + 1);
limited
.read_to_end(&mut out)
.map_err(|e| ArchiveError::Read {
format: "tar",
detail: e.to_string(),
})?;
if out.len() as u64 > cap {
return Err(ArchiveError::TooLarge { cap });
}
Ok(out)
}
fn read_zip_member(
archive: &mut zip_core::ZipArchive<std::io::Cursor<Vec<u8>>>,
index: usize,
) -> Result<Vec<u8>> {
use std::io::Read;
let zf = archive.by_index(index).map_err(|e| ArchiveError::Read {
format: "zip",
detail: e.to_string(),
})?;
let mut out = Vec::new();
let mut limited = zf.take(MAX_INFLATED + 1);
limited
.read_to_end(&mut out)
.map_err(|e| ArchiveError::Read {
format: "zip",
detail: e.to_string(),
})?;
if out.len() as u64 > MAX_INFLATED {
return Err(ArchiveError::TooLarge { cap: MAX_INFLATED });
}
Ok(out)
}
fn read_7z_member(
reader: &mut sevenz_rust2::ArchiveReader<std::io::Cursor<Vec<u8>>>,
name: &str,
declared_size: u64,
) -> Result<Vec<u8>> {
if declared_size > MAX_INFLATED {
return Err(ArchiveError::TooLarge { cap: MAX_INFLATED });
}
let out = reader.read_file(name).map_err(|e| ArchiveError::Read {
format: "7z",
detail: e.to_string(),
})?;
if out.len() as u64 > MAX_INFLATED {
return Err(ArchiveError::TooLarge { cap: MAX_INFLATED });
}
Ok(out)
}
fn copy_capped(
reader: impl std::io::Read,
out: &mut dyn std::io::Write,
cap: u64,
format: &'static str,
) -> Result<u64> {
let mut limited = reader.take(cap + 1);
let n = std::io::copy(&mut limited, out).map_err(|e| ArchiveError::Read {
format,
detail: e.to_string(),
})?;
if n > cap {
return Err(ArchiveError::TooLarge { cap });
}
Ok(n)
}
fn stream_tar_member(
compressed: &[u8],
outer: Format,
index: usize,
out: &mut dyn std::io::Write,
cap: u64,
) -> Result<u64> {
let mut ar = tar::Archive::new(tar_stream(compressed, outer));
let mut iter = ar.entries().map_err(|e| ArchiveError::Read {
format: "tar",
detail: e.to_string(),
})?;
let entry = iter
.nth(index)
.ok_or(ArchiveError::IndexOutOfRange {
index,
count: index,
})?
.map_err(|e| ArchiveError::Read {
format: "tar",
detail: e.to_string(),
})?;
copy_capped(entry, out, cap, "tar")
}
fn stream_zip_member(
archive: &mut zip_core::ZipArchive<std::io::Cursor<Vec<u8>>>,
index: usize,
out: &mut dyn std::io::Write,
cap: u64,
) -> Result<u64> {
let zf = archive.by_index(index).map_err(|e| ArchiveError::Read {
format: "zip",
detail: e.to_string(),
})?;
copy_capped(zf, out, cap, "zip")
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
fn build_tar(members: &[(&str, Vec<u8>)]) -> Vec<u8> {
let mut b = tar::Builder::new(Vec::new());
for (name, data) in members {
let mut h = tar::Header::new_gnu();
h.set_size(data.len() as u64);
h.set_mode(0o644);
h.set_cksum();
b.append_data(&mut h, name, data.as_slice()).unwrap();
}
b.into_inner().unwrap()
}
fn gzip(data: &[u8]) -> Vec<u8> {
let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
e.write_all(data).unwrap();
e.finish().unwrap()
}
#[test]
fn streaming_targz_extracts_each_member_under_whole_tar_cap() {
let a = vec![0xAA_u8; 1000];
let b = vec![0xBB_u8; 1000];
let tar = build_tar(&[("a.bin", a.clone()), ("b.bin", b.clone())]);
assert!(
tar.len() as u64 > 1200,
"combined tar ({}) must exceed the per-member cap",
tar.len()
);
let targz = gzip(&tar);
assert_eq!(
extract_tar_member_streaming(&targz, Format::TarGz, 0, 1200).unwrap(),
a
);
assert_eq!(
extract_tar_member_streaming(&targz, Format::TarGz, 1, 1200).unwrap(),
b
);
}
#[test]
fn streaming_plain_tar_extracts_member() {
let a = vec![0x11_u8; 800];
let b = vec![0x22_u8; 800];
let tar = build_tar(&[("a.bin", a.clone()), ("b.bin", b.clone())]);
assert_eq!(
extract_tar_member_streaming(&tar, Format::Tar, 0, MAX_INFLATED).unwrap(),
a
);
assert_eq!(
extract_tar_member_streaming(&tar, Format::Tar, 1, MAX_INFLATED).unwrap(),
b
);
}
#[test]
fn streaming_member_over_cap_fails_loud() {
let targz = gzip(&build_tar(&[("a.bin", vec![0xAA_u8; 1000])]));
match extract_tar_member_streaming(&targz, Format::TarGz, 0, 500) {
Err(ArchiveError::TooLarge { cap }) => assert_eq!(cap, 500),
other => panic!("expected TooLarge, got {other:?}"),
}
}
#[test]
fn streaming_bad_index_fails_loud() {
let tar = build_tar(&[("only.bin", vec![0x33_u8; 10])]);
assert!(matches!(
extract_tar_member_streaming(&tar, Format::Tar, 99, MAX_INFLATED),
Err(ArchiveError::IndexOutOfRange { .. })
));
}
const PAYLOAD_TBZ2: &[u8] = include_bytes!("../../tests/data/fixtures/payload.tbz2");
#[test]
fn member_access_out_of_range_fails_loud() {
let mut a = Archive::open(PAYLOAD_TBZ2, Some("payload.tbz2"))
.unwrap()
.unwrap();
assert!(matches!(
a.member_access(9999),
Err(ArchiveError::IndexOutOfRange { .. })
));
}
#[test]
fn streaming_tbz2_reads_member_via_same_path() {
let mut a = Archive::open(PAYLOAD_TBZ2, Some("payload.tbz2"))
.unwrap()
.unwrap();
assert_eq!(a.format(), Format::TarBz2);
let ia = a
.entries()
.iter()
.position(|e| e.name == "a.txt" && !e.is_dir)
.unwrap();
assert_eq!(a.read(ia).unwrap(), b"alpha member contents\n");
}
}