diff --git a/native/src/boot/compress.rs b/native/src/boot/compress.rs index 9af777489..0d66400cc 100644 --- a/native/src/boot/compress.rs +++ b/native/src/boot/compress.rs @@ -1,24 +1,26 @@ use crate::ffi::{FileFormat, check_fmt}; use base::libc::{O_RDONLY, O_TRUNC, O_WRONLY}; -use base::{Chunker, FileOrStd, LoggedResult, Utf8CStr, Utf8CString, WriteExt, log_err}; -use bytemuck::bytes_of_mut; -use bzip2::{Compression as BzCompression, write::BzDecoder, write::BzEncoder}; -use flate2::{Compression as GzCompression, write::GzEncoder, write::MultiGzDecoder}; +use base::{ + Chunker, FileOrStd, LoggedResult, ReadExt, ResultExt, Utf8CStr, Utf8CString, WriteExt, log_err, +}; +use bzip2::{Compression as BzCompression, read::BzDecoder, write::BzEncoder}; +use flate2::{Compression as GzCompression, read::MultiGzDecoder, write::GzEncoder}; use lz4::{ - BlockMode, BlockSize, ContentChecksum, Encoder as LZ4FrameEncoder, + BlockMode, BlockSize, ContentChecksum, Decoder as LZ4FrameDecoder, Encoder as LZ4FrameEncoder, EncoderBuilder as LZ4FrameEncoderBuilder, block::CompressionMode, liblz4::BlockChecksum, }; -use std::cell::Cell; +use std::cmp::min; use std::fmt::Write as FmtWrite; use std::fs::File; -use std::io::{BufWriter, Read, Write}; +use std::io::{BufWriter, Cursor, Read, Write}; use std::mem::ManuallyDrop; use std::num::NonZeroU64; use std::ops::DerefMut; use std::os::fd::{FromRawFd, RawFd}; use xz2::{ + read::XzDecoder, stream::{Check as LzmaCheck, Filters as LzmaFilters, LzmaOptions, Stream as LzmaStream}, - write::{XzDecoder, XzEncoder}, + write::XzEncoder, }; use zopfli::{BlockType, GzipEncoder as ZopFliEncoder, Options as ZopfliOptions}; @@ -38,19 +40,7 @@ macro_rules! finish_impl { )*} } -finish_impl!(GzEncoder, MultiGzDecoder, BzEncoder, XzEncoder); - -macro_rules! finish_impl_ref { - ($($t:ty),*) => {$( - impl WriteFinish for $t { - fn finish(mut self: Box) -> std::io::Result { - Self::finish(self.as_mut()) - } - } - )*} -} - -finish_impl_ref!(BzDecoder, XzDecoder); +finish_impl!(GzEncoder, BzEncoder, XzEncoder); impl WriteFinish for BufWriter> { fn finish(self: Box) -> std::io::Result { @@ -67,89 +57,6 @@ impl WriteFinish for LZ4FrameEncoder { } } -// Adapt Reader to Writer - -// In case some decoders don't support the Write trait, instead of pushing data into the -// decoder, we have no choice but to pull data out of it. So first, we create a "fake" reader -// that does not own any data as a placeholder. In the Writer adapter struct, when data -// is fed in, we call FakeReader::set_data to forward this data as the "source" of the -// decoder. Next, we pull data out of the decoder, and finally, forward the decoded data to output. - -struct FakeReader(Cell<&'static [u8]>); - -impl FakeReader { - fn new() -> FakeReader { - FakeReader(Cell::new(&[])) - } - - // SAFETY: the lifetime of the buffer is between the invocation of - // this method and the invocation of FakeReader::clear. There is currently - // no way to represent this with Rust's lifetime marker, so we transmute all - // lifetimes away and make the users of this struct manually manage the lifetime. - // It is the responsibility of the caller to ensure the underlying reference does not - // live longer than it should. - unsafe fn set_data(&self, data: &[u8]) { - let buf: &'static [u8] = unsafe { std::mem::transmute(data) }; - self.0.set(buf) - } - - fn clear(&self) { - self.0.set(&[]) - } -} - -impl Read for FakeReader { - fn read(&mut self, buf: &mut [u8]) -> std::io::Result { - let data = self.0.get(); - let len = std::cmp::min(buf.len(), data.len()); - buf[..len].copy_from_slice(&data[..len]); - self.0.set(&data[len..]); - Ok(len) - } -} - -// LZ4FrameDecoder - -struct LZ4FrameDecoder { - write: W, - decoder: lz4::Decoder, -} - -impl LZ4FrameDecoder { - fn new(write: W) -> Self { - let fake = FakeReader::new(); - let decoder = lz4::Decoder::new(fake).unwrap(); - LZ4FrameDecoder { write, decoder } - } -} - -impl Write for LZ4FrameDecoder { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.write_all(buf)?; - Ok(buf.len()) - } - - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } - - fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { - // SAFETY: buf is removed from the reader immediately after usage - unsafe { self.decoder.reader().set_data(buf) }; - std::io::copy(&mut self.decoder, &mut self.write)?; - self.decoder.reader().clear(); - Ok(()) - } -} - -impl WriteFinish for LZ4FrameDecoder { - fn finish(self: Box) -> std::io::Result { - let (_, r) = self.decoder.finish(); - r?; - Ok(self.write) - } -} - // LZ4BlockArchive format // // len: | 4 | 4 | n | ... | 4 | @@ -159,7 +66,7 @@ impl WriteFinish for LZ4FrameDecoder { const LZ4_BLOCK_SIZE: usize = 0x800000; const LZ4HC_CLEVEL_MAX: i32 = 12; -const LZ4_MAGIC: &[u8] = b"\x02\x21\x4c\x18"; +const LZ4_MAGIC: u32 = 0x184c2102; struct LZ4BlockEncoder { write: W, @@ -208,7 +115,7 @@ impl Write for LZ4BlockEncoder { fn write_all(&mut self, mut buf: &[u8]) -> std::io::Result<()> { if self.total == 0 { // Write header - self.write.write_all(LZ4_MAGIC)?; + self.write.write_pod(&LZ4_MAGIC)?; } self.total += buf.len() as u32; @@ -238,88 +145,72 @@ impl WriteFinish for LZ4BlockEncoder { // LZ4BlockDecoder -struct LZ4BlockDecoder { - write: W, - chunker: Chunker, +struct LZ4BlockDecoder { + read: R, + in_buf: Box<[u8]>, out_buf: Box<[u8]>, - curr_block_size: usize, + out_len: usize, + out_pos: usize, } -impl LZ4BlockDecoder { - fn new(write: W) -> Self { - LZ4BlockDecoder { - write, - chunker: Chunker::new(size_of::()), - // SAFETY: all bytes will be initialized before it is used +impl LZ4BlockDecoder { + fn new(read: R) -> Self { + let compressed_sz = lz4::block::compress_bound(LZ4_BLOCK_SIZE).unwrap_or(LZ4_BLOCK_SIZE); + Self { + read, + in_buf: unsafe { Box::new_uninit_slice(compressed_sz).assume_init() }, out_buf: unsafe { Box::new_uninit_slice(LZ4_BLOCK_SIZE).assume_init() }, - curr_block_size: 0, + out_len: 0, + out_pos: 0, } } - - fn decode_block(write: &mut W, out_buf: &mut [u8], chunk: &[u8]) -> std::io::Result<()> { - let decompressed_size = - lz4::block::decompress_to_buffer(chunk, Some(LZ4_BLOCK_SIZE as i32), out_buf)?; - write.write_all(&out_buf[..decompressed_size]) - } } -impl Write for LZ4BlockDecoder { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.write_all(buf)?; - Ok(buf.len()) - } - - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } - - fn write_all(&mut self, mut buf: &[u8]) -> std::io::Result<()> { - while !buf.is_empty() { - let (b, chunk) = self.chunker.add_data(buf); - buf = b; - if let Some(chunk) = chunk { - if chunk == LZ4_MAGIC { - // Skip magic, read next u32 - continue; - } - if self.curr_block_size == 0 { - // We haven't got the current block size yet, try read it - let mut next_u32: u32 = 0; - bytes_of_mut(&mut next_u32).copy_from_slice(chunk); - - if next_u32 > lz4::block::compress_bound(LZ4_BLOCK_SIZE)? as u32 { - // This is the LG format trailer, EOF - continue; - } - - // Update chunker to read next block - self.curr_block_size = next_u32 as usize; - self.chunker.set_chunk_size(self.curr_block_size); - continue; - } - - // Actually decode the block - Self::decode_block(&mut self.write, &mut self.out_buf, chunk)?; - - // Reset for the next block - self.curr_block_size = 0; - self.chunker.set_chunk_size(size_of::()); +impl Read for LZ4BlockDecoder { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + if self.out_pos == self.out_len { + let mut block_size: u32 = 0; + if let Err(e) = self.read.read_pod(&mut block_size) { + return if e.kind() == std::io::ErrorKind::UnexpectedEof { + Ok(0) + } else { + Err(e) + }; + } + if block_size == LZ4_MAGIC { + self.read.read_pod(&mut block_size)?; } - } - Ok(()) - } -} -impl WriteFinish for LZ4BlockDecoder { - fn finish(mut self: Box) -> std::io::Result { - let chunk = self.chunker.get_available(); - if !chunk.is_empty() { - return Err(std::io::Error::new( - std::io::ErrorKind::Interrupted, - "Finish ran before end of compressed stream", - )); + let block_size = block_size as usize; + + if block_size > self.in_buf.len() { + // This may be the LG format trailer, EOF + return Ok(0); + } + + // Read the entire compressed block + let compressed_block = &mut self.in_buf[..block_size]; + if let Ok(len) = self.read.read(compressed_block) { + if len == 0 { + // We hit EOF, that's fine + return Ok(0); + } else if len != block_size { + let remain = &mut compressed_block[len..]; + self.read.read_exact(remain)?; + } + } + + self.out_len = lz4::block::decompress_to_buffer( + compressed_block, + Some(LZ4_BLOCK_SIZE as i32), + &mut self.out_buf, + )?; + self.out_pos = 0; } - Ok(self.write) + let copy_len = min(buf.len(), self.out_len - self.out_pos); + buf[..copy_len].copy_from_slice(&self.out_buf[self.out_pos..self.out_pos + copy_len]); + self.out_pos += copy_len; + Ok(copy_len) } } @@ -368,16 +259,16 @@ pub fn get_encoder<'a, W: Write + 'a>(format: FileFormat, w: W) -> Box(format: FileFormat, w: W) -> Box + 'a> { +pub fn get_decoder<'a, R: Read + 'a>(format: FileFormat, r: R) -> Box { match format { FileFormat::XZ | FileFormat::LZMA => { let stream = LzmaStream::new_auto_decoder(u64::MAX, 0).unwrap(); - Box::new(XzDecoder::new_stream(w, stream)) + Box::new(XzDecoder::new_stream(r, stream)) } - FileFormat::BZIP2 => Box::new(BzDecoder::new(w)), - FileFormat::LZ4 => Box::new(LZ4FrameDecoder::new(w)), - FileFormat::LZ4_LG | FileFormat::LZ4_LEGACY => Box::new(LZ4BlockDecoder::new(w)), - FileFormat::ZOPFLI | FileFormat::GZIP => Box::new(MultiGzDecoder::new(w)), + FileFormat::BZIP2 => Box::new(BzDecoder::new(r)), + FileFormat::LZ4 => Box::new(LZ4FrameDecoder::new(r).unwrap()), + FileFormat::LZ4_LG | FileFormat::LZ4_LEGACY => Box::new(LZ4BlockDecoder::new(r)), + FileFormat::ZOPFLI | FileFormat::GZIP => Box::new(MultiGzDecoder::new(r)), _ => unreachable!(), } } @@ -397,11 +288,8 @@ pub fn compress_bytes(format: FileFormat, in_bytes: &[u8], out_fd: RawFd) { pub fn decompress_bytes(format: FileFormat, in_bytes: &[u8], out_fd: RawFd) { let mut out_file = unsafe { ManuallyDrop::new(File::from_raw_fd(out_fd)) }; - let mut decoder = get_decoder(format, out_file.deref_mut()); - let _: LoggedResult<()> = try { - decoder.write_all(in_bytes)?; - decoder.finish()?; - }; + let mut decoder = get_decoder(format, in_bytes); + std::io::copy(decoder.as_mut(), out_file.deref_mut()).log_ok(); } // Command-line entry points @@ -463,10 +351,8 @@ pub(crate) fn decompress_cmd( FileOrStd::File(outfile.create(O_WRONLY | O_TRUNC, 0o644)?) }; - let mut decoder = get_decoder(format, output.as_file()); - decoder.write_all(buf)?; - std::io::copy(&mut input.as_file(), decoder.as_mut())?; - decoder.finish()?; + let mut decoder = get_decoder(format, Cursor::new(buf).chain(input.as_file())); + std::io::copy(decoder.as_mut(), &mut output.as_file())?; if rm_in { infile.remove()?; diff --git a/native/src/boot/cpio.rs b/native/src/boot/cpio.rs index 5cdeb2cf5..6d2266996 100644 --- a/native/src/boot/cpio.rs +++ b/native/src/boot/cpio.rs @@ -4,7 +4,7 @@ use std::cmp::Ordering; use std::collections::{BTreeMap, HashMap}; use std::fmt::{Display, Formatter}; use std::fs::File; -use std::io::{Read, Write}; +use std::io::{Cursor, Read, Write}; use std::mem::size_of; use std::process::exit; use std::str; @@ -705,10 +705,11 @@ impl CpioEntry { return false; } - let mut decoder = get_decoder(FileFormat::XZ, Vec::new()); let Ok(data): std::io::Result> = (try { - decoder.write_all(&self.data)?; - decoder.finish()? + let mut decoder = get_decoder(FileFormat::XZ, Cursor::new(&self.data)); + let mut data = Vec::new(); + std::io::copy(decoder.as_mut(), &mut data)?; + data }) else { eprintln!("xz compression failed"); return false; diff --git a/native/src/boot/payload.rs b/native/src/boot/payload.rs index ff4d25b17..b9c675582 100644 --- a/native/src/boot/payload.rs +++ b/native/src/boot/payload.rs @@ -1,16 +1,16 @@ +use crate::compress::get_decoder; +use crate::ffi::check_fmt; +use crate::proto::update_metadata::{DeltaArchiveManifest, mod_InstallOperation::Type}; +use base::{LoggedError, LoggedResult, ReadSeekExt, ResultExt, WriteExt, error}; use byteorder::{BigEndian, ReadBytesExt}; use quick_protobuf::{BytesReader, MessageRead}; +use std::io::Cursor; use std::{ fs::File, io::{BufReader, Read, Seek, SeekFrom, Write}, os::fd::FromRawFd, }; -use crate::compress::get_decoder; -use crate::ffi::check_fmt; -use crate::proto::update_metadata::{DeltaArchiveManifest, mod_InstallOperation::Type}; -use base::{LoggedError, LoggedResult, ReadSeekExt, ResultExt, WriteExt, error}; - macro_rules! bad_payload { ($msg:literal) => {{ error!(concat!("Invalid payload: ", $msg)); @@ -164,9 +164,11 @@ pub fn extract_boot_from_payload( } Type::REPLACE_BZ | Type::REPLACE_XZ => { out_file.seek(SeekFrom::Start(out_offset))?; - let mut decoder = get_decoder(check_fmt(data), &mut out_file); + let fmt = check_fmt(data); + + let mut decoder = get_decoder(fmt, Cursor::new(data)); let Ok(_): std::io::Result<()> = (try { - decoder.write_all(data)?; + std::io::copy(decoder.as_mut(), &mut out_file)?; }) else { return Err(bad_payload!("decompression failed")); };