use alloc::vec::Vec;
use crate::{
Error, ErrorKind, Level, Limits, Progress, Status,
checksum::{Adler32, adler32},
deflate::{Deflater, Flush, Inflater},
};
#[must_use]
pub fn compress(data: &[u8]) -> Vec<u8> {
compress_with(data, Level::DEFAULT)
}
#[must_use]
pub fn compress_with(data: &[u8], level: Level) -> Vec<u8> {
let mut encoder = Encoder::new(level);
let mut out = Vec::with_capacity(data.len() / 2 + 64);
let _ = encoder.compress_vec(data, &mut out, Flush::Finish);
out
}
pub fn decompress(data: &[u8]) -> Result<Vec<u8>, Error> {
decompress_with(data, Limits::new())
}
pub fn decompress_with(data: &[u8], limits: Limits) -> Result<Vec<u8>, Error> {
let mut decoder = Decoder::one_shot(limits);
let mut out = Vec::new();
let reserve = data.len().saturating_mul(4).clamp(64, 1 << 20);
let cap = usize::try_from(limits.max_output()).unwrap_or(usize::MAX);
if out.try_reserve(reserve.min(cap)).is_err() {
return Err(Error::config(ErrorKind::AllocationFailed));
}
let p = decoder.finish_vec(data, &mut out)?;
if p.consumed != data.len() {
return Err(Error::new(ErrorKind::TrailingData, p.consumed as u64));
}
Ok(out)
}
pub(crate) fn header_ok(cmf: u8, flg: u8) -> bool {
cmf & 0x0F == 8 && cmf >> 4 <= 7 && ((u16::from(cmf) << 8) | u16::from(flg)) % 31 == 0
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum State {
Header,
DictId,
Body,
Trailer,
Done,
}
#[derive(Debug)]
pub struct Decoder {
inflater: Inflater,
state: State,
held: [u8; 4],
held_len: usize,
adler: Adler32,
dictionary: Option<Vec<u8>>,
dictionary_id: Option<u32>,
max_window_bits: u8,
consumed_total: u64,
failed: Option<Error>,
}
impl Default for Decoder {
fn default() -> Self {
Self::new()
}
}
impl Decoder {
#[must_use]
pub fn new() -> Self {
Self::from_inflater(Inflater::new(), 15)
}
pub(crate) fn one_shot(limits: Limits) -> Self {
Self::from_inflater(Inflater::without_history(limits), 15)
}
fn from_inflater(inflater: Inflater, max_window_bits: u8) -> Self {
Self {
inflater,
state: State::Header,
held: [0; 4],
held_len: 0,
adler: Adler32::new(),
dictionary: None,
dictionary_id: None,
max_window_bits,
consumed_total: 0,
failed: None,
}
}
pub fn with_window_bits(mut self, bits: u8) -> Result<Self, Error> {
self.inflater = self.inflater.with_window_bits(bits)?;
self.max_window_bits = bits;
Ok(self)
}
#[must_use]
pub fn with_limits(mut self, limits: Limits) -> Self {
self.inflater = self.inflater.with_limits(limits);
self
}
pub fn set_dictionary(&mut self, dictionary: &[u8]) -> Result<(), Error> {
if !matches!(self.state, State::Header | State::DictId) {
return Err(Error::config(ErrorKind::InvalidState));
}
let mut stored = Vec::new();
if stored.try_reserve_exact(dictionary.len()).is_err() {
return Err(Error::config(ErrorKind::AllocationFailed));
}
stored.extend_from_slice(dictionary);
self.dictionary = Some(stored);
Ok(())
}
#[must_use]
pub fn dictionary_id(&self) -> Option<u32> {
self.dictionary_id
}
#[must_use]
pub fn is_done(&self) -> bool {
self.state == State::Done
}
#[must_use]
pub fn total_in(&self) -> u64 {
self.consumed_total
}
#[must_use]
pub fn total_out(&self) -> u64 {
self.inflater.total_out()
}
pub fn reset(&mut self) {
self.inflater.reset();
self.state = State::Header;
self.held_len = 0;
self.adler = Adler32::new();
self.dictionary_id = None;
self.consumed_total = 0;
self.failed = None;
}
#[must_use]
pub fn memory_in_use(&self) -> usize {
self.inflater.memory_in_use() + self.dictionary.as_ref().map_or(0, Vec::capacity)
}
pub fn decompress(&mut self, input: &[u8], output: &mut [u8]) -> Result<Progress, Error> {
self.run(input, Sink::Slice(output), false)
}
pub fn finish(&mut self, input: &[u8], output: &mut [u8]) -> Result<Progress, Error> {
self.run(input, Sink::Slice(output), true)
}
pub fn decompress_vec(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<Progress, Error> {
self.run(input, Sink::Vec(out), false)
}
pub fn finish_vec(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<Progress, Error> {
self.run(input, Sink::Vec(out), true)
}
fn fail(&mut self, kind: ErrorKind, at: usize) -> Error {
let e = Error::new(kind, self.consumed_total + at as u64);
self.failed = Some(e);
e
}
fn run(&mut self, input: &[u8], mut sink: Sink<'_>, last: bool) -> Result<Progress, Error> {
if let Some(e) = self.failed {
return Err(e);
}
let mut ip = 0;
let mut produced = 0;
let status = loop {
match self.state {
State::Header | State::DictId | State::Trailer => {
let need = if self.state == State::Header { 2 } else { 4 };
while self.held_len < need && ip < input.len() {
self.held[self.held_len] = input[ip];
self.held_len += 1;
ip += 1;
}
if self.held_len < need {
break Status::NeedsInput;
}
self.held_len = 0;
if let Err(kind) = self.on_field() {
return Err(self.fail(kind, ip));
}
}
State::Body => {
let r = match &mut sink {
Sink::Slice(out) => self
.inflater
.decompress(&input[ip..], &mut out[produced..])
.inspect(|p| {
self.adler.update(&out[produced..produced + p.produced]);
}),
Sink::Vec(out) => {
let at = out.len();
self.inflater
.decompress_vec(&input[ip..], out)
.inspect(|_| {
self.adler.update(&out[at..]);
})
}
};
let p = match r {
Ok(p) => p,
Err(e) => {
self.failed = Some(e);
return Err(e);
}
};
ip += p.consumed;
produced += p.produced;
match p.status {
Status::Done => self.state = State::Trailer,
Status::NeedsOutput => break Status::NeedsOutput,
_ => break Status::NeedsInput,
}
}
State::Done => break Status::Done,
}
};
self.consumed_total += ip as u64;
if status == Status::NeedsInput && last {
return Err(self.fail(ErrorKind::UnexpectedEnd, 0));
}
Ok(Progress::new(ip, produced, status))
}
fn on_field(&mut self) -> Result<(), ErrorKind> {
match self.state {
State::Header => {
let (cmf, flg) = (self.held[0], self.held[1]);
if !header_ok(cmf, flg) {
return Err(ErrorKind::InvalidHeader);
}
if (cmf >> 4) + 8 > self.max_window_bits {
return Err(ErrorKind::InvalidWindowBits);
}
self.state = if flg & 0x20 != 0 {
State::DictId
} else {
State::Body
};
}
State::DictId => {
let id = u32::from_be_bytes(self.held);
self.dictionary_id = Some(id);
let Some(dict) = &self.dictionary else {
return Err(ErrorKind::DictionaryRequired);
};
if adler32(dict) != id {
return Err(ErrorKind::WrongDictionary);
}
self.inflater.set_dictionary(dict).map_err(|e| e.kind())?;
self.state = State::Body;
}
State::Trailer => {
if u32::from_be_bytes(self.held) != self.adler.finish() {
return Err(ErrorKind::ChecksumMismatch);
}
self.state = State::Done;
}
State::Body | State::Done => {}
}
Ok(())
}
}
pub(crate) enum Sink<'a> {
Slice(&'a mut [u8]),
Vec(&'a mut Vec<u8>),
}
#[derive(Debug)]
pub struct Encoder {
deflater: Deflater,
adler: Adler32,
dictionary_id: Option<u32>,
pending: [u8; 6],
pending_start: usize,
pending_len: usize,
header_written: bool,
trailer_written: bool,
}
impl Default for Encoder {
fn default() -> Self {
Self::new(Level::DEFAULT)
}
}
impl Encoder {
#[must_use]
pub fn new(level: Level) -> Self {
Self {
deflater: Deflater::new(level),
adler: Adler32::new(),
dictionary_id: None,
pending: [0; 6],
pending_start: 0,
pending_len: 0,
header_written: false,
trailer_written: false,
}
}
pub fn with_window_bits(mut self, bits: u8) -> Result<Self, Error> {
self.deflater = self.deflater.with_window_bits(bits)?;
Ok(self)
}
pub fn set_dictionary(&mut self, dictionary: &[u8]) -> Result<(), Error> {
if self.header_written {
return Err(Error::config(ErrorKind::InvalidState));
}
self.deflater.set_dictionary(dictionary)?;
self.dictionary_id = Some(adler32(dictionary));
Ok(())
}
pub fn reset(&mut self) {
self.deflater.reset();
self.adler = Adler32::new();
self.dictionary_id = None;
self.pending_start = 0;
self.pending_len = 0;
self.header_written = false;
self.trailer_written = false;
}
#[must_use]
pub fn total_in(&self) -> u64 {
self.deflater.total_in()
}
#[must_use]
pub fn memory_in_use(&self) -> usize {
self.deflater.memory_in_use()
}
fn header(&self) -> [u8; 2] {
let cmf = ((self.deflater.window_bits() - 8) << 4) | 8;
let level_bits = match self.deflater.level().get() {
0 | 1 => 0,
2..=5 => 1,
6 => 2,
_ => 3,
};
let mut flg = (level_bits << 6)
| if self.dictionary_id.is_some() {
0x20
} else {
0
};
let rem = ((u16::from(cmf) << 8) | u16::from(flg)) % 31;
if rem != 0 {
flg += (31 - rem) as u8;
}
[cmf, flg]
}
fn drain(&mut self, output: &mut [u8], produced: &mut usize) {
let n = self.pending_len.min(output.len() - *produced);
output[*produced..*produced + n]
.copy_from_slice(&self.pending[self.pending_start..self.pending_start + n]);
self.pending_start += n;
self.pending_len -= n;
*produced += n;
}
pub fn compress(
&mut self,
input: &[u8],
output: &mut [u8],
flush: Flush,
) -> Result<Progress, Error> {
let mut produced = 0;
if !self.header_written {
let h = self.header();
self.pending[..2].copy_from_slice(&h);
if let Some(id) = self.dictionary_id {
self.pending[2..6].copy_from_slice(&id.to_be_bytes());
self.pending_len = 6;
} else {
self.pending_len = 2;
}
self.pending_start = 0;
self.header_written = true;
}
self.drain(output, &mut produced);
if self.pending_len > 0 {
return Ok(Progress::new(0, produced, Status::NeedsOutput));
}
if self.trailer_written {
if !input.is_empty() {
return Err(Error::config(ErrorKind::InvalidState));
}
return Ok(Progress::new(0, produced, Status::Done));
}
let p = self
.deflater
.compress(input, &mut output[produced..], flush)?;
self.adler.update(&input[..p.consumed]);
produced += p.produced;
if p.status == Status::Done && flush == Flush::Finish {
self.pending[..4].copy_from_slice(&self.adler.finish().to_be_bytes());
self.pending_start = 0;
self.pending_len = 4;
self.trailer_written = true;
self.drain(output, &mut produced);
let status = if self.pending_len == 0 {
Status::Done
} else {
Status::NeedsOutput
};
return Ok(Progress::new(p.consumed, produced, status));
}
Ok(Progress::new(p.consumed, produced, p.status))
}
pub fn compress_vec(
&mut self,
input: &[u8],
out: &mut Vec<u8>,
flush: Flush,
) -> Result<Progress, Error> {
let start = out.len();
let mut consumed = 0;
loop {
let have = out.len();
let room = (input.len() - consumed) / 2 + 1024;
if out.capacity() - have < room && out.try_reserve(room).is_err() {
return Err(Error::config(ErrorKind::AllocationFailed));
}
let spare = out.capacity() - have;
out.resize(have + spare, 0);
match self.compress(&input[consumed..], &mut out[have..], flush) {
Ok(p) => {
out.truncate(have + p.produced);
consumed += p.consumed;
if p.status != Status::NeedsOutput {
return Ok(Progress::new(consumed, out.len() - start, p.status));
}
}
Err(e) => {
out.truncate(have);
return Err(e);
}
}
}
}
}
#[cfg(test)]
mod tests {
use alloc::vec;
use super::*;
#[test]
fn test_zlib_header_values_match_zlib() {
assert_eq!(&compress_with(b"", Level::NONE)[..2], &[0x78, 0x01]);
assert_eq!(&compress_with(b"", Level::FAST)[..2], &[0x78, 0x01]);
assert_eq!(
&compress_with(b"", Level::new(4).unwrap())[..2],
&[0x78, 0x5E]
);
assert_eq!(&compress(b"")[..2], &[0x78, 0x9C]);
assert_eq!(&compress_with(b"", Level::BEST)[..2], &[0x78, 0xDA]);
}
#[test]
fn test_zlib_known_stream_from_python_zlib() {
let packed = [
0x78, 0x9C, 0xCB, 0x48, 0xCD, 0xC9, 0xC9, 0x07, 0x00, 0x06, 0x2C, 0x02, 0x15,
];
assert_eq!(decompress(&packed).unwrap(), b"hello");
}
#[test]
fn test_zlib_bad_headers_rejected() {
assert_eq!(
decompress(&[0x77, 0x01, 0x03, 0x00]).unwrap_err().kind(),
ErrorKind::InvalidHeader
);
assert_eq!(
decompress(&[0x88, 0x1C, 0x03, 0x00]).unwrap_err().kind(),
ErrorKind::InvalidHeader
);
assert_eq!(
decompress(&[0x78, 0x9D, 0x03, 0x00]).unwrap_err().kind(),
ErrorKind::InvalidHeader
);
}
#[test]
fn test_zlib_dictionary_required_and_wrong() {
let mut enc = Encoder::new(Level::BEST);
enc.set_dictionary(b"the dictionary").unwrap();
let mut packed = Vec::new();
enc.compress_vec(b"the dictionary again", &mut packed, Flush::Finish)
.unwrap();
assert_eq!(
decompress(&packed).unwrap_err().kind(),
ErrorKind::DictionaryRequired
);
let mut dec = Decoder::new();
dec.set_dictionary(b"another one").unwrap();
assert_eq!(
dec.finish_vec(&packed, &mut Vec::new()).unwrap_err().kind(),
ErrorKind::WrongDictionary
);
let mut dec = Decoder::new();
dec.set_dictionary(b"the dictionary").unwrap();
let mut out = Vec::new();
dec.finish_vec(&packed, &mut out).unwrap();
assert_eq!(out, b"the dictionary again");
assert_eq!(dec.dictionary_id(), Some(adler32(b"the dictionary")));
}
#[test]
fn test_zlib_every_split_and_buffer() {
let data: Vec<u8> = (0..3000u32).map(|i| (i % 97) as u8).collect();
let packed = compress(&data);
for step in [1usize, 2, 5, 13] {
for out_size in [1usize, 3, 64] {
let mut dec = Decoder::new();
let mut out = Vec::new();
let mut buf = vec![0u8; out_size];
let mut pos = 0;
loop {
let end = (pos + step).min(packed.len());
let p = if end == packed.len() {
dec.finish(&packed[pos..end], &mut buf).unwrap()
} else {
dec.decompress(&packed[pos..end], &mut buf).unwrap()
};
out.extend_from_slice(&buf[..p.produced]);
pos += p.consumed;
if p.status == Status::Done {
break;
}
}
assert_eq!(out, data, "{step}/{out_size}");
}
}
}
#[test]
fn test_zlib_trailer_and_truncation() {
let packed = compress(b"checked");
for n in 0..packed.len() {
assert_eq!(
decompress(&packed[..n]).unwrap_err().kind(),
ErrorKind::UnexpectedEnd,
"prefix {n}"
);
}
let mut extra = packed.clone();
extra.push(0);
assert_eq!(
decompress(&extra).unwrap_err().kind(),
ErrorKind::TrailingData
);
let mut bad = packed;
let n = bad.len();
bad[n - 2] ^= 0x10;
assert_eq!(
decompress(&bad).unwrap_err().kind(),
ErrorKind::ChecksumMismatch
);
}
}