From c81aee5b253da3e905aad0d14315bdaf2ebe6b37 Mon Sep 17 00:00:00 2001 From: Kristof Date: Mon, 9 Sep 2024 16:34:02 +0200 Subject: [PATCH] work --- src/lib.rs | 6 ++-- src/preflate_container.rs | 59 +++++++++++++++++++++++++++++++-------- src/process.rs | 5 +--- src/scan_deflate.rs | 4 +-- tests/end_to_end.rs | 2 +- 5 files changed, 54 insertions(+), 22 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 048f54f..8ace1f1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -60,7 +60,7 @@ pub struct DecompressResult { pub compressed_size: usize, /// the parameters that were used to compress the stream (informational) - pub parameters: Option, + pub parameters: PreflateParameters, } impl core::fmt::Debug for DecompressResult { @@ -116,7 +116,7 @@ pub fn decompress_deflate_stream( plain_text: contents.plain_text, prediction_corrections: cabac_encoded, compressed_size: contents.compressed_size, - parameters: Some(params), + parameters: params, }) } @@ -180,7 +180,7 @@ pub fn decompress_deflate_stream_assert( plain_text: contents.plain_text, prediction_corrections: cabac_encoded, compressed_size: contents.compressed_size, - parameters: Some(params), + parameters: params, }) } diff --git a/src/preflate_container.rs b/src/preflate_container.rs index 29f8116..7405b99 100644 --- a/src/preflate_container.rs +++ b/src/preflate_container.rs @@ -1,15 +1,11 @@ -use anyhow::Result; use byteorder::ReadBytesExt; -use std::io::{Cursor, Read, Write}; +use std::io::{Read, Write}; use crate::{ idat_parse::{recreate_idat, IdatContents}, preflate_error::PreflateError, - preflate_parameter_estimator::PreflateParameters, - process::DeflateContents, recompress_deflate_stream, scan_deflate::{split_into_deflate_streams, BlockChunk}, - DecompressResult, }; const COMPRESSED_WRAPPER_VERSION_1: u8 = 1; @@ -75,11 +71,11 @@ fn test_variant_roundtrip() { } fn write_chunk_block( - signature: BlockChunk, + block: BlockChunk, literal_data: &[u8], destination: &mut impl Write, ) -> std::io::Result { - match signature { + match block { BlockChunk::Literal(content_size) => { destination.write_all(&[LITERAL_CHUNK])?; write_varint(destination, content_size as u32)?; @@ -106,7 +102,7 @@ fn write_chunk_block( write_varint(destination, res.prediction_corrections.len() as u32)?; destination.write_all(&res.prediction_corrections)?; - Ok(res.compressed_size) + Ok(idat.total_chunk_length) } } } @@ -152,7 +148,7 @@ fn read_chunk_block( if let Some(idat) = idat { recreate_idat(&idat, &recompressed[..], destination) - .map_err(|e| PreflateError::InvalidCompressedWrapper)?; + .map_err(|_e| PreflateError::InvalidCompressedWrapper)?; } else { destination.write_all(&recompressed)?; } @@ -168,7 +164,7 @@ fn roundtrip_chunk_block_literal() { write_chunk_block(BlockChunk::Literal(5), b"hello", &mut buffer).unwrap(); - let mut read_cursor = Cursor::new(buffer); + let mut read_cursor = std::io::Cursor::new(buffer); let mut destination = Vec::new(); read_chunk_block(&mut read_cursor, &mut destination).unwrap(); @@ -184,7 +180,7 @@ fn roundtrip_chunk_block_deflate() { write_chunk_block(BlockChunk::DeflateStream(results), &[], &mut buffer).unwrap(); - let mut read_cursor = Cursor::new(buffer); + let mut read_cursor = std::io::Cursor::new(buffer); let mut destination = Vec::new(); read_chunk_block(&mut read_cursor, &mut destination).unwrap(); @@ -210,7 +206,7 @@ fn roundtrip_chunk_block_png() { ) .unwrap(); - let mut read_cursor = Cursor::new(buffer); + let mut read_cursor = std::io::Cursor::new(buffer); let mut destination = Vec::new(); read_chunk_block(&mut read_cursor, &mut destination).unwrap(); @@ -224,6 +220,7 @@ pub fn expand_zlib_chunks(compressed_data: &[u8]) -> std::result::Result let mut locations_found = Vec::new(); split_into_deflate_streams(compressed_data, &mut locations_found); + println!("locations found: {:?}", locations_found); let mut plain_text = Vec::new(); plain_text.push(COMPRESSED_WRAPPER_VERSION_1); // version 1 of format. Definitely will improved. @@ -255,3 +252,41 @@ pub fn recreated_zlib_chunks( Ok(()) } + +#[cfg(test)] +fn roundtrip_deflate_chunks(filename: &str) { + let f = crate::process::read_file(filename); + + let expanded = expand_zlib_chunks(&f).unwrap(); + + let mut read_cursor = std::io::Cursor::new(expanded); + + let mut destination = Vec::new(); + recreated_zlib_chunks(&mut read_cursor, &mut destination).unwrap(); + + assert_eq!(destination.len(), f.len()); + for i in 0..destination.len() { + assert_eq!(destination[i], f[i], "Mismatch at index {}", i); + } + assert!(destination == f); +} + +#[test] +fn roundtrip_png_chunks() { + roundtrip_deflate_chunks("treegdi.png"); +} + +#[test] +fn roundtrip_zip_chunks() { + roundtrip_deflate_chunks("samplezip.zip"); +} + +#[test] +fn roundtrip_gz_chunks() { + roundtrip_deflate_chunks("sample1.bin.gz"); +} + +#[test] +fn roundtrip_pdf_chunks() { + roundtrip_deflate_chunks("starcontrol.samplesave"); +} diff --git a/src/process.rs b/src/process.rs index 22b8c70..3fe5450 100644 --- a/src/process.rs +++ b/src/process.rs @@ -164,11 +164,8 @@ fn recreate_blocks( Ok(output_blocks) } +#[allow(dead_code)] pub fn write_file(filename: &str, data: &[u8]) { - use std::fs::File; - use std::io::Write; - use std::path::Path; - let mut writecomp = std::fs::File::create(filename).unwrap(); std::io::Write::write_all(&mut writecomp, &data).unwrap(); } diff --git a/src/scan_deflate.rs b/src/scan_deflate.rs index 9a6b0a9..26ff09f 100644 --- a/src/scan_deflate.rs +++ b/src/scan_deflate.rs @@ -1,4 +1,4 @@ -use std::{io::Cursor, thread::panicking}; +use std::io::Cursor; use crate::{ decompress_deflate_stream, @@ -123,7 +123,7 @@ pub fn split_into_deflate_streams(src: &[u8], locations_found: &mut Vec