From 793b3b72d26178c422c213913d87f01c89cf155e Mon Sep 17 00:00:00 2001 From: wcampbell Date: Wed, 18 Oct 2023 00:26:42 -0400 Subject: [PATCH] Use new deku writer --- Cargo.lock | 4 +-- Cargo.toml | 2 +- src/bin/unsquashfs.rs | 2 +- src/filesystem/writer.rs | 58 +++++++++++++++++++++------------------- src/inode.rs | 10 +++---- src/metadata.rs | 20 +++++++------- 6 files changed, 49 insertions(+), 47 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 64c50366..2fb0d33f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -468,7 +468,7 @@ dependencies = [ [[package]] name = "deku" version = "0.16.0" -source = "git+https://github.com/sharksforarms/deku?branch=impl-reader#b32ef5d4374a790a731767367ff594bcbdd0d203" +source = "git+https://github.com/sharksforarms/deku?branch=impl-writer#91bb702b70231f0f8f7ea2ee0cbfc61c46a659de" dependencies = [ "bitvec", "deku_derive", @@ -478,7 +478,7 @@ dependencies = [ [[package]] name = "deku_derive" version = "0.16.0" -source = "git+https://github.com/sharksforarms/deku?branch=impl-reader#b32ef5d4374a790a731767367ff594bcbdd0d203" +source = "git+https://github.com/sharksforarms/deku?branch=impl-writer#91bb702b70231f0f8f7ea2ee0cbfc61c46a659de" dependencies = [ "darling", "proc-macro-crate", diff --git a/Cargo.toml b/Cargo.toml index 2b6279d4..b9e073ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ categories = ["filesystem", "parsing"] [dependencies] # for lib -deku = { git = "https://github.com/sharksforarms/deku", branch = "impl-reader" } +deku = { git = "https://github.com/sharksforarms/deku", branch = "impl-writer" } tracing = "0.1.37" thiserror = "1.0.37" flate2 = { version = "1.0.24", optional = true } diff --git a/src/bin/unsquashfs.rs b/src/bin/unsquashfs.rs index 096b92b9..366f8fdd 100644 --- a/src/bin/unsquashfs.rs +++ b/src/bin/unsquashfs.rs @@ -175,7 +175,7 @@ fn main() -> ExitCode { let mut file = BufReader::new(File::open(args.filesystem.as_ref().unwrap()).unwrap()); let blue_bold: console::Style = console::Style::new().blue().bold(); - let red_bold: console::Style = console::Style::new().blue().bold(); + let red_bold: console::Style = console::Style::new().red().bold(); let pb = ProgressBar::new_spinner(); if args.auto_offset { diff --git a/src/filesystem/writer.rs b/src/filesystem/writer.rs index dd5a015d..4132c56e 100644 --- a/src/filesystem/writer.rs +++ b/src/filesystem/writer.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use std::sync::Mutex; use std::time::{SystemTime, UNIX_EPOCH}; -use deku::bitvec::BitVec; -use deku::DekuWrite; +use deku::writer::Writer; +use deku::DekuWriter; use tracing::{error, info, instrument, trace}; use super::node::{InnerNode, Nodes}; @@ -573,12 +573,11 @@ impl<'a, 'b> FilesystemWriter<'a, 'b> { trace!("WRITING DIR: {block_offset:#02x?}"); let mut total_size: usize = 3; for dir in Entry::into_dir(entries) { - let mut bv = BitVec::new(); - dir.write(&mut bv, kind.inner.type_endian)?; - let bytes = bv.as_raw_slice(); - dir_writer.write_all(bv.as_raw_slice())?; - + let mut bytes = vec![]; + let mut writer = Writer::new(&mut bytes); + dir.to_writer(&mut writer, kind.inner.type_endian)?; total_size += bytes.len(); + dir_writer.write_all(&bytes)?; } let entry = Entry::path( filename, @@ -622,17 +621,24 @@ impl<'a, 'b> FilesystemWriter<'a, 'b> { // Write compression options, if any if let Some(options) = &self.fs_compressor.options { superblock.flags |= Flags::CompressorOptionsArePresent as u16; - let mut buf = BitVec::new(); + let mut compression_opt_buf_out = vec![]; + let mut writer = Writer::new(&mut compression_opt_buf_out); match options { CompressionOptions::Gzip(gzip) => { - gzip.write(&mut buf, self.kind.inner.type_endian)? + gzip.to_writer(&mut writer, self.kind.inner.type_endian)? + } + CompressionOptions::Lz4(lz4) => { + lz4.to_writer(&mut writer, self.kind.inner.type_endian)? } - CompressionOptions::Lz4(lz4) => lz4.write(&mut buf, self.kind.inner.type_endian)?, CompressionOptions::Zstd(zstd) => { - zstd.write(&mut buf, self.kind.inner.type_endian)? + zstd.to_writer(&mut writer, self.kind.inner.type_endian)? + } + CompressionOptions::Xz(xz) => { + xz.to_writer(&mut writer, self.kind.inner.type_endian)? + } + CompressionOptions::Lzo(lzo) => { + lzo.to_writer(&mut writer, self.kind.inner.type_endian)? } - CompressionOptions::Xz(xz) => xz.write(&mut buf, self.kind.inner.type_endian)?, - CompressionOptions::Lzo(lzo) => lzo.write(&mut buf, self.kind.inner.type_endian)?, CompressionOptions::Lzma => {} } let mut metadata = MetadataWriter::new( @@ -642,7 +648,7 @@ impl<'a, 'b> FilesystemWriter<'a, 'b> { inner: self.kind.inner.clone(), }, ); - metadata.write_all(buf.as_raw_slice())?; + metadata.write_all(&compression_opt_buf_out)?; metadata.finalize(w)?; } @@ -754,9 +760,9 @@ impl<'a, 'b> FilesystemWriter<'a, 'b> { // Seek back the beginning and write the superblock info!("Writing Superblock"); w.rewind()?; - let mut bv = BitVec::new(); - superblock.write( - &mut bv, + let mut writer = Writer::new(w); + superblock.to_writer( + &mut writer, ( self.kind.inner.magic, self.kind.inner.version_major, @@ -764,7 +770,6 @@ impl<'a, 'b> FilesystemWriter<'a, 'b> { self.kind.inner.type_endian, ), )?; - w.write_all(bv.as_raw_slice())?; info!("Writing Finished"); @@ -800,7 +805,7 @@ impl<'a, 'b> FilesystemWriter<'a, 'b> { /// │└────────────────────────────┘│ /// └──────────────────────────────┘ /// ``` - fn write_lookup_table, W: Write + Seek>( + fn write_lookup_table, W: Write + Seek>( &self, w: &mut W, table: &Vec, @@ -810,20 +815,18 @@ impl<'a, 'b> FilesystemWriter<'a, 'b> { let mut table_bytes = Vec::with_capacity(table.len() * element_size); let mut iter = table.iter().peekable(); while let Some(t) = iter.next() { + let mut table_writer = Writer::new(&mut table_bytes); // convert fragment ptr to bytes - let mut bv = BitVec::new(); - t.write(&mut bv, self.kind.inner.type_endian)?; - table_bytes.write_all(bv.as_raw_slice())?; + t.to_writer(&mut table_writer, self.kind.inner.type_endian)?; // once table_bytes + next is over the maximum size of a metadata block, write if ((table_bytes.len() + element_size) > METADATA_MAXSIZE) || iter.peek().is_none() { ptrs.push(w.stream_position()?); - let mut bv = BitVec::new(); // write metadata len let len = metadata::set_if_uncompressed(table_bytes.len() as u16); - len.write(&mut bv, self.kind.inner.data_endian)?; - w.write_all(bv.as_raw_slice())?; + let mut writer = Writer::new(w); + len.to_writer(&mut writer, self.kind.inner.data_endian)?; // write metadata bytes w.write_all(&table_bytes)?; @@ -836,9 +839,8 @@ impl<'a, 'b> FilesystemWriter<'a, 'b> { // write ptr for ptr in ptrs { - let mut bv = BitVec::new(); - ptr.write(&mut bv, self.kind.inner.type_endian)?; - w.write_all(bv.as_raw_slice())?; + let mut writer = Writer::new(w); + ptr.to_writer(&mut writer, self.kind.inner.type_endian)?; } Ok((table_position, count)) diff --git a/src/inode.rs b/src/inode.rs index a83bccca..23f331a4 100644 --- a/src/inode.rs +++ b/src/inode.rs @@ -3,7 +3,6 @@ use core::fmt; use std::io::Write; -use deku::bitvec::{BitVec, Msb0}; use deku::prelude::*; use crate::data::DataSize; @@ -36,9 +35,10 @@ impl Inode { superblock: &SuperBlock, kind: &Kind, ) -> Entry<'a> { - let mut bytes = BitVec::::new(); - self.write( - &mut bytes, + let mut inode_bytes = vec![]; + let mut writer = Writer::new(&mut inode_bytes); + self.to_writer( + &mut writer, ( 0xffff_ffff_ffff_ffff, // bytes_used is unused for ctx. set to max superblock.block_size, @@ -49,7 +49,7 @@ impl Inode { .unwrap(); let start = m_writer.metadata_start; let offset = m_writer.uncompressed_bytes.len() as u16; - m_writer.write_all(bytes.as_raw_slice()).unwrap(); + m_writer.write_all(&inode_bytes).unwrap(); Entry { start, diff --git a/src/metadata.rs b/src/metadata.rs index b8341ce3..264f1530 100644 --- a/src/metadata.rs +++ b/src/metadata.rs @@ -1,7 +1,6 @@ use std::collections::VecDeque; use std::io::{self, Cursor, Read, Seek, Write}; -use deku::bitvec::BitVec; use deku::prelude::*; use tracing::{instrument, trace}; @@ -41,6 +40,7 @@ impl MetadataWriter { #[instrument(skip_all)] fn add_block(&mut self) -> io::Result<()> { + trace!("adding block"); // uncompress data that will create the metablock let uncompressed_len = self.uncompressed_bytes.len().min(METADATA_MAXSIZE); if uncompressed_len == 0 { @@ -63,9 +63,11 @@ impl MetadataWriter { // Remove the data consumed, if the uncompressed data is smalled, use it. let (compressed, metadata) = if compressed.len() > uncompressed_len { + trace!("using uncompressed"); let uncompressed = self.uncompressed_bytes.drain(0..uncompressed_len).collect(); (false, uncompressed) } else { + trace!("using compressed"); self.uncompressed_bytes.drain(0..uncompressed_len); (true, compressed) }; @@ -75,7 +77,6 @@ impl MetadataWriter { trace!("new metadata start: {:#02x?}", self.metadata_start); self.final_bytes.push((compressed, metadata)); - trace!("LEN: {:02x?}", self.uncompressed_bytes.len()); Ok(()) } @@ -87,15 +88,14 @@ impl MetadataWriter { } // write all the metadata blocks - for (compressed, cb) in &self.final_bytes { - trace!("len: {:02x?}", cb.len()); - //trace!("total: {:02x?}", out.len()); - let mut bv = BitVec::new(); + for (compressed, compressed_bytes) in &self.final_bytes { + trace!("len: {:02x?}", compressed_bytes.len()); // if uncompressed, set the highest bit of len - let len = cb.len() as u16 | if *compressed { 0 } else { 1 << (u16::BITS - 1) }; - len.write(&mut bv, self.kind.inner.data_endian)?; - out.write_all(bv.as_raw_slice())?; - out.write_all(cb)?; + let len = + compressed_bytes.len() as u16 | if *compressed { 0 } else { 1 << (u16::BITS - 1) }; + let mut writer = Writer::new(out); + len.to_writer(&mut writer, self.kind.inner.data_endian)?; + out.write_all(compressed_bytes)?; } Ok(())