1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

extern crate libc;

use libc::{c_int, size_t};

#[link(name = "snappy")]
extern {
    fn snappy_compress(input: *const u8,
                       input_length: size_t,
                       compressed: *mut u8,
                       compressed_length: *mut size_t) -> c_int;
    fn snappy_uncompress(compressed: *const u8,
                         compressed_length: size_t,
                         uncompressed: *mut u8,
                         uncompressed_length: *mut size_t) -> c_int;
    fn snappy_max_compressed_length(source_length: size_t) -> size_t;
    fn snappy_uncompressed_length(compressed: *const u8,
                                  compressed_length: size_t,
                                  result: *mut size_t) -> c_int;
    fn snappy_validate_compressed_buffer(compressed: *const u8,
                                         compressed_length: size_t) -> c_int;
}

pub fn validate_compressed_buffer(src: &[u8]) -> bool {
    unsafe {
        snappy_validate_compressed_buffer(src.as_ptr(), src.len() as size_t) == 0
    }
}

pub fn compress(src: &[u8]) -> Vec<u8> {
    unsafe {
        let srclen = src.len() as size_t;
        let psrc = src.as_ptr();

        let mut dstlen = snappy_max_compressed_length(srclen);
        let mut dst = Vec::with_capacity(dstlen as usize);
        let pdst = dst.as_mut_ptr();

        snappy_compress(psrc, srclen, pdst, &mut dstlen);
        dst.set_len(dstlen as usize);
        dst
    }
}

/// Uncompress 'src' into a newly allocated vector.
pub fn uncompress(src: &[u8]) -> Result<Vec<u8>, ()> {
    let mut out = Vec::new();
    uncompress_to(src, &mut out).map(|_| out)
}

/// Uncompress 'src' directly appended to 'dst' and return the number
/// of bytes produced. Return an error upon an invalid 'src'.
pub fn uncompress_to(src: &[u8], dst: &mut Vec<u8>) -> Result<usize, ()> {
    unsafe {
        let src_len = src.len() as size_t;
        let src_ptr = src.as_ptr();

        let dst_cur_len = dst.len();
        let mut dst_add_len: size_t = 0;
        snappy_uncompressed_length(src_ptr, src_len, &mut dst_add_len);

        // now make sure the vector is large enough
        dst.reserve(dst_add_len as usize);
        // try to uncompress
        let dst_ptr = dst[dst_cur_len..].as_mut_ptr();
        if snappy_uncompress(src_ptr, src_len, dst_ptr, &mut dst_add_len) == 0 {
            dst.set_len(dst_cur_len + dst_add_len as usize);
            Ok(dst_add_len as usize)
        } else {
            dst.set_len(dst_cur_len); // make sure not to leak partial output
            Err(()) // SNAPPY_INVALID_INPUT
        }
    }
}

#[cfg(test)]
mod tests {
    use std::str;
    use super::*;

    #[test]
    fn valid() {
        let d = vec![0xde, 0xad, 0xd0, 0x0d];
        let c = compress(&d);
        assert!(validate_compressed_buffer(&c));
        assert!(uncompress(&c) == Ok(d));
    }

    #[test]
    fn invalid() {
        let d = vec![0, 0, 0, 0];
        assert!(!validate_compressed_buffer(&d));
        assert!(uncompress(&d).is_err());
    }

    #[test]
    fn empty() {
        let d = vec![];
        assert!(!validate_compressed_buffer(&d));
        assert!(uncompress(&d).is_err());

        let c = compress(&d);
        assert!(validate_compressed_buffer(&c));
        assert!(uncompress(&c) == Ok(d));
    }

    #[test]
    fn uncompress_to_appends() {
        // "This is test"
        let compressed = &[12, 44, 84, 104, 105, 115, 32, 105, 115, 32, 116, 101, 115, 116];

        let mut out = vec![b'a', b'b', b'c', b'>'];
        uncompress_to(compressed, &mut out).unwrap();
        let s = str::from_utf8(&out[..]).unwrap();
        assert_eq!(s, "abc>This is test");
    }
}