Skip to main content

initramfs_builder/initramfs/
compress.rs

1use anyhow::{Context, Result};
2use flate2::write::GzEncoder;
3use flate2::Compression as GzCompression;
4use std::fs::File;
5use std::io::{BufWriter, Write};
6use std::path::Path;
7use tracing::info;
8
9#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
10pub enum Compression {
11    #[default]
12    Gzip,
13    Zstd,
14    None,
15}
16
17impl std::str::FromStr for Compression {
18    type Err = String;
19
20    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
21        match s.to_lowercase().as_str() {
22            "gzip" | "gz" => Ok(Compression::Gzip),
23            "zstd" | "zst" => Ok(Compression::Zstd),
24            "none" | "raw" => Ok(Compression::None),
25            _ => Err(format!("Unknown compression: {}", s)),
26        }
27    }
28}
29
30impl std::fmt::Display for Compression {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        match self {
33            Compression::Gzip => write!(f, "gzip"),
34            Compression::Zstd => write!(f, "zstd"),
35            Compression::None => write!(f, "none"),
36        }
37    }
38}
39
40/// Compress data and write to output path
41pub fn compress_archive(data: &[u8], output_path: &Path, compression: Compression) -> Result<u64> {
42    info!(
43        "Compressing {} bytes with {} to {:?}",
44        data.len(),
45        compression,
46        output_path
47    );
48
49    let file = File::create(output_path)
50        .with_context(|| format!("Failed to create output file: {:?}", output_path))?;
51    let mut writer = BufWriter::new(file);
52
53    match compression {
54        Compression::Gzip => {
55            let mut encoder = GzEncoder::new(&mut writer, GzCompression::default());
56            encoder.write_all(data)?;
57            encoder.finish()?;
58        }
59        Compression::Zstd => {
60            let mut encoder = zstd::stream::Encoder::new(&mut writer, 3)?;
61            encoder.write_all(data)?;
62            encoder.finish()?;
63        }
64        Compression::None => {
65            writer.write_all(data)?;
66        }
67    }
68
69    writer.flush()?;
70
71    let output_size = std::fs::metadata(output_path)?.len();
72    info!(
73        "Compressed {} bytes -> {} bytes ({:.1}% ratio)",
74        data.len(),
75        output_size,
76        (output_size as f64 / data.len() as f64) * 100.0
77    );
78
79    Ok(output_size)
80}