use crate::aead::encrypt;
use crate::dh::DhKeyPair;
use aes_gcm_siv::aead::AeadInPlace;
use aes_gcm_siv::{Aes256GcmSiv, KeyInit, Nonce};
use serde::{Deserialize, Serialize};
use x25519_dalek::PublicKey;
use zeroize::Zeroize;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
#[derive(Serialize, Deserialize, Debug, Zeroize, Clone, PartialEq, Eq)]
#[zeroize(drop)]
pub struct Header {
ad: Vec<u8>,
pub public_key: PublicKey,
pub pn: usize,
pub n: usize,
}
impl Header {
pub fn new(dh_pair: &DhKeyPair, pn: usize, n: usize) -> Self {
Header {
ad: Vec::new(),
public_key: dh_pair.public_key,
pn,
n,
}
}
pub fn concat(&self, ad: &[u8]) -> Vec<u8> {
let mut header = self.clone();
header.ad = ad.to_vec();
postcard::to_allocvec(&header).expect("Failed to serialize Header")
}
pub fn encrypt(&self, hk: &[u8; 32], ad: &[u8]) -> EncryptedHeader {
let header_data = self.concat(ad);
let enc_header = encrypt(hk, &header_data, b"");
EncryptedHeader(enc_header.0, enc_header.1)
}
}
impl From<Vec<u8>> for Header {
fn from(d: Vec<u8>) -> Self {
postcard::from_bytes(&d).unwrap()
}
}
impl From<&[u8]> for Header {
fn from(d: &[u8]) -> Self {
postcard::from_bytes(d).unwrap()
}
}
impl From<Header> for Vec<u8> {
fn from(s: Header) -> Self {
s.concat(b"")
}
}
pub struct EncryptedHeader(Vec<u8>, [u8; 12]);
impl EncryptedHeader {
pub fn decrypt(&self, hk: &Option<[u8; 32]>) -> Option<Header> {
let key_d = match hk {
None => return None,
Some(d) => d,
};
let cipher = match Aes256GcmSiv::new_from_slice(key_d) {
Ok(v) => v,
Err(_) => return None,
};
let nonce = Nonce::from_slice(&self.1);
let mut buffer = Vec::new();
buffer.extend_from_slice(&self.0);
match cipher.decrypt_in_place(nonce, b"", &mut buffer) {
Ok(_) => {}
Err(_) => return None,
};
Some(Header::from(buffer))
}
}
#[cfg(test)]
pub fn gen_header() -> Header {
let dh_pair = DhKeyPair::new();
let pn = 10;
let n = 50;
Header::new(&dh_pair, pn, n)
}
#[cfg(test)]
mod tests {
use x25519_dalek::PublicKey;
use crate::aead::{decrypt, encrypt};
use crate::header::{gen_header, Header};
use crate::kdf_chain::gen_mk;
#[test]
fn ser_des() {
let ad = b"";
let header = gen_header();
let serialized = header.concat(ad);
let created = Header::from(serialized.as_slice());
assert_eq!(header, created)
}
#[test]
fn enc_header() {
let header = gen_header();
let mk = gen_mk();
let header_data = header.concat(b"");
let data = include_bytes!("aead.rs");
let (encrypted, nonce) = encrypt(&mk, data, &header_data);
let decrypted = decrypt(&mk, &encrypted, &header_data, &nonce);
assert_eq!(decrypted, data.to_vec())
}
#[test]
fn test_eq_header() {
let header1 = gen_header();
let header2 = gen_header();
assert_ne!(header1, header2)
}
#[test]
fn debug_header() {
let header = gen_header();
let _string = alloc::format!("{:?}", header);
}
#[test]
fn gen_ex_header() {
let ex_header = Header {
ad: alloc::vec![0],
public_key: PublicKey::from([1; 32]),
pn: 0,
n: 0,
};
let _string = alloc::format!("{:?}", ex_header);
}
#[test]
fn dec_header() {
let header = gen_header();
let encrypted = header.encrypt(&[0; 32], &[0]);
let decrypted = encrypted.decrypt(&Some([1u8; 32]));
assert_eq!(None, decrypted)
}
}