You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

245 lines
7.4 KiB

use std::ptr;
use std::ffi::{CStr, CString};
use std::io::Write;
use std::str::FromStr;
use squash::*;
#[derive(Clone, Debug)]
pub enum Compression {
Snappy(()),
Deflate(u8),
Brotli(u8),
Lzma2(u8),
ZStd(u8)
}
serde_impl!(Compression(u64) {
Snappy(()) => 0,
Deflate(u8) => 1,
Brotli(u8) => 2,
Lzma2(u8) => 3,
ZStd(u8) => 4
});
impl Compression {
#[inline]
pub fn to_string(&self) -> String {
if let Some(level) = self.level() {
format!("{}/{}", self.name(), level)
} else {
self.name().to_string()
}
}
#[inline]
pub fn from_string(name: &str) -> Result<Self, &'static str> {
let (name, level) = if let Some(pos) = name.find("/") {
let level = try!(u8::from_str(&name[pos+1..]).map_err(|_| "Level must be a number"));
let name = &name[..pos];
(name, level)
} else {
(name, 5)
};
match name {
"snappy" => Ok(Compression::Snappy(())),
"zstd" => Ok(Compression::ZStd(level)),
"deflate" | "zlib" | "gzip" => Ok(Compression::Deflate(level)),
"brotli" => Ok(Compression::Brotli(level)),
"lzma2" => Ok(Compression::Lzma2(level)),
_ => Err("Unsupported codec")
}
}
#[inline]
pub fn name(&self) -> &'static str {
match *self {
Compression::Snappy(_) => "snappy",
Compression::Deflate(_) => "deflate",
Compression::Brotli(_) => "brotli",
Compression::Lzma2(_) => "lzma2",
Compression::ZStd(_) => "zstd",
}
}
#[inline]
fn codec(&self) -> Result<*mut SquashCodec, &'static str> {
let name = CString::new(self.name().as_bytes()).unwrap();
let codec = unsafe { squash_get_codec(name.as_ptr()) };
if codec.is_null() {
return Err("Unsupported algorithm")
}
Ok(codec)
}
#[inline]
pub fn level(&self) -> Option<u8> {
match *self {
Compression::Snappy(_) => None,
Compression::Deflate(lvl) |
Compression::Brotli(lvl) |
Compression::ZStd(lvl) |
Compression::Lzma2(lvl) => Some(lvl),
}
}
fn options(&self) -> Result<*mut SquashOptions, &'static str> {
let codec = try!(self.codec());
let options = unsafe { squash_options_new(codec, ptr::null::<()>()) };
if let Some(level) = self.level() {
if options.is_null() {
return Err("Algorithm does not support a level")
}
let option = CString::new("level");
let value = CString::new(format!("{}", 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("Failed to set compression level")
}
}
Ok(options)
}
#[inline]
fn error(code: SquashStatus) -> &'static str {
unsafe { CStr::from_ptr(squash_status_to_string(code)).to_str().unwrap() }
}
pub fn compress(&self, data: &[u8]) -> Result<Vec<u8>, &'static str> {
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)
}
pub fn decompress(&self, data: &[u8]) -> Result<Vec<u8>, &'static str> {
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)
}
#[inline]
pub fn compress_stream(&self) -> Result<CompressionStream, &'static str> {
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() {
return Err("Failed to create stream");
}
Ok(CompressionStream::new(unsafe { Box::from_raw(stream) }))
}
#[inline]
pub fn decompress_stream(&self) -> Result<CompressionStream, &'static str> {
let codec = try!(self.codec());
let stream = unsafe { squash_stream_new(
codec, SQUASH_STREAM_DECOMPRESS, ptr::null::<()>()
) };
if stream.is_null() {
return Err("Failed to create stream");
}
Ok(CompressionStream::new(unsafe { Box::from_raw(stream) }))
}
}
pub struct CompressionStream {
stream: Box<SquashStream>,
buffer: [u8; 16*1024]
}
impl CompressionStream {
#[inline]
fn new(stream: Box<SquashStream>) -> Self {
CompressionStream {
stream: stream,
buffer: [0; 16*1024]
}
}
pub fn process<W: Write>(&mut self, input: &[u8], output: &mut W) -> Result<(), &'static str> {
let mut stream = &mut *self.stream;
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;
try!(output.write_all(&self.buffer[..output_size]).map_err(|_| "Failed to write to output"));
if res != SQUASH_PROCESSING {
break
}
}
Ok(())
}
pub fn finish<W: Write>(mut self, output: &mut W) -> Result<(), &'static str> {
let mut stream = &mut *self.stream;
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;
try!(output.write_all(&self.buffer[..output_size]).map_err(|_| "Failed to write to output"));
if res != SQUASH_PROCESSING {
break
}
}
Ok(())
}
}