1use crate::block::{
4 compress_block, decompress_block_into, resolve_compression_level, CompressedBlock,
5};
6use crate::config::{EngineConfiguration, ProgressEvent, ProgressPhase};
7use crate::format::{BlockHeader, BlockIndexEntry, FileFlags, FileFooter, FileHeader, IndexHeader};
8use crate::index::load_index;
9use crush_core::error::{CrushError, Result};
10use libdeflater::{Compressor, Decompressor};
11use rayon::prelude::*;
12use std::io::{Cursor, Read, Seek, SeekFrom, Write};
13use std::path::Path;
14
15fn with_pool<T: Send>(config: &EngineConfiguration, f: impl FnOnce() -> T + Send) -> T {
22 match &config.thread_pool {
23 Some(pool) => pool.install(f),
24 None => f(),
25 }
26}
27
28#[allow(clippy::too_many_lines)] pub fn compress(input: &[u8], config: &EngineConfiguration) -> Result<Vec<u8>> {
40 let block_size = config.block_size as usize;
41 let blocks: Vec<&[u8]> = input.chunks(block_size).collect();
42 let total_blocks = blocks.len() as u64;
43
44 let lvl = resolve_compression_level(config.compression_level)?;
46
47 let results: Vec<Result<CompressedBlock>> = with_pool(config, || {
51 blocks
52 .par_iter()
53 .enumerate()
54 .map_init(
55 || Compressor::new(lvl),
56 |compressor, (i, chunk)| compress_block(compressor, chunk, i, config),
57 )
58 .collect()
59 });
60
61 #[allow(clippy::cast_possible_truncation)]
63 let mut compressed_blocks: Vec<CompressedBlock> = Vec::with_capacity(total_blocks as usize);
64 for r in results {
65 compressed_blocks.push(r?);
66 }
67
68 let mut flags = FileFlags::default();
73 if config.checksums {
74 flags = flags.with_checksums();
75 }
76
77 let file_header = FileHeader::new(
78 config.block_size,
79 config.compression_level,
80 flags,
81 input.len() as u64,
82 total_blocks,
83 );
84
85 let header_size = FileHeader::SIZE;
86 let blocks_region_size: usize = compressed_blocks
87 .iter()
88 .map(|b| BlockHeader::SIZE + b.payload.len())
89 .sum();
90 let index_region_size = IndexHeader::SIZE + compressed_blocks.len() * BlockIndexEntry::SIZE;
91 let footer_size = FileFooter::SIZE;
92 let total_size = header_size + blocks_region_size + index_region_size + footer_size;
93
94 let mut out = vec![0u8; total_size];
95
96 let mut block_offsets: Vec<usize> = Vec::with_capacity(compressed_blocks.len());
98 let mut cursor = header_size;
99 for b in &compressed_blocks {
100 block_offsets.push(cursor);
101 cursor += BlockHeader::SIZE + b.payload.len();
102 }
103 debug_assert_eq!(cursor, header_size + blocks_region_size);
104
105 out[..header_size].copy_from_slice(&file_header.to_bytes());
107
108 {
111 let (_header_region, rest) = out.split_at_mut(header_size);
112 let (blocks_region, _trailing) = rest.split_at_mut(blocks_region_size);
113
114 let mut block_slices: Vec<&mut [u8]> = Vec::with_capacity(compressed_blocks.len());
115 let mut remaining: &mut [u8] = blocks_region;
116 for b in &compressed_blocks {
117 let chunk_len = BlockHeader::SIZE + b.payload.len();
118 let (chunk, rest) = remaining.split_at_mut(chunk_len);
119 block_slices.push(chunk);
120 remaining = rest;
121 }
122 debug_assert!(remaining.is_empty());
123
124 block_slices
125 .into_par_iter()
126 .zip(compressed_blocks.par_iter())
127 .for_each(|(slice, block)| {
128 let hdr_bytes = block.header.to_bytes();
129 slice[..BlockHeader::SIZE].copy_from_slice(&hdr_bytes);
130 slice[BlockHeader::SIZE..].copy_from_slice(&block.payload);
131 });
132 }
133
134 let mut index_entries = Vec::with_capacity(compressed_blocks.len());
136 for (i, b) in compressed_blocks.iter().enumerate() {
137 index_entries.push(BlockIndexEntry {
138 block_offset: block_offsets[i] as u64,
139 compressed_size: b.header.compressed_size,
140 uncompressed_size: b.header.uncompressed_size,
141 checksum: b.header.checksum,
142 });
143 }
144
145 let index_offset = header_size + blocks_region_size;
147 let entry_count = u32::try_from(index_entries.len())
148 .map_err(|_| CrushError::InvalidConfig("too many blocks for index".to_owned()))?;
149 let ih = IndexHeader {
150 entry_count,
151 index_flags: 0,
152 };
153 out[index_offset..index_offset + IndexHeader::SIZE].copy_from_slice(&ih.to_bytes());
154 let mut entry_cursor = index_offset + IndexHeader::SIZE;
155 for e in &index_entries {
156 out[entry_cursor..entry_cursor + BlockIndexEntry::SIZE].copy_from_slice(&e.to_bytes());
157 entry_cursor += BlockIndexEntry::SIZE;
158 }
159 debug_assert_eq!(entry_cursor, index_offset + index_region_size);
160
161 let index_size_u32 = u32::try_from(index_region_size)
163 .map_err(|_| CrushError::InvalidConfig("index too large for footer".to_owned()))?;
164 let footer = FileFooter::new(index_offset as u64, index_size_u32);
165 out[total_size - FileFooter::SIZE..].copy_from_slice(&footer.to_bytes());
166
167 if let Some(cb_arc) = &config.progress {
170 let mut cb = cb_arc.lock().map_err(|_| {
171 CrushError::InvalidConfig("progress callback mutex poisoned".to_owned())
172 })?;
173 let mut bytes_processed: u64 = 0;
174 for (i, block) in compressed_blocks.iter().enumerate() {
175 bytes_processed += u64::from(block.header.uncompressed_size);
176 let event = ProgressEvent {
177 bytes_processed,
178 blocks_completed: i as u64 + 1,
179 total_blocks: Some(total_blocks),
180 phase: ProgressPhase::Compressing,
181 };
182 if !cb(event) {
183 return Err(CrushError::Cancelled);
184 }
185 }
186 }
187
188 Ok(out)
189}
190
191pub fn compress_file(path: &Path, config: &EngineConfiguration) -> Result<Vec<u8>> {
207 let file = std::fs::File::open(path)?;
208 let mmap = unsafe { memmap2::MmapOptions::new().map(&file)? };
210 compress(&mmap, config)
211}
212
213pub fn compress_to_writer<W: Write>(
223 input: &[u8],
224 mut writer: W,
225 config: &EngineConfiguration,
226) -> Result<u64> {
227 let out = compress(input, config)?;
228 let len = out.len() as u64;
229 writer.write_all(&out)?;
230 Ok(len)
231}
232
233pub fn compress_stream<R: Read, W: Write>(
249 mut reader: R,
250 mut writer: W,
251 config: &EngineConfiguration,
252) -> Result<u64> {
253 let mut input = Vec::new();
254 reader.read_to_end(&mut input)?;
255 let out = compress(&input, config)?;
256 let len = out.len() as u64;
257 writer.write_all(&out)?;
258 Ok(len)
259}
260
261pub fn decompress(input: &[u8], config: &EngineConfiguration) -> Result<Vec<u8>> {
279 let mut cursor = Cursor::new(input);
280 decompress_from_reader(&mut cursor, config)
281}
282
283#[allow(clippy::too_many_lines)] pub fn decompress_from_reader<R: Read + Seek>(
305 reader: &mut R,
306 config: &EngineConfiguration,
307) -> Result<Vec<u8>> {
308 let index = load_index(reader)?;
309
310 #[allow(
313 clippy::cast_precision_loss,
314 clippy::cast_possible_truncation,
315 clippy::cast_sign_loss
316 )]
317 let limit = {
318 let file_size = reader.seek(SeekFrom::End(0))?;
319 (file_size as f64 * config.max_decompression_ratio) as u64
320 };
321 let total_uncompressed = index.total_uncompressed_size();
322 if total_uncompressed > limit {
323 return Err(CrushError::ExpansionLimitExceeded { block_index: 0 });
324 }
325
326 let total_blocks = index.len();
327 let checksums_enabled = index.checksums_enabled;
328
329 let raw_blocks: Vec<(BlockHeader, Vec<u8>)> = index
331 .entries
332 .iter()
333 .enumerate()
334 .map(|(i, entry)| -> Result<(BlockHeader, Vec<u8>)> {
335 reader.seek(SeekFrom::Start(entry.block_offset))?;
336 let mut hdr_buf = [0u8; BlockHeader::SIZE];
337 reader.read_exact(&mut hdr_buf).map_err(|e| {
338 CrushError::InvalidFormat(format!("block {i} header read error: {e}"))
339 })?;
340 let header = BlockHeader::from_bytes(&hdr_buf);
341 let mut payload = vec![0u8; header.compressed_size as usize];
342 reader.read_exact(&mut payload).map_err(|e| {
343 CrushError::InvalidFormat(format!("block {i} payload read error: {e}"))
344 })?;
345 Ok((header, payload))
346 })
347 .collect::<Result<Vec<_>>>()?;
348
349 let total_usize = usize::try_from(total_uncompressed).map_err(|_| {
351 CrushError::InvalidConfig(format!(
352 "total uncompressed size {total_uncompressed} overflows usize"
353 ))
354 })?;
355 let mut output = vec![0u8; total_usize];
356
357 let mut per_block_sizes: Vec<usize> = Vec::with_capacity(raw_blocks.len());
360 let mut running: usize = 0;
361 for (header, _) in &raw_blocks {
362 let sz = header.uncompressed_size as usize;
363 per_block_sizes.push(sz);
364 running = running.checked_add(sz).ok_or_else(|| {
365 CrushError::InvalidFormat(
366 "block uncompressed sizes overflow usize during offset calculation".to_owned(),
367 )
368 })?;
369 }
370 debug_assert_eq!(running, total_usize);
371
372 let mut output_slices: Vec<&mut [u8]> = Vec::with_capacity(raw_blocks.len());
374 {
375 let mut remaining: &mut [u8] = &mut output;
376 for &sz in &per_block_sizes {
377 let (chunk, rest) = remaining.split_at_mut(sz);
378 output_slices.push(chunk);
379 remaining = rest;
380 }
381 debug_assert!(remaining.is_empty());
382 }
383
384 #[allow(clippy::type_complexity)]
387 let pairs: Vec<(&mut [u8], &(BlockHeader, Vec<u8>))> =
388 output_slices.into_iter().zip(raw_blocks.iter()).collect();
389
390 let results: Vec<Result<()>> = with_pool(config, || {
391 pairs
392 .into_par_iter()
393 .enumerate()
394 .map_init(
395 Decompressor::new,
396 |decompressor, (i, (slice, (header, payload)))| {
397 decompress_block_into(
398 decompressor,
399 header,
400 payload,
401 slice,
402 i as u64,
403 checksums_enabled,
404 )
405 },
406 )
407 .collect()
408 });
409
410 for r in results {
411 r?;
412 }
413
414 if let Some(cb_arc) = &config.progress {
416 let mut cb = cb_arc
417 .lock()
418 .map_err(|_| CrushError::InvalidConfig("progress mutex poisoned".to_owned()))?;
419 let mut bytes_processed: u64 = 0;
420 for (i, (header, _)) in raw_blocks.iter().enumerate() {
421 bytes_processed += u64::from(header.uncompressed_size);
422 let event = ProgressEvent {
423 bytes_processed,
424 blocks_completed: i as u64 + 1,
425 total_blocks: Some(total_blocks),
426 phase: ProgressPhase::Decompressing,
427 };
428 if !cb(event) {
429 return Err(CrushError::Cancelled);
430 }
431 }
432 }
433
434 Ok(output)
435}
436
437#[cfg(test)]
442#[allow(
443 clippy::expect_used,
444 clippy::unwrap_used,
445 clippy::panic,
446 clippy::cast_possible_truncation,
447 clippy::missing_panics_doc
448)]
449mod tests {
450 use super::*;
451 use crate::format::FORMAT_VERSION;
452 use std::io::Write as IoWrite;
453 use tempfile::NamedTempFile;
454
455 fn default_config() -> EngineConfiguration {
456 EngineConfiguration::builder()
457 .block_size(65_536) .build()
459 .expect("config")
460 }
461
462 #[test]
463 fn test_compress_roundtrip_small() {
464 let data: Vec<u8> = b"hello world"
465 .iter()
466 .cycle()
467 .take(200_000)
468 .copied()
469 .collect();
470 let config = default_config();
471 let compressed = compress(&data, &config).expect("compress");
472 let recovered = decompress(&compressed, &config).expect("decompress");
473 assert_eq!(data, recovered);
474 }
475
476 #[test]
477 fn test_compress_incompressible_stored() {
478 let data: Vec<u8> = b"hello world!"
483 .iter()
484 .cycle()
485 .take(200_000)
486 .copied()
487 .collect();
488 let config = EngineConfiguration::builder()
489 .block_size(65_536)
490 .max_expansion_ratio(0.001) .build()
492 .expect("config");
493 let compressed = compress(&data, &config).expect("compress");
494 let recovered = decompress(&compressed, &config).expect("decompress");
496 assert_eq!(data, recovered);
497 let mut cursor = Cursor::new(&compressed);
499 let index = load_index(&mut cursor).expect("load_index");
500 cursor
502 .seek(SeekFrom::Start(index.entries[0].block_offset))
503 .expect("seek");
504 let mut hdr = [0u8; BlockHeader::SIZE];
505 cursor.read_exact(&mut hdr).expect("read hdr");
506 let header = BlockHeader::from_bytes(&hdr);
507 assert!(
508 header.flags.stored(),
509 "expected stored flag on incompressible data"
510 );
511 }
512
513 #[test]
514 fn test_compress_output_valid_crsh_format() {
515 let data: Vec<u8> = b"test".iter().cycle().take(100_000).copied().collect();
516 let config = default_config();
517 let compressed = compress(&data, &config).expect("compress");
518 let hdr_bytes: [u8; FileHeader::SIZE] = compressed[..FileHeader::SIZE]
520 .try_into()
521 .expect("hdr bytes");
522 let hdr = FileHeader::from_bytes(&hdr_bytes).expect("parse header");
523 assert_eq!(hdr.magic, crate::format::CRSH_MAGIC);
524 assert_eq!(hdr.format_version, FORMAT_VERSION);
525 let footer_bytes: [u8; FileFooter::SIZE] = compressed
527 [compressed.len() - FileFooter::SIZE..]
528 .try_into()
529 .expect("footer bytes");
530 let footer = FileFooter::from_bytes(&footer_bytes).expect("parse footer");
531 assert_eq!(footer.magic, crate::format::CRSH_MAGIC);
532 let mut cursor = Cursor::new(&compressed);
534 let index = load_index(&mut cursor).expect("load_index");
535 assert_eq!(index.len(), hdr.block_count);
536 }
537
538 #[test]
539 fn test_progress_callback_invoked_per_block() {
540 use std::sync::{Arc, Mutex};
541 let data: Vec<u8> = b"abc".iter().cycle().take(300_000).copied().collect();
542 let count = Arc::new(Mutex::new(0u64));
543 let count_clone = count.clone();
544 let cb: crate::config::ProgressCallback = Box::new(move |_event| {
545 let mut c = count_clone.lock().expect("lock");
546 *c += 1;
547 true
548 });
549 let config = EngineConfiguration::builder()
550 .block_size(65_536)
551 .progress(Arc::new(Mutex::new(cb)))
552 .build()
553 .expect("config");
554 compress(&data, &config).expect("compress");
555 let final_count = *count.lock().expect("lock");
556 assert!(final_count >= 1, "progress callback was not invoked");
558 }
559
560 #[test]
561 fn test_cancel_halts_at_block_boundary() {
562 use std::sync::{Arc, Mutex};
563 let data: Vec<u8> = b"xyz".iter().cycle().take(1_000_000).copied().collect();
564 let cb: crate::config::ProgressCallback = Box::new(|_event| false); let config = EngineConfiguration::builder()
566 .block_size(65_536)
567 .progress(Arc::new(Mutex::new(cb)))
568 .build()
569 .expect("config");
570 let result = compress(&data, &config);
571 assert!(result.is_err());
572 assert!(result.unwrap_err().is_cancelled());
573 }
574
575 #[test]
576 fn test_compress_file_roundtrip() {
577 let data: Vec<u8> = b"file data".iter().cycle().take(200_000).copied().collect();
578 let mut tmp = NamedTempFile::new().expect("temp file");
579 tmp.write_all(&data).expect("write");
580 let config = default_config();
581 let compressed = compress_file(tmp.path(), &config).expect("compress_file");
582 let recovered = decompress(&compressed, &config).expect("decompress");
583 assert_eq!(data, recovered);
584 }
585
586 #[test]
587 fn test_decompress_roundtrip() {
588 let data: Vec<u8> = b"decompress me"
589 .iter()
590 .cycle()
591 .take(500_000)
592 .copied()
593 .collect();
594 let config = default_config();
595 let compressed = compress(&data, &config).expect("compress");
596 let recovered = decompress(&compressed, &config).expect("decompress");
597 assert_eq!(data, recovered);
598 }
599
600 #[test]
601 fn test_decompress_corrupt_block_detected() {
602 let data: Vec<u8> = b"corrupt test"
603 .iter()
604 .cycle()
605 .take(200_000)
606 .copied()
607 .collect();
608 let config = default_config();
609 let mut compressed = compress(&data, &config).expect("compress");
610
611 let mut cursor = Cursor::new(&compressed);
613 let index = load_index(&mut cursor).expect("load_index");
614 let block0_offset = index.entries[0].block_offset as usize;
615
616 let payload_start = block0_offset + BlockHeader::SIZE;
618 if payload_start < compressed.len() {
619 compressed[payload_start] ^= 0xFF;
620 }
621
622 let result = decompress(&compressed, &config);
623 assert!(result.is_err());
624 let err = result.unwrap_err();
625 assert!(
626 matches!(err, CrushError::ChecksumMismatch { block_index: 0, .. })
627 || matches!(err, CrushError::InvalidFormat(_)),
628 "expected checksum or format error, got {err:?}"
629 );
630 }
631
632 #[test]
633 fn test_version_mismatch_rejected() {
634 let data: Vec<u8> = b"version test"
635 .iter()
636 .cycle()
637 .take(100_000)
638 .copied()
639 .collect();
640 let config = default_config();
641 let mut compressed = compress(&data, &config).expect("compress");
642
643 let footer_start = compressed.len() - FileFooter::SIZE;
645 compressed[footer_start + 16..footer_start + 20].copy_from_slice(&9999u32.to_le_bytes());
647 let result = decompress(&compressed, &config);
651 assert!(result.is_err());
652 }
653
654 #[test]
655 fn test_expansion_limit_exceeded() {
656 let data: Vec<u8> = b"test data".iter().cycle().take(100_000).copied().collect();
657 let compress_config = default_config();
658 let compressed = compress(&data, &compress_config).expect("compress");
659
660 let decompress_config = EngineConfiguration::builder()
662 .block_size(65_536)
663 .max_decompression_ratio(0.000_001)
664 .build()
665 .expect("config");
666 let result = decompress(&compressed, &decompress_config);
667 assert!(result.is_err());
668 assert!(matches!(
669 result.unwrap_err(),
670 CrushError::ExpansionLimitExceeded { .. }
671 ));
672 }
673
674 #[test]
675 fn test_truncated_footer_rejected() {
676 let data: Vec<u8> = b"truncated".iter().cycle().take(100_000).copied().collect();
677 let config = default_config();
678 let mut compressed = compress(&data, &config).expect("compress");
679 compressed.truncate(compressed.len() - FileFooter::SIZE);
681 let result = decompress(&compressed, &config);
682 assert!(result.is_err());
683 }
684
685 proptest::proptest! {
690 #![proptest_config(proptest::prelude::ProptestConfig::with_cases(50))]
691
692 #[test]
693 fn proptest_compress_decompress_roundtrip(
694 data in proptest::collection::vec(proptest::prelude::any::<u8>(), 0..200_000),
695 block_kb in proptest::prelude::prop_oneof![
696 proptest::prelude::Just(64usize),
697 proptest::prelude::Just(256),
698 proptest::prelude::Just(1024)
699 ],
700 level in 0u8..=9,
701 ) {
702 let block_size = u32::try_from(block_kb * 1024).unwrap();
703 let config = EngineConfiguration::builder()
704 .block_size(block_size)
705 .compression_level(level)
706 .build()
707 .unwrap();
708 let compressed = compress(&data, &config).unwrap();
709 let recovered = decompress(&compressed, &config).unwrap();
710 proptest::prop_assert_eq!(data, recovered);
711 }
712 }
713}