diff --git a/Cargo.toml b/Cargo.toml index fd63d79f..8b238560 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,8 @@ thiserror = "1.0" lazy_static = "1.4.0" tracing = "0.1" +stream-unpack = "1.0.1" + serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" diff --git a/src/archive/mod.rs b/src/archive/mod.rs index 40d086b1..770a5467 100644 --- a/src/archive/mod.rs +++ b/src/archive/mod.rs @@ -7,6 +7,8 @@ pub mod sevenz; pub mod zip; pub mod tar; +pub mod stream; + use crate::updater::UpdaterExt; use updater::BasicUpdater; diff --git a/src/archive/stream.rs b/src/archive/stream.rs new file mode 100644 index 00000000..a8c53f15 --- /dev/null +++ b/src/archive/stream.rs @@ -0,0 +1,358 @@ +use std::cell::{Cell, RefCell}; +use std::fs::OpenOptions; +use std::io::Write; +use std::os::unix::fs::FileExt; +use std::path::{Path, PathBuf}; +use std::thread::JoinHandle; + +use stream_unpack::zip::structures::CompressionMethod; +use stream_unpack::zip::structures::central_directory::CentralDirectory; +use stream_unpack::zip::{read_cd, ZipUnpacker, DecoderError, ZipDecodedData, ZipPosition}; + +use crate::updater::UpdaterExt; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("Failed to send message through the flume channel: {0}")] + FlumeSendError(#[from] flume::SendError), + + #[error("IO error: {0}")] + Io(#[from] std::io::Error), + + #[error("Failed to fetch data: {0}")] + Minreq(#[from] minreq::Error), + + #[error("No Content-Length in response")] + NoSize, + + #[error("Invalid Content-Length in response: {0}")] + InvalidSize(#[from] std::num::ParseIntError), + + #[error("Failed to process central directory: {0}")] + CentralDirectory(#[from] read_cd::CentralDirectoryReadError), + + #[error("This archive contains elements which currently cannot be stream unpacked")] + UnsupportedArchive, + + #[error("Could not map position to archive volumes")] + InvalidPosition, + + #[error("Decoder error: {0}")] + Decoder(#[from] DecoderError) +} + +struct StreamArchiveDisk { + uri: String, + size: usize +} + +impl StreamArchiveDisk { + fn from_uri(uri: impl AsRef) -> Result { + let uri = uri.as_ref(); + + let size = minreq::head(uri) + .send()?.headers.get("content-length") + .ok_or(Error::NoSize)? + .parse()?; + + Ok(Self { + uri: uri.to_owned(), + size + }) + } +} + +pub struct StreamArchive { + archives: Vec, + central_directory: CentralDirectory, + is_cut: bool +} + +impl StreamArchive { + pub fn from_uris(uris: &[&str], is_cut: bool) -> Result { + let archives = uris.iter() + .map(StreamArchiveDisk::from_uri) + .collect::, _>>()?; + + let central_directory = read_cd::from_provider( + archives.iter().map(|d| d.size).collect::>(), + is_cut, + |pos, length| { + let start = pos.offset; + let end = start + length - 1; + + let bytes = minreq::get(&archives[pos.disk].uri) + .with_header("range", format!("bytes={start}-{end}")) + .send()?.into_bytes(); + + Ok(bytes) + } + )?; + + // None - no compression - is supported + let is_supported = central_directory.headers_ref().iter() + .flat_map(|h| &h.compression_method) + .all(CompressionMethod::is_supported); + + if !is_supported { + return Err(Error::UnsupportedArchive); + } + + Ok(Self { + archives, + central_directory, + is_cut + }) + } + + pub fn from_uri(uri: &str) -> Result { + Self::from_uris(&[uri], false) + } + + pub fn total_size(&self) -> usize { + self.archives.iter() + .map(|a| a.size) + .sum() + } + + pub fn uncompressed_size(&self) -> usize { + self.central_directory.headers_ref().iter() + .map(|h| h.uncompressed_size as usize) + .sum() + } + + pub fn stream_unpack(self, folder: impl AsRef, status_file: impl AsRef) -> Result { + let folder = folder.as_ref().to_owned(); + let status_file = status_file.as_ref().to_owned(); + + let (send, recv) = flume::unbounded(); + + Ok(StreamArchiveUpdater { + incrementer: recv, + + current: Cell::new(0), + total: self.total_size(), + + status_updater_result: None, + + status_updater: Some(std::thread::spawn(move || -> Result<(), Error> { + let sorted_cd = self.central_directory.sort(); + let disk_sizes = self.archives.iter().map(|a| a.size).collect::>(); + let file = RefCell::new(None); + + let status_file_path = folder.join(status_file); + let status_file_exists = status_file_path.exists(); + let status_file = OpenOptions::new() + .create(true) + .read(true) + .write(true) + .open(&status_file_path)?; + + let virtual_disk_sizes = if !self.is_cut { + disk_sizes.clone() + } else { + vec![disk_sizes.iter().sum()] + }; + + let (pos, mut unpacker) = if !status_file_exists { + let pos = sorted_cd.headers_ref()[0].header_position(); + StreamArchiveUpdater::write_status_file(&status_file, pos)?; + + (pos, ZipUnpacker::new(sorted_cd, virtual_disk_sizes)) + } else { + let pos = StreamArchiveUpdater::read_status_file(&status_file)?; + + (pos, ZipUnpacker::resume(sorted_cd, virtual_disk_sizes, pos)?) + }; + + unpacker.set_callback(|data| { + match data { + ZipDecodedData::FileHeader(h, _) => { + let mut path = PathBuf::from(&folder); + path.push(&h.filename); + + if !h.is_directory() { + std::fs::create_dir_all(path.parent().unwrap())?; + + *file.borrow_mut() = Some( + OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(path)? + ); + } else { + std::fs::create_dir_all(path)?; + } + + StreamArchiveUpdater::write_status_file(&status_file, h.header_position())?; + }, + + ZipDecodedData::FileData(data) => { + file.borrow().as_ref().unwrap().write_all(data)?; + } + } + + Ok(()) + }); + + let start_pos = if !self.is_cut { + pos + } else { + Self::offset_to_position(&disk_sizes, pos.offset)? + }; + + if start_pos != ZipPosition::default() { + let completed = disk_sizes.iter().take(start_pos.disk).sum::() + start_pos.offset; + send.send(completed)?; + } + + let mut buf = Vec::with_capacity((1 << 16) + 4096); + for (i, archive) in self.archives.iter().enumerate().skip(start_pos.disk) { + let mut request = minreq::get(&archive.uri); + if start_pos.disk == i && start_pos.offset != 0 { + request = request.with_header("range", format!("bytes={}-", start_pos.offset)); + } + let response = request.send_lazy()?; + + for byte in response { + let (byte, _) = byte?; + + buf.push(byte); + + if buf.len() >= (1 << 16) { + let (advanced, reached_end) = unpacker.update(&buf)?; + + if advanced != 0 { + buf.drain(..advanced); + send.send(advanced)?; + } + + if reached_end { + break; + } + } + } + } + + if !buf.is_empty() { + unpacker.update(&buf)?; + send.send(buf.len())?; + } + + drop(unpacker); + + drop(status_file); + std::fs::remove_file(status_file_path)?; + + Ok(()) + })) + }) + } + + pub fn offset_to_position(disk_sizes: impl AsRef<[usize]>, offset: usize) -> Result { + let disk_sizes = disk_sizes.as_ref(); + + let mut left = offset; + for (i, size) in disk_sizes.iter().enumerate() { + if left < *size { + return Ok(ZipPosition::new(i, left)); + } else { + left -= *size; + } + } + + Err(Error::InvalidPosition) + } +} + +pub struct StreamArchiveUpdater { + status_updater: Option>>, + status_updater_result: Option>, + + incrementer: flume::Receiver, + + current: Cell, + total: usize +} + +impl StreamArchiveUpdater { + fn read_status_file(file: &std::fs::File) -> Result { + let mut bytes = [0u8; 12]; + file.read_exact_at(&mut bytes, 0)?; + + Ok(ZipPosition::new( + u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize, + u64::from_le_bytes(bytes[4..12].try_into().unwrap()) as usize + )) + } + + fn write_status_file(file: &std::fs::File, pos: ZipPosition) -> Result<(), std::io::Error> { + file.write_all_at(&((pos.disk as u32).to_le_bytes()), 0)?; + file.write_all_at(&((pos.offset as u64).to_le_bytes()), 4)?; + + Ok(()) + } +} + +impl UpdaterExt for StreamArchiveUpdater { + type Error = Error; + type Status = bool; + type Result = (); + + #[inline] + fn status(&mut self) -> Result { + if let Some(status_updater) = self.status_updater.take() { + if !status_updater.is_finished() { + self.status_updater = Some(status_updater); + + return Ok(false); + } + + self.status_updater_result = Some(status_updater.join().expect("Failed to join thread")); + } + + match &self.status_updater_result { + Some(Ok(_)) => Ok(true), + Some(Err(err)) => Err(err), + + None => unreachable!() + } + } + + #[inline] + fn wait(mut self) -> Result { + if let Some(worker) = self.status_updater.take() { + return worker.join().expect("Failed to join thread"); + } + + else if let Some(result) = self.status_updater_result.take() { + return result; + } + + unreachable!() + } + + #[inline] + fn is_finished(&mut self) -> bool { + matches!(self.status(), Ok(true) | Err(_)) + } + + #[inline] + fn current(&self) -> u64 { + let mut current = self.current.get(); + + while let Ok(increment) = self.incrementer.try_recv() { + current += increment; + } + + self.current.set(current); + + current as u64 + } + + #[inline] + fn total(&self) -> u64 { + self.total as u64 + } +}