use std::borrow::Cow;
use std::convert::TryFrom;
use std::fmt;
use std::io::{BufReader, Error, Read, Result, Write};
use std::str::FromStr;
mod lz4_standard;
use self::lz4_standard::*;
#[cfg(feature = "zran")]
pub mod zlib_random;
const COMPRESSION_MINIMUM_RATIO: usize = 100;
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Algorithm {
None,
Lz4Block,
GZip,
Zstd,
}
impl Default for Algorithm {
fn default() -> Self {
Self::None
}
}
impl fmt::Display for Algorithm {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{:?}", self)
}
}
impl FromStr for Algorithm {
type Err = Error;
fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
match s {
"none" => Ok(Self::None),
"lz4_block" => Ok(Self::Lz4Block),
"gzip" => Ok(Self::GZip),
"zstd" => Ok(Self::Zstd),
_ => Err(einval!("compression algorithm should be none or lz4_block")),
}
}
}
impl TryFrom<u32> for Algorithm {
type Error = ();
fn try_from(value: u32) -> std::result::Result<Self, Self::Error> {
if value == Algorithm::None as u32 {
Ok(Algorithm::None)
} else if value == Algorithm::Lz4Block as u32 {
Ok(Algorithm::Lz4Block)
} else if value == Algorithm::GZip as u32 {
Ok(Algorithm::GZip)
} else if value == Algorithm::Zstd as u32 {
Ok(Algorithm::Zstd)
} else {
Err(())
}
}
}
impl TryFrom<u64> for Algorithm {
type Error = ();
fn try_from(value: u64) -> std::result::Result<Self, Self::Error> {
if value == Algorithm::None as u64 {
Ok(Algorithm::None)
} else if value == Algorithm::Lz4Block as u64 {
Ok(Algorithm::Lz4Block)
} else if value == Algorithm::GZip as u64 {
Ok(Algorithm::GZip)
} else if value == Algorithm::Zstd as u64 {
Ok(Algorithm::Zstd)
} else {
Err(())
}
}
}
impl Algorithm {
pub fn is_none(self) -> bool {
self == Self::None
}
}
pub fn compress(src: &[u8], algorithm: Algorithm) -> Result<(Cow<[u8]>, bool)> {
let src_size = src.len();
if src_size == 0 {
return Ok((Cow::Borrowed(src), false));
}
let compressed = match algorithm {
Algorithm::None => return Ok((Cow::Borrowed(src), false)),
Algorithm::Lz4Block => lz4_compress(src)?,
Algorithm::GZip => {
let dst: Vec<u8> = Vec::new();
let mut gz = flate2::write::GzEncoder::new(dst, flate2::Compression::default());
gz.write_all(src)?;
gz.finish()?
}
Algorithm::Zstd => zstd_compress(src)?,
};
if (COMPRESSION_MINIMUM_RATIO == 100 && compressed.len() >= src_size)
|| ((100 * compressed.len() / src_size) >= COMPRESSION_MINIMUM_RATIO)
{
Ok((Cow::Borrowed(src), false))
} else {
Ok((Cow::Owned(compressed), true))
}
}
pub fn decompress(src: &[u8], dst: &mut [u8], algorithm: Algorithm) -> Result<usize> {
match algorithm {
Algorithm::None => {
assert_eq!(src.len(), dst.len());
dst.copy_from_slice(src);
Ok(dst.len())
}
Algorithm::Lz4Block => lz4_decompress(src, dst),
Algorithm::GZip => {
let mut gz = flate2::bufread::GzDecoder::new(src);
gz.read_exact(dst)?;
Ok(dst.len())
}
Algorithm::Zstd => zstd::bulk::decompress_to_buffer(src, dst),
}
}
#[allow(clippy::large_enum_variant)]
pub enum Decoder<'a, R: Read> {
None(R),
Gzip(flate2::bufread::GzDecoder<BufReader<R>>),
Zstd(zstd::stream::Decoder<'a, BufReader<R>>),
}
impl<'a, R: Read> Decoder<'a, R> {
pub fn new(reader: R, algorithm: Algorithm) -> Result<Self> {
let decoder = match algorithm {
Algorithm::None => Decoder::None(reader),
Algorithm::GZip => {
Decoder::Gzip(flate2::bufread::GzDecoder::new(BufReader::new(reader)))
}
Algorithm::Lz4Block => panic!("Decoder doesn't support lz4_block"),
Algorithm::Zstd => Decoder::Zstd(zstd::stream::Decoder::new(reader)?),
};
Ok(decoder)
}
}
impl<'a, R: Read> Read for Decoder<'a, R> {
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
match self {
Decoder::None(r) => r.read(buf),
Decoder::Gzip(r) => r.read(buf),
Decoder::Zstd(r) => r.read(buf),
}
}
}
pub struct ZlibDecoder<R> {
stream: flate2::bufread::GzDecoder<BufReader<R>>,
}
impl<R: Read> ZlibDecoder<R> {
pub fn new(reader: R) -> Self {
ZlibDecoder {
stream: flate2::bufread::GzDecoder::new(BufReader::new(reader)),
}
}
}
impl<R: Read> Read for ZlibDecoder<R> {
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
self.stream.read(buf)
}
}
pub fn compute_compressed_gzip_size(size: usize, max_size: usize) -> usize {
let size = size + 10 + 8 + 5 + (size / (16 << 10)) * 5 + 128;
std::cmp::min(size, max_size)
}
fn zstd_compress(src: &[u8]) -> Result<Vec<u8>> {
zstd::bulk::compress(src, zstd::DEFAULT_COMPRESSION_LEVEL)
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs::OpenOptions;
use std::io::{Seek, SeekFrom};
use std::path::Path;
use vmm_sys_util::tempfile::TempFile;
#[test]
fn test_compress_algorithm_gzip() {
let buf = vec![0x2u8; 4095];
let compressed = compress(&buf, Algorithm::GZip).unwrap();
assert!(compressed.1);
let (compressed, _) = compressed;
assert_ne!(compressed.len(), 0);
let mut decompressed = vec![0; buf.len()];
let sz = decompress(&compressed, decompressed.as_mut_slice(), Algorithm::GZip).unwrap();
assert_eq!(sz, 4095);
assert_eq!(buf, decompressed);
let mut tmp_file = TempFile::new().unwrap().into_file();
tmp_file.write_all(&compressed).unwrap();
tmp_file.seek(SeekFrom::Start(0)).unwrap();
let mut decompressed = vec![0; buf.len()];
let mut decoder = Decoder::new(tmp_file, Algorithm::GZip).unwrap();
decoder.read_exact(decompressed.as_mut_slice()).unwrap();
assert_eq!(sz, 4095);
assert_eq!(buf, decompressed);
}
#[test]
fn test_compress_algorithm_none() {
let buf = [
0x1u8, 0x2u8, 0x3u8, 0x4u8, 0x1u8, 0x2u8, 0x3u8, 0x4u8, 0x1u8, 0x2u8, 0x3u8, 0x4u8,
0x1u8, 0x2u8, 0x3u8, 0x4u8,
];
let (compressed, _) = compress(&buf, Algorithm::None).unwrap();
assert_eq!(buf.to_vec(), compressed.to_vec());
}
#[test]
fn test_lz4_compress_decompress_1_byte() {
let buf = vec![0x1u8];
let compressed = lz4_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(
&compressed,
decompressed.as_mut_slice(),
Algorithm::Lz4Block,
)
.unwrap();
assert_eq!(sz, 1);
assert_eq!(buf, decompressed);
}
#[test]
fn test_lz4_compress_decompress_2_bytes() {
let buf = vec![0x2u8, 0x3u8];
let compressed = lz4_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(
&compressed,
decompressed.as_mut_slice(),
Algorithm::Lz4Block,
)
.unwrap();
assert_eq!(sz, 2);
assert_eq!(buf, decompressed);
}
#[test]
fn test_lz4_compress_decompress_16_bytes() {
let buf = [
0x1u8, 0x2u8, 0x3u8, 0x4u8, 0x1u8, 0x2u8, 0x3u8, 0x4u8, 0x1u8, 0x2u8, 0x3u8, 0x4u8,
0x1u8, 0x2u8, 0x3u8, 0x4u8,
];
let compressed = lz4_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(
&compressed,
decompressed.as_mut_slice(),
Algorithm::Lz4Block,
)
.unwrap();
assert_eq!(sz, 16);
assert_eq!(&buf, decompressed.as_slice());
}
#[test]
fn test_lz4_compress_decompress_4095_bytes() {
let buf = vec![0x2u8; 4095];
let compressed = lz4_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(
&compressed,
decompressed.as_mut_slice(),
Algorithm::Lz4Block,
)
.unwrap();
assert_eq!(sz, 4095);
assert_eq!(buf, decompressed);
}
#[test]
fn test_lz4_compress_decompress_4096_bytes() {
let buf = vec![0x2u8; 4096];
let compressed = lz4_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(
&compressed,
decompressed.as_mut_slice(),
Algorithm::Lz4Block,
)
.unwrap();
assert_eq!(sz, 4096);
assert_eq!(buf, decompressed);
}
#[test]
fn test_lz4_compress_decompress_4097_bytes() {
let buf = vec![0x2u8; 4097];
let compressed = lz4_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(
&compressed,
decompressed.as_mut_slice(),
Algorithm::Lz4Block,
)
.unwrap();
assert_eq!(sz, 4097);
assert_eq!(buf, decompressed);
}
#[test]
fn test_zstd_compress_decompress_1_byte() {
let buf = vec![0x1u8];
let compressed = zstd_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(&compressed, decompressed.as_mut_slice(), Algorithm::Zstd).unwrap();
assert_eq!(sz, 1);
assert_eq!(buf, decompressed);
}
#[test]
fn test_zstd_compress_decompress_2_bytes() {
let buf = vec![0x2u8, 0x3u8];
let compressed = zstd_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(&compressed, decompressed.as_mut_slice(), Algorithm::Zstd).unwrap();
assert_eq!(sz, 2);
assert_eq!(buf, decompressed);
}
#[test]
fn test_zstd_compress_decompress_16_bytes() {
let buf = [
0x1u8, 0x2u8, 0x3u8, 0x4u8, 0x1u8, 0x2u8, 0x3u8, 0x4u8, 0x1u8, 0x2u8, 0x3u8, 0x4u8,
0x1u8, 0x2u8, 0x3u8, 0x4u8,
];
let compressed = zstd_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(&compressed, decompressed.as_mut_slice(), Algorithm::Zstd).unwrap();
assert_eq!(sz, 16);
assert_eq!(&buf, decompressed.as_slice());
}
#[test]
fn test_zstd_compress_decompress_4095_bytes() {
let buf = vec![0x2u8; 4095];
let compressed = zstd_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(&compressed, decompressed.as_mut_slice(), Algorithm::Zstd).unwrap();
assert_eq!(sz, 4095);
assert_eq!(buf, decompressed);
}
#[test]
fn test_zstd_compress_decompress_4096_bytes() {
let buf = vec![0x2u8; 4096];
let compressed = zstd_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(&compressed, decompressed.as_mut_slice(), Algorithm::Zstd).unwrap();
assert_eq!(sz, 4096);
assert_eq!(buf, decompressed);
}
#[test]
fn test_zstd_compress_decompress_4097_bytes() {
let buf = vec![0x2u8; 4097];
let compressed = zstd_compress(&buf).unwrap();
let mut decompressed = vec![0; buf.len()];
let sz = decompress(&compressed, decompressed.as_mut_slice(), Algorithm::Zstd).unwrap();
assert_eq!(sz, 4097);
assert_eq!(buf, decompressed);
}
#[test]
fn test_new_decoder_none() {
let buf = b"This is a test";
let mut decoder = Decoder::new(buf.as_slice(), Algorithm::None).unwrap();
let mut buf2 = vec![0u8; 1024];
let res = decoder.read(&mut buf2).unwrap();
assert_eq!(res, 14);
assert_eq!(&buf2[0..14], buf.as_slice());
}
#[test]
fn test_gzip_decoder() {
let root_dir = &std::env::var("CARGO_MANIFEST_DIR").expect("$CARGO_MANIFEST_DIR");
let path = Path::new(root_dir).join("../tests/texture/zran/zlib_sample.txt.gz");
let file = OpenOptions::new().read(true).open(&path).unwrap();
let mut decoder = Decoder::new(file, Algorithm::GZip).unwrap();
let mut buf = [0u8; 8];
decoder.read_exact(&mut buf).unwrap();
assert_eq!(&String::from_utf8_lossy(&buf), "This is ");
decoder.read_exact(&mut buf).unwrap();
assert_eq!(&String::from_utf8_lossy(&buf), "a test f");
decoder.read_exact(&mut buf).unwrap();
assert_eq!(&String::from_utf8_lossy(&buf), "ile for ");
let ret = decoder.read(&mut buf).unwrap();
assert_eq!(ret, 6);
assert_eq!(&String::from_utf8_lossy(&buf[0..6]), "zlib.\n");
let ret = decoder.read(&mut buf).unwrap();
assert_eq!(ret, 0);
}
#[test]
fn test_zlib_decoder() {
let root_dir = &std::env::var("CARGO_MANIFEST_DIR").expect("$CARGO_MANIFEST_DIR");
let path = Path::new(root_dir).join("../tests/texture/zran/zlib_sample.txt.gz");
let file = OpenOptions::new().read(true).open(&path).unwrap();
let mut decoder = ZlibDecoder::new(file);
let mut buf = [0u8; 8];
decoder.read_exact(&mut buf).unwrap();
assert_eq!(&String::from_utf8_lossy(&buf), "This is ");
decoder.read_exact(&mut buf).unwrap();
assert_eq!(&String::from_utf8_lossy(&buf), "a test f");
decoder.read_exact(&mut buf).unwrap();
assert_eq!(&String::from_utf8_lossy(&buf), "ile for ");
let ret = decoder.read(&mut buf).unwrap();
assert_eq!(ret, 6);
assert_eq!(&String::from_utf8_lossy(&buf[0..6]), "zlib.\n");
let ret = decoder.read(&mut buf).unwrap();
assert_eq!(ret, 0);
}
}