zvault/src/util/compression.rs

277 lines
8.3 KiB
Rust
Raw Normal View History

2017-03-10 11:43:32 +00:00
use std::ptr;
use std::ffi::{CStr, CString};
2017-03-16 08:42:30 +00:00
use std::io::{self, Write};
2017-03-14 15:07:52 +00:00
use std::str::FromStr;
2017-03-10 11:43:32 +00:00
2017-03-23 14:58:49 +00:00
use libc;
2017-03-10 11:43:32 +00:00
use squash::*;
2017-03-16 08:42:30 +00:00
quick_error!{
#[derive(Debug)]
pub enum CompressionError {
UnsupportedCodec(name: String) {
description("Unsupported codec")
display("Unsupported codec: {}", name)
}
InitializeCodec {
description("Failed to initialize codec")
}
InitializeOptions {
description("Failed to set codec options")
}
InitializeStream {
description("Failed to create stream")
}
Operation(reason: &'static str) {
description("Operation failed")
display("Operation failed: {}", reason)
}
Output(err: io::Error) {
from()
cause(err)
description("Failed to write to output")
}
}
}
2017-04-02 16:55:53 +00:00
#[derive(Clone, Debug, Copy, Eq, PartialEq)]
2017-04-03 12:05:16 +00:00
pub enum CompressionMethod {
2017-03-17 10:42:17 +00:00
Deflate, // Standardized
Brotli, // Good speed and ratio
2017-04-03 12:05:16 +00:00
Lzma, // Very good ratio, slow
2017-03-17 10:42:17 +00:00
Lz4 // Very fast, low ratio
}
2017-04-03 12:05:16 +00:00
serde_impl!(CompressionMethod(u8) {
2017-03-17 10:42:17 +00:00
Deflate => 0,
Brotli => 1,
2017-04-03 12:05:16 +00:00
Lzma => 2,
2017-03-17 10:42:17 +00:00
Lz4 => 3
});
2017-03-16 08:42:30 +00:00
2017-04-02 16:55:53 +00:00
#[derive(Clone, Debug, Eq, PartialEq)]
2017-03-17 10:42:17 +00:00
pub struct Compression {
2017-04-03 12:05:16 +00:00
method: CompressionMethod,
2017-03-17 10:42:17 +00:00
level: u8
}
impl Default for Compression {
fn default() -> Self {
2017-04-03 12:05:16 +00:00
Compression { method: CompressionMethod::Brotli, level: 3 }
2017-03-17 10:42:17 +00:00
}
2017-03-10 11:43:32 +00:00
}
serde_impl!(Compression(u64) {
2017-04-03 12:05:16 +00:00
method: CompressionMethod => 0,
2017-03-17 10:42:17 +00:00
level: u8 => 1
2017-03-10 11:43:32 +00:00
});
impl Compression {
2017-03-14 15:07:52 +00:00
#[inline]
pub fn to_string(&self) -> String {
2017-03-17 10:42:17 +00:00
format!("{}/{}", self.name(), self.level)
2017-03-14 15:07:52 +00:00
}
2017-03-16 08:42:30 +00:00
pub fn from_string(name: &str) -> Result<Self, CompressionError> {
2017-03-15 07:27:27 +00:00
let (name, level) = if let Some(pos) = name.find('/') {
2017-03-16 08:42:30 +00:00
let level = try!(u8::from_str(&name[pos+1..]).map_err(|_| CompressionError::UnsupportedCodec(name.to_string())));
2017-03-14 15:07:52 +00:00
let name = &name[..pos];
(name, level)
} else {
(name, 5)
};
2017-04-03 12:05:16 +00:00
let method = match name {
"deflate" | "zlib" | "gzip" => CompressionMethod::Deflate,
"brotli" => CompressionMethod::Brotli,
"lzma" | "lzma2" | "xz" => CompressionMethod::Lzma,
"lz4" => CompressionMethod::Lz4,
2017-03-17 10:42:17 +00:00
_ => return Err(CompressionError::UnsupportedCodec(name.to_string()))
};
2017-04-03 12:05:16 +00:00
Ok(Compression { method: method, level: level })
2017-03-14 15:07:52 +00:00
}
2017-03-10 11:43:32 +00:00
pub fn name(&self) -> &'static str {
2017-04-03 12:05:16 +00:00
match self.method {
CompressionMethod::Deflate => "deflate",
CompressionMethod::Brotli => "brotli",
CompressionMethod::Lzma => "lzma",
CompressionMethod::Lz4 => "lz4",
2017-03-10 11:43:32 +00:00
}
}
2017-03-16 08:42:30 +00:00
fn codec(&self) -> Result<*mut SquashCodec, CompressionError> {
2017-03-10 11:43:32 +00:00
let name = CString::new(self.name().as_bytes()).unwrap();
let codec = unsafe { squash_get_codec(name.as_ptr()) };
if codec.is_null() {
2017-03-16 08:42:30 +00:00
return Err(CompressionError::InitializeCodec)
2017-03-10 11:43:32 +00:00
}
Ok(codec)
}
#[inline]
2017-03-17 10:42:17 +00:00
pub fn level(&self) -> u8 {
self.level
2017-03-10 11:43:32 +00:00
}
2017-03-16 08:42:30 +00:00
fn options(&self) -> Result<*mut SquashOptions, CompressionError> {
2017-03-10 11:43:32 +00:00
let codec = try!(self.codec());
let options = unsafe { squash_options_new(codec, ptr::null::<()>()) };
2017-03-17 10:42:17 +00:00
if options.is_null() {
return Err(CompressionError::InitializeOptions)
}
let option = CString::new("level");
let value = CString::new(format!("{}", self.level));
let res = unsafe { squash_options_parse_option(
options,
option.unwrap().as_ptr(),
value.unwrap().as_ptr()
)};
if res != SQUASH_OK {
//panic!(unsafe { CStr::from_ptr(squash_status_to_string(res)).to_str().unwrap() });
return Err(CompressionError::InitializeOptions)
2017-03-10 11:43:32 +00:00
}
Ok(options)
}
#[inline]
2017-03-16 08:42:30 +00:00
fn error(code: SquashStatus) -> CompressionError {
CompressionError::Operation(unsafe { CStr::from_ptr(squash_status_to_string(code)).to_str().unwrap() })
2017-03-10 11:43:32 +00:00
}
2017-03-16 08:42:30 +00:00
pub fn compress(&self, data: &[u8]) -> Result<Vec<u8>, CompressionError> {
2017-03-10 11:43:32 +00:00
let codec = try!(self.codec());
let options = try!(self.options());
let mut size = data.len() * 2 + 500;
// The following does not work for all codecs
/*unsafe { squash_codec_get_max_compressed_size(
codec,
data.len() as usize
)};*/
let mut buf = Vec::with_capacity(size as usize);
let res = unsafe { squash_codec_compress_with_options(
codec,
&mut size,
buf.as_mut_ptr(),
data.len(),
data.as_ptr(),
options)
};
if res != SQUASH_OK {
println!("{:?}", data);
println!("{}, {}", data.len(), size);
return Err(Self::error(res))
}
unsafe { buf.set_len(size) };
Ok(buf)
}
2017-03-16 08:42:30 +00:00
pub fn decompress(&self, data: &[u8]) -> Result<Vec<u8>, CompressionError> {
2017-03-10 11:43:32 +00:00
let codec = try!(self.codec());
let mut size = unsafe { squash_codec_get_uncompressed_size(
codec,
data.len(),
data.as_ptr()
)};
if size == 0 {
size = 100 * data.len();
}
let mut buf = Vec::with_capacity(size);
let res = unsafe { squash_codec_decompress(
codec,
&mut size,
buf.as_mut_ptr(),
data.len(),
data.as_ptr(),
ptr::null_mut::<()>())
};
if res != SQUASH_OK {
return Err(Self::error(res))
}
unsafe { buf.set_len(size) };
Ok(buf)
}
2017-03-16 08:42:30 +00:00
pub fn compress_stream(&self) -> Result<CompressionStream, CompressionError> {
2017-03-10 11:43:32 +00:00
let codec = try!(self.codec());
let options = try!(self.options());
let stream = unsafe { squash_stream_new_with_options(
codec, SQUASH_STREAM_COMPRESS, options
) };
if stream.is_null() {
2017-03-16 08:42:30 +00:00
return Err(CompressionError::InitializeStream);
2017-03-10 11:43:32 +00:00
}
2017-03-23 14:58:49 +00:00
Ok(CompressionStream::new(stream))
2017-03-10 11:43:32 +00:00
}
2017-03-16 08:42:30 +00:00
pub fn decompress_stream(&self) -> Result<CompressionStream, CompressionError> {
2017-03-10 11:43:32 +00:00
let codec = try!(self.codec());
let stream = unsafe { squash_stream_new(
codec, SQUASH_STREAM_DECOMPRESS, ptr::null::<()>()
) };
if stream.is_null() {
2017-03-16 08:42:30 +00:00
return Err(CompressionError::InitializeStream);
2017-03-10 11:43:32 +00:00
}
2017-03-23 14:58:49 +00:00
Ok(CompressionStream::new(stream))
2017-03-10 11:43:32 +00:00
}
}
pub struct CompressionStream {
2017-03-23 14:58:49 +00:00
stream: *mut SquashStream,
2017-03-10 11:43:32 +00:00
buffer: [u8; 16*1024]
}
impl CompressionStream {
#[inline]
2017-03-23 14:58:49 +00:00
fn new(stream: *mut SquashStream) -> Self {
2017-03-10 11:43:32 +00:00
CompressionStream {
stream: stream,
buffer: [0; 16*1024]
}
}
2017-03-16 08:42:30 +00:00
pub fn process<W: Write>(&mut self, input: &[u8], output: &mut W) -> Result<(), CompressionError> {
2017-03-23 14:58:49 +00:00
let stream = unsafe { &mut (*self.stream) };
2017-03-10 11:43:32 +00:00
stream.next_in = input.as_ptr();
stream.avail_in = input.len();
loop {
stream.next_out = self.buffer.as_mut_ptr();
stream.avail_out = self.buffer.len();
let res = unsafe { squash_stream_process(stream) };
if res < 0 {
return Err(Compression::error(res))
}
let output_size = self.buffer.len() - stream.avail_out;
2017-03-16 08:42:30 +00:00
try!(output.write_all(&self.buffer[..output_size]));
2017-03-10 11:43:32 +00:00
if res != SQUASH_PROCESSING {
break
}
}
Ok(())
}
2017-03-16 08:42:30 +00:00
pub fn finish<W: Write>(mut self, output: &mut W) -> Result<(), CompressionError> {
2017-03-23 14:58:49 +00:00
let stream = unsafe { &mut (*self.stream) };
2017-03-10 11:43:32 +00:00
loop {
stream.next_out = self.buffer.as_mut_ptr();
stream.avail_out = self.buffer.len();
let res = unsafe { squash_stream_finish(stream) };
if res < 0 {
return Err(Compression::error(res))
}
let output_size = self.buffer.len() - stream.avail_out;
2017-03-16 08:42:30 +00:00
try!(output.write_all(&self.buffer[..output_size]));
2017-03-10 11:43:32 +00:00
if res != SQUASH_PROCESSING {
break
}
}
Ok(())
}
}
2017-03-23 14:58:49 +00:00
impl Drop for CompressionStream {
fn drop(&mut self) {
unsafe { squash_object_unref(self.stream as *mut libc::c_void); }
}
}