use super::{Compression, Encryption, Key, HMAC};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use std::cmp;
#[derive(Error, Debug)]
pub enum ChunkError {
#[error("Compression Error")]
CompressionError(#[from] super::CompressionError),
#[error("Encryption Error")]
EncryptionError(#[from] super::EncryptionError),
#[error("Key Error")]
KeyError(#[from] super::KeyError),
#[error("HMAC Vailidation Failed")]
HMACValidationFailed,
}
type Result<T> = std::result::Result<T, ChunkError>;
#[derive(PartialEq, Eq, Copy, Clone, Serialize, Deserialize, Hash, Debug)]
pub struct ChunkID {
id: [u8; 32],
}
impl ChunkID {
pub fn new(input_id: &[u8]) -> ChunkID {
let mut id: [u8; 32] = [0; 32];
id[..cmp::min(32, input_id.len())]
.clone_from_slice(&input_id[..cmp::min(32, input_id.len())]);
ChunkID { id }
}
#[cfg_attr(tarpaulin, skip)]
pub fn get_id(&self) -> &[u8] {
&self.id
}
pub fn verify(&self, slice: &[u8]) -> bool {
if slice.len() < self.id.len() {
false
} else {
let mut equal = true;
for (i, val) in self.id.iter().enumerate() {
if *val != slice[i] {
equal = false;
}
}
equal
}
}
pub fn manifest_id() -> ChunkID {
ChunkID { id: [0_u8; 32] }
}
pub fn random_id() -> ChunkID {
let id = rand::random();
ChunkID { id }
}
}
#[derive(Serialize, Deserialize, Clone, Debug, Copy, PartialEq, Eq)]
pub struct ChunkSettings {
pub compression: Compression,
pub encryption: Encryption,
pub hmac: HMAC,
}
impl ChunkSettings {
pub fn lightweight() -> ChunkSettings {
ChunkSettings {
compression: Compression::NoCompression,
encryption: Encryption::NoEncryption,
hmac: HMAC::Blake2b,
}
}
}
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub struct ChunkHeader {
compression: Compression,
encryption: Encryption,
hmac: HMAC,
mac: Vec<u8>,
id: ChunkID,
}
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ChunkBody(pub Vec<u8>);
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Chunk {
#[serde(with = "serde_bytes")]
data: Vec<u8>,
compression: Compression,
encryption: Encryption,
hmac: HMAC,
#[serde(with = "serde_bytes")]
mac: Vec<u8>,
id: ChunkID,
}
impl Chunk {
pub fn pack(
data: Vec<u8>,
compression: Compression,
encryption: Encryption,
hmac: HMAC,
key: &Key,
) -> Chunk {
let id_mac = hmac.id(&data, key);
let id = ChunkID::new(&id_mac);
Chunk::pack_with_id(data, compression, encryption, hmac, key, id)
}
pub fn from_parts(
data: Vec<u8>,
compression: Compression,
encryption: Encryption,
hmac: HMAC,
mac: Vec<u8>,
id: ChunkID,
) -> Chunk {
Chunk {
data,
compression,
encryption,
hmac,
mac,
id,
}
}
pub fn pack_with_id(
data: Vec<u8>,
compression: Compression,
mut encryption: Encryption,
hmac: HMAC,
key: &Key,
id: ChunkID,
) -> Chunk {
let compressed_data = compression.compress(data);
let data = encryption.encrypt(&compressed_data, key);
let mac = hmac.mac(&data, key);
Chunk {
data,
compression,
encryption,
hmac,
mac,
id,
}
}
pub fn unpack(&self, key: &Key) -> Result<Vec<u8>> {
if self.hmac.verify_hmac(&self.mac, &self.data, key) {
let decrypted_data = self.encryption.decrypt(&self.data, key)?;
let decompressed_data = self.compression.decompress(decrypted_data)?;
Ok(decompressed_data)
} else {
Err(ChunkError::HMACValidationFailed)
}
}
#[cfg_attr(tarpaulin, skip)]
pub fn len(&self) -> usize {
self.data.len()
}
#[cfg_attr(tarpaulin, skip)]
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
#[cfg_attr(tarpaulin, skip)]
pub fn get_bytes(&self) -> &[u8] {
&self.data
}
pub fn get_id(&self) -> ChunkID {
self.id
}
pub fn mac(&self) -> Vec<u8> {
self.mac.clone()
}
pub fn split(self) -> (ChunkHeader, ChunkBody) {
let header = ChunkHeader {
compression: self.compression,
encryption: self.encryption,
hmac: self.hmac,
mac: self.mac,
id: self.id,
};
let body = ChunkBody(self.data);
(header, body)
}
pub fn unsplit(header: ChunkHeader, body: ChunkBody) -> Chunk {
Chunk {
data: body.0,
compression: header.compression,
encryption: header.encryption,
hmac: header.hmac,
mac: header.mac,
id: header.id,
}
}
pub fn encryption(&self) -> Encryption {
self.encryption
}
#[cfg(test)]
#[cfg_attr(tarpaulin, skip)]
pub fn break_data(&mut self, index: usize) {
let val = self.data[index];
if val == 0 {
self.data[index] = 1;
} else {
self.data[index] = 0;
}
}
}
impl PartialEq for Chunk {
fn eq(&self, other: &Self) -> bool {
self.id == other.id && self.data == other.data
}
}
impl Eq for Chunk {}
#[cfg(test)]
mod tests {
use super::*;
fn chunk_with_settings(compression: Compression, encryption: Encryption, hmac: HMAC) {
let data_string =
"Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.";
let data_bytes = data_string.as_bytes().to_vec();
println!("Data: \n:{:X?}", data_bytes);
println!("{:?} {:?} {:?}", compression, hmac, encryption);
let key = Key::random(32);
let packed = Chunk::pack(data_bytes, compression, encryption, hmac, &key);
let output_bytes = packed.unpack(&key).expect("Failed to unpack output bytes");
assert_eq!(data_string.as_bytes().to_vec(), output_bytes);
}
#[test]
fn all_combos() {
let compressions = [
Compression::NoCompression,
Compression::ZStd { level: 1 },
Compression::LZ4 { level: 1 },
Compression::LZMA { level: 1 },
];
let encryptions = [
Encryption::NoEncryption,
Encryption::new_aes256ctr(),
Encryption::new_chacha20(),
];
let hmacs = [
HMAC::SHA256,
HMAC::Blake2b,
HMAC::Blake2bp,
HMAC::Blake3,
HMAC::SHA3,
];
for c in compressions.iter() {
for e in encryptions.iter() {
for h in hmacs.iter() {
chunk_with_settings(*c, *e, *h);
}
}
}
}
#[test]
fn detect_bad_data() {
let data_string = "I am but a humble test string";
let data_bytes = data_string.as_bytes().to_vec();
let compression = Compression::NoCompression;
let encryption = Encryption::NoEncryption;
let hmac = HMAC::SHA256;
let key = Key::random(32);
let mut packed = Chunk::pack(data_bytes, compression, encryption, hmac, &key);
packed.break_data(5);
let result = packed.unpack(&key);
assert!(result.is_err());
}
#[test]
fn chunk_id_equality() {
let data1 = [1_u8; 64];
let data2 = [2_u8; 64];
let id = ChunkID::new(&data1);
assert!(id.verify(&data1));
assert!(!id.verify(&data2));
}
#[test]
fn split_unsplit() {
let data_string = "I am but a humble test string";
let data_bytes = data_string.as_bytes().to_vec();
let compression = Compression::LZ4 { level: 1 };
let encryption = Encryption::new_aes256ctr();
let hmac = HMAC::SHA256;
let key = Key::random(32);
let packed = Chunk::pack(data_bytes, compression, encryption, hmac, &key);
let (header, body) = packed.split();
let packed = Chunk::unsplit(header, body);
let result = packed.unpack(&key);
assert!(result.is_ok());
}
}