use crate::cipher::aes::encrypt_block;
use crate::cipher::block;
use crate::cipher::chacha::encrypt_block_chacha20;
use crate::cipher::key::SymmetricCipherKey;
use crate::hkdf::KeyType;
use crate::{derive_debug_via_id, error, hkdf};
pub struct HeaderProtectionKey {
inner: SymmetricCipherKey,
algorithm: &'static Algorithm,
}
impl From<hkdf::Okm<'_, &'static Algorithm>> for HeaderProtectionKey {
fn from(okm: hkdf::Okm<&'static Algorithm>) -> Self {
let mut key_bytes = [0; super::MAX_KEY_LEN];
let algorithm = *okm.len();
let key_bytes = &mut key_bytes[..algorithm.key_len()];
okm.fill(key_bytes).unwrap();
Self::new(algorithm, key_bytes).unwrap()
}
}
impl HeaderProtectionKey {
pub fn new(
algorithm: &'static Algorithm,
key_bytes: &[u8],
) -> Result<Self, error::Unspecified> {
Ok(Self {
inner: (algorithm.init)(key_bytes)?,
algorithm,
})
}
#[inline]
pub fn new_mask(&self, sample: &[u8]) -> Result<[u8; 5], error::Unspecified> {
let sample = <&[u8; SAMPLE_LEN]>::try_from(sample)?;
cipher_new_mask(&self.inner, *sample)
}
#[inline]
#[must_use]
pub fn algorithm(&self) -> &'static Algorithm {
self.algorithm
}
}
const SAMPLE_LEN: usize = super::TAG_LEN;
pub type Sample = [u8; SAMPLE_LEN];
pub struct Algorithm {
init: fn(key: &[u8]) -> Result<SymmetricCipherKey, error::Unspecified>,
key_len: usize,
id: AlgorithmID,
}
impl KeyType for &'static Algorithm {
#[inline]
fn len(&self) -> usize {
self.key_len()
}
}
impl Algorithm {
#[inline]
#[must_use]
pub fn key_len(&self) -> usize {
self.key_len
}
#[inline]
#[must_use]
pub fn sample_len(&self) -> usize {
SAMPLE_LEN
}
}
derive_debug_via_id!(Algorithm);
#[derive(Debug, Eq, PartialEq)]
#[allow(non_camel_case_types)]
enum AlgorithmID {
AES_128,
AES_256,
CHACHA20,
}
impl PartialEq for Algorithm {
#[inline]
fn eq(&self, other: &Self) -> bool {
self.id == other.id
}
}
impl Eq for Algorithm {}
pub static AES_128: Algorithm = Algorithm {
key_len: 16,
init: SymmetricCipherKey::aes128,
id: AlgorithmID::AES_128,
};
pub static AES_256: Algorithm = Algorithm {
key_len: 32,
init: SymmetricCipherKey::aes256,
id: AlgorithmID::AES_256,
};
pub static CHACHA20: Algorithm = Algorithm {
key_len: 32,
init: SymmetricCipherKey::chacha20,
id: AlgorithmID::CHACHA20,
};
#[inline]
fn cipher_new_mask(
cipher_key: &SymmetricCipherKey,
sample: Sample,
) -> Result<[u8; 5], error::Unspecified> {
let block = block::Block::from(sample);
let encrypted_block = match cipher_key {
SymmetricCipherKey::Aes128 { enc_key, .. } | SymmetricCipherKey::Aes256 { enc_key, .. } => {
encrypt_block(enc_key, block)
}
SymmetricCipherKey::ChaCha20 { raw_key } => {
let plaintext = block.as_ref();
let counter_bytes: &[u8; 4] = plaintext[0..=3]
.try_into()
.map_err(|_| error::Unspecified)?;
let nonce: &[u8; 12] = plaintext[4..=15]
.try_into()
.map_err(|_| error::Unspecified)?;
let input = block::Block::zero();
unsafe {
let counter = core::mem::transmute::<[u8; 4], u32>(*counter_bytes).to_le();
encrypt_block_chacha20(raw_key, input, nonce, counter)?
}
}
};
let mut out: [u8; 5] = [0; 5];
out.copy_from_slice(&encrypted_block.as_ref()[..5]);
Ok(out)
}
#[cfg(test)]
mod test {
use crate::aead::quic::{Algorithm, HeaderProtectionKey};
use crate::test;
#[test]
fn test_types() {
test::compile_time_assert_send::<Algorithm>();
test::compile_time_assert_sync::<Algorithm>();
test::compile_time_assert_send::<HeaderProtectionKey>();
test::compile_time_assert_sync::<HeaderProtectionKey>();
}
}