use super::cipher::Cipher;
use super::error::HandshakeError;
use crate::zeroize::zeroize_array;
use std::marker::PhantomData;
pub(crate) const MAX_MESSAGE_LEN: usize = 65535;
pub struct CipherState<Ci: Cipher> {
k: Option<[u8; 32]>,
n: u64,
_cipher: PhantomData<Ci>,
}
impl<Ci: Cipher> CipherState<Ci> {
pub fn empty() -> Self {
Self {
k: None,
n: 0,
_cipher: PhantomData,
}
}
pub(crate) fn from_key(key: [u8; 32]) -> Self {
Self {
k: Some(key),
n: 0,
_cipher: PhantomData,
}
}
pub fn has_key(&self) -> bool {
self.k.is_some()
}
#[cfg(test)]
pub(crate) fn set_nonce_for_test(&mut self, n: u64) {
self.n = n;
}
pub fn encrypt_with_ad(
&mut self,
ad: &[u8],
plaintext: &[u8],
output: &mut [u8],
) -> Result<usize, HandshakeError> {
match self.k {
None => {
let len = plaintext.len();
if len > MAX_MESSAGE_LEN {
return Err(HandshakeError::MessageTooLong { len });
}
output[..len].copy_from_slice(plaintext);
Ok(len)
}
Some(ref key) => {
if self.n == u64::MAX {
return Err(HandshakeError::NonceOverflow);
}
let msg_len = plaintext.len() + Ci::TAG_SIZE;
if msg_len > MAX_MESSAGE_LEN {
return Err(HandshakeError::MessageTooLong { len: msg_len });
}
let len = Ci::encrypt(key, self.n, ad, plaintext, output)?;
self.n += 1;
Ok(len)
}
}
}
pub fn rekey(&mut self) -> Result<(), HandshakeError> {
let key = self.k.as_ref().ok_or(HandshakeError::RekeyWithoutKey)?;
let zeros = [0u8; 32];
let mut output = [0u8; 48];
Ci::encrypt(key, u64::MAX, &[], &zeros, &mut output)?;
let mut new_key = [0u8; 32];
new_key.copy_from_slice(&output[..32]);
if let Some(ref mut old) = self.k {
zeroize_array(old);
}
self.k = Some(new_key);
zeroize_array(&mut new_key);
Ok(())
}
pub fn decrypt_with_ad(
&mut self,
ad: &[u8],
ciphertext: &[u8],
output: &mut [u8],
) -> Result<usize, HandshakeError> {
if ciphertext.len() > MAX_MESSAGE_LEN {
return Err(HandshakeError::MessageTooLong {
len: ciphertext.len(),
});
}
match self.k {
None => {
let len = ciphertext.len();
output[..len].copy_from_slice(ciphertext);
Ok(len)
}
Some(ref key) => {
if self.n == u64::MAX {
return Err(HandshakeError::NonceOverflow);
}
let len = Ci::decrypt(key, self.n, ad, ciphertext, output)?;
self.n += 1;
Ok(len)
}
}
}
}
impl<Ci: Cipher> Drop for CipherState<Ci> {
fn drop(&mut self) {
if let Some(ref mut key) = self.k {
zeroize_array(key);
}
self.n = 0;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::noise::cipher::ChaChaPoly;
type Cs = CipherState<ChaChaPoly>;
#[test]
fn encrypt_overflows_at_max_nonce() {
let mut cs = Cs::from_key([0u8; 32]);
cs.set_nonce_for_test(u64::MAX - 1);
let plaintext = b"x";
let mut out = [0u8; 1 + <ChaChaPoly as Cipher>::TAG_SIZE];
cs.encrypt_with_ad(&[], plaintext, &mut out)
.expect("encrypt at u64::MAX - 1 should succeed");
let err = cs
.encrypt_with_ad(&[], plaintext, &mut out)
.expect_err("encrypt at u64::MAX must overflow");
assert!(matches!(err, HandshakeError::NonceOverflow));
}
#[test]
fn decrypt_guard_fires_at_max_nonce() {
let mut cs = Cs::from_key([0u8; 32]);
cs.set_nonce_for_test(u64::MAX);
let ciphertext = [0u8; <ChaChaPoly as Cipher>::TAG_SIZE];
let mut out = [0u8; 1];
let err = cs
.decrypt_with_ad(&[], &ciphertext, &mut out)
.expect_err("decrypt at u64::MAX must overflow");
assert!(matches!(err, HandshakeError::NonceOverflow));
}
}