titan-api-codec 1.2.9

Helpers for encoding and decoding Titan API messages
Documentation
//! Defines transforms that utilize [brotli] to compress and decompress data.
//!
//! [brotli]: https://github.com/google/brotli

use super::common::BinaryTransform;

use brotli::enc::{BrotliCompress, BrotliCompressCustomAlloc, BrotliEncoderParams, StandardAlloc};
use brotli::{BrotliDecompress, BrotliDecompressCustomAlloc};
use bytes::{Buf, BufMut, Bytes, BytesMut};

const INPUT_BUFFER_SIZE: usize = 4096;
const OUTPUT_BUFFER_SIZE: usize = 4096;

/// Transform that applies brotli compression to the input.
pub struct BrotliCompressor {
    input_buffer: Vec<u8>,
    output_buffer: Vec<u8>,
    params: BrotliEncoderParams,
    alloc: StandardAlloc,
}

impl Default for BrotliCompressor {
    fn default() -> Self {
        Self {
            input_buffer: vec![0u8; INPUT_BUFFER_SIZE],
            output_buffer: vec![0u8; OUTPUT_BUFFER_SIZE],
            params: BrotliEncoderParams::default(),
            alloc: StandardAlloc::default(),
        }
    }
}

impl BinaryTransform for BrotliCompressor {
    fn transform_mut(&mut self, data: Bytes) -> Result<Bytes, std::io::Error> {
        let mut buffer = BytesMut::new().writer();
        BrotliCompressCustomAlloc(
            &mut data.reader(),
            &mut buffer,
            &mut self.input_buffer,
            &mut self.output_buffer,
            &self.params,
            self.alloc,
        )?;
        Ok(buffer.into_inner().into())
    }

    fn transform(&self, data: Bytes) -> Result<Bytes, std::io::Error> {
        let mut buffer = BytesMut::new().writer();
        BrotliCompress(&mut data.reader(), &mut buffer, &self.params)?;
        Ok(buffer.into_inner().into())
    }
}

/// Transform that decompress brotli-compressed data back to its original content.
pub struct BrotliDecompressor {
    input_buffer: Vec<u8>,
    output_buffer: Vec<u8>,
    alloc: StandardAlloc,
}

impl Default for BrotliDecompressor {
    fn default() -> Self {
        Self {
            input_buffer: vec![0u8; INPUT_BUFFER_SIZE],
            output_buffer: vec![0u8; OUTPUT_BUFFER_SIZE],
            alloc: StandardAlloc::default(),
        }
    }
}

impl BinaryTransform for BrotliDecompressor {
    fn transform_mut(&mut self, data: Bytes) -> Result<Bytes, std::io::Error> {
        let mut buffer = BytesMut::new().writer();
        BrotliDecompressCustomAlloc(
            &mut data.reader(),
            &mut buffer,
            &mut self.input_buffer,
            &mut self.output_buffer,
            self.alloc,
            self.alloc,
            self.alloc,
        )?;
        Ok(buffer.into_inner().into())
    }

    fn transform(&self, data: Bytes) -> Result<Bytes, std::io::Error> {
        let mut buffer = BytesMut::new().writer();
        BrotliDecompress(&mut data.reader(), &mut buffer)?;
        Ok(buffer.into_inner().into())
    }
}

#[cfg(test)]
mod test {
    use super::{BrotliCompressor, BrotliDecompressor};
    use crate::transform::BinaryTransform;
    use bytes::Bytes;
    use lipsum::lipsum;

    #[test]
    fn test_roundtrip_default() {
        let compressor = BrotliCompressor::default();
        let decompressor = BrotliDecompressor::default();

        let data = Bytes::from(lipsum(1000));

        let compressed = compressor
            .transform(data.clone())
            .expect("should compress via brotli");
        let uncompressed = decompressor
            .transform(compressed)
            .expect("should decompress from brotli");

        assert_eq!(data, uncompressed);
    }

    #[test]
    fn test_roundtrip_mut_default() {
        let mut compressor = BrotliCompressor::default();
        let mut decompressor = BrotliDecompressor::default();

        let data = Bytes::from(lipsum(1000));

        let compressed = compressor
            .transform_mut(data.clone())
            .expect("should compress via brotli");
        let uncompressed = decompressor
            .transform_mut(compressed)
            .expect("should decompress from brotli");

        assert_eq!(data, uncompressed);
    }
}