src/security/lib/scrutiny/utils/src/zstd.rs

// Copyright 2020 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use anyhow::{Result, anyhow};
use libc::{c_char, c_int, c_void, size_t};
use std::ffi::CStr;

// "C" externs provided from the zstd third_party library.
unsafe extern "C" {
    fn ZSTD_compress(
        dst: *mut c_void,
        dst_capacity: size_t,
        src: *const c_void,
        src_size: size_t,
        compression_level: c_int,
    ) -> size_t;
    fn ZSTD_decompress(
        dst: *mut c_void,
        dst_capacity: size_t,
        src: *const c_void,
        compressed_size: size_t,
    ) -> size_t;
    fn ZSTD_isError(code: size_t) -> u32;
    fn ZSTD_getErrorName(code: size_t) -> *const c_char;
}

// "C" externs for ffi compression functions.
unsafe extern "C" {
    fn zstd_chunked_decompress(
        src: *const c_void,
        src_len: size_t,
        dst: *mut c_void,
        dst_capacity: size_t,
    ) -> size_t;
}

/// Attempts to compress the `src` buffer returning a Vec<u8> with a capacity which
/// is at most `dst_capacity`.
pub fn compress(src: &[u8], dst_capacity: u32, level: i32) -> Result<Vec<u8>> {
    unsafe {
        let src_len = src.len() as size_t;
        let src_ptr = src.as_ptr() as *const c_void;
        let dst_len = dst_capacity as size_t;
        let mut dst = Vec::with_capacity(dst_capacity as usize);
        let dst_ptr = dst.as_mut_ptr() as *mut c_void;
        let compression_level = level as c_int;

        let code = ZSTD_compress(dst_ptr, dst_len, src_ptr, src_len, compression_level);

        if ZSTD_isError(code) != 0 {
            let error: *const c_char = ZSTD_getErrorName(code);
            let error_str = CStr::from_ptr(error).to_str().unwrap().to_owned();
            Err(anyhow!(error_str))
        } else {
            dst.set_len(code as usize);
            Ok(dst)
        }
    }
}

/// Attempts to decompress the `src` buffer returning a Vec<u8> with a capacity which
/// is at most `dst_capacity`. This is intended to be used when the dst_capacity
/// can be predicted (such as through the zbi.extra flag).
pub fn decompress(src: &[u8], dst_capacity: u32) -> Result<Vec<u8>> {
    unsafe {
        let src_len = src.len() as size_t;
        let src_ptr = src.as_ptr() as *const c_void;
        let dst_len = dst_capacity as size_t;
        let mut dst = Vec::with_capacity(dst_capacity as usize);
        let dst_ptr = dst.as_mut_ptr() as *mut c_void;

        let code = ZSTD_decompress(dst_ptr, dst_len, src_ptr, src_len);

        if ZSTD_isError(code) != 0 {
            let error: *const c_char = ZSTD_getErrorName(code);
            let error_str = CStr::from_ptr(error).to_str().unwrap().to_owned();
            Err(anyhow!(error_str))
        } else {
            dst.set_len(code as usize);
            Ok(dst)
        }
    }
}

/// Attempts to use zstd_chunk decompression on the provided source block into
/// a buffer that is of size `dst_capacity`. Any decompression errors will
/// result in an empty vector. This function acts as a one-shot decompression
/// attempting to decompress all compressed chunks in the provided source data.
pub fn chunked_decompress(src: &[u8], dst_capacity: u32) -> Vec<u8> {
    unsafe {
        let src_len = src.len() as size_t;
        let src_ptr = src.as_ptr() as *const c_void;
        let dst_len = dst_capacity as size_t;
        let mut dst = Vec::with_capacity(dst_capacity as usize);
        let dst_ptr = dst.as_mut_ptr() as *mut c_void;

        let bytes_written = zstd_chunked_decompress(src_ptr, src_len, dst_ptr, dst_len);
        dst.set_len(bytes_written as usize);
        dst
    }
}

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

    #[test]
    fn test_compress_decompress() {
        let dst_capacity = 100;
        let to_compress = hex::decode("aaaabbbbcccc").unwrap();
        let compressed = compress(&to_compress, dst_capacity, 3).unwrap();
        let decompressed = decompress(&compressed, dst_capacity).unwrap();
        assert_eq!(decompressed, to_compress);
    }

    #[test]
    fn test_chunked_decompress_invalid_data() {
        let invalid_data = vec![0u8; 128];
        assert_eq!(chunked_decompress(&invalid_data, 256).len(), 0);
    }
}