pretty_good/
compression.rs

1use std::io::prelude::*;
2
3use bzip2::read::{BzDecoder, BzEncoder};
4use failure::Error;
5use flate2::read::{DeflateDecoder, DeflateEncoder, ZlibDecoder, ZlibEncoder};
6use nom::{be_u8, rest};
7use nom::{ErrorKind, IResult};
8use nom::Err as NomErr;
9
10use types::NomError;
11
12named!(uncompressed_data<Result<CompressedDataPacket, Error>>,
13    map!(rest, |x| Ok(CompressedDataPacket::Uncompressed(Vec::from(x))))
14);
15
16named!(zip_data<Result<CompressedDataPacket, Error>>,
17    map!(rest, |x| CompressedDataPacket::from_zip(x))
18);
19
20named!(zlib_data<Result<CompressedDataPacket, Error>>,
21    map!(rest, |x| CompressedDataPacket::from_zlib(x))
22);
23
24named!(bzip2_data<Result<CompressedDataPacket, Error>>,
25    map!(rest, |x| CompressedDataPacket::from_bzip2(x))
26);
27
28named!(
29    compressed_data<Result<CompressedDataPacket, Error>>,
30    switch!(be_u8,
31        0 => call!(uncompressed_data) |
32        1 => call!(zip_data) |
33        2 => call!(zlib_data) |
34        3 => call!(bzip2_data)
35    )
36);
37
38#[derive(Clone, Debug)]
39pub enum CompressedDataPacket {
40    Uncompressed(Vec<u8>),
41    Zip(Vec<u8>),
42    Zlib(Vec<u8>),
43    Bzip2(Vec<u8>),
44}
45
46impl CompressedDataPacket {
47    fn from_zip(data: &[u8]) -> Result<CompressedDataPacket, Error> {
48        let mut deflater = DeflateDecoder::new(data);
49        let mut out = Vec::new();
50        deflater.read_to_end(&mut out)?;
51        Ok(CompressedDataPacket::Zip(out))
52    }
53
54    fn from_zlib(data: &[u8]) -> Result<CompressedDataPacket, Error> {
55        let mut deflater = ZlibDecoder::new(data);
56        let mut out = Vec::new();
57        deflater.read_to_end(&mut out)?;
58        Ok(CompressedDataPacket::Zlib(out))
59    }
60
61    fn from_bzip2(data: &[u8]) -> Result<CompressedDataPacket, Error> {
62        let mut deflater = BzDecoder::new(data);
63        let mut out = Vec::new();
64        deflater.read_to_end(&mut out)?;
65        Ok(CompressedDataPacket::Bzip2(out))
66    }
67
68    pub fn from_bytes(data: &[u8]) -> Result<CompressedDataPacket, Error> {
69        match compressed_data(data) {
70            IResult::Done(_, cdata) => cdata,
71            IResult::Error(NomErr::Code(ErrorKind::Custom(e))) => {
72                let e = NomError::from(e);
73
74                bail!(CompressionError::InvalidFormat {
75                    reason: format!("{:?}", e),
76                })
77            }
78            IResult::Error(e) => bail!(CompressionError::InvalidFormat {
79                reason: format!("{:?}", e),
80            }),
81            IResult::Incomplete(i) => bail!(CompressionError::InvalidFormat {
82                reason: format!("{:?}", i),
83            }),
84        }
85    }
86
87    pub fn to_bytes(&self) -> Result<Vec<u8>, Error> {
88        let mut out = Vec::new();
89
90        let mut deflater: Box<Read> = match self {
91            &CompressedDataPacket::Uncompressed(ref data) => {
92                out.push(0);
93                Box::new(&data[..])
94            }
95            &CompressedDataPacket::Zip(ref data) => {
96                out.push(1);
97                Box::new(DeflateEncoder::new(&data[..], ::flate2::Compression::best()))
98            }
99            &CompressedDataPacket::Zlib(ref data) => {
100                out.push(2);
101                Box::new(ZlibEncoder::new(&data[..], ::flate2::Compression::best()))
102            }
103            &CompressedDataPacket::Bzip2(ref data) => {
104                out.push(3);
105                Box::new(BzEncoder::new(&data[..], ::bzip2::Compression::Best))
106            }
107        };
108
109        deflater.read_to_end(&mut out)?;
110
111        Ok(out)
112    }
113}
114
115#[derive(Debug, Fail)]
116pub enum CompressionError {
117    #[fail(display = "Invalid compressed data: {}", reason)]
118    InvalidFormat { reason: String },
119}