rustls-ccm 0.2.0

use std::marker::PhantomData;

use ccm::aead::array::Array;
use ccm::aead::{AeadInOut, KeyInit, Tag};
use rustls::crypto::cipher::{
    AeadKey, InboundOpaqueMessage, InboundPlainMessage, Iv, KeyBlockShape, MessageDecrypter,
    MessageEncrypter, Nonce, OutboundOpaqueMessage, OutboundPlainMessage, PrefixedPayload,
    Tls12AeadAlgorithm, UnsupportedOperationError, make_tls12_aad,
};
use rustls::{ConnectionTrafficSecrets, Error};

use crate::CcmVariant;

const EXPLICIT_NONCE_LEN: usize = 8;
const MAX_FRAGMENT_LEN: usize = 16384;

pub(crate) struct Tls12CcmAead<V: CcmVariant>(PhantomData<V>);

impl<V: CcmVariant> Tls12CcmAead<V> {
    pub(crate) const NEW: Self = Self(PhantomData);
}

impl<V: CcmVariant> Tls12AeadAlgorithm for Tls12CcmAead<V> {
    fn encrypter(
        &self,
        key: AeadKey,
        write_iv: &[u8],
        explicit: &[u8],
    ) -> Box<dyn MessageEncrypter> {
        let cipher = V::Cipher::new_from_slice(key.as_ref()).unwrap();
        let iv = build_iv(write_iv, explicit);
        Box::new(Tls12CcmEncrypter::<V> {
            cipher,
            iv,
            _v: PhantomData,
        })
    }

    fn decrypter(&self, key: AeadKey, dec_iv: &[u8]) -> Box<dyn MessageDecrypter> {
        let cipher = V::Cipher::new_from_slice(key.as_ref()).unwrap();
        let mut salt = [0u8; 4];
        debug_assert_eq!(dec_iv.len(), 4);
        salt.copy_from_slice(dec_iv);
        Box::new(Tls12CcmDecrypter::<V> {
            cipher,
            salt,
            _v: PhantomData,
        })
    }

    fn key_block_shape(&self) -> KeyBlockShape {
        KeyBlockShape {
            enc_key_len: V::KEY_LEN,
            fixed_iv_len: 4,
            explicit_nonce_len: 8,
        }
    }

    fn extract_keys(
        &self,
        _key: AeadKey,
        _iv: &[u8],
        _explicit: &[u8],
    ) -> Result<ConnectionTrafficSecrets, UnsupportedOperationError> {
        Err(UnsupportedOperationError)
    }
}

struct Tls12CcmEncrypter<V: CcmVariant> {
    cipher: V::Cipher,
    iv: Iv,
    _v: PhantomData<V>,
}

impl<V: CcmVariant> MessageEncrypter for Tls12CcmEncrypter<V> {
    fn encrypt(
        &mut self,
        msg: OutboundPlainMessage<'_>,
        seq: u64,
    ) -> Result<OutboundOpaqueMessage, Error> {
        let total_len = self.encrypted_payload_len(msg.payload.len());
        let mut payload = PrefixedPayload::with_capacity(total_len);

        let nonce = Nonce::new(&self.iv, seq);
        let aad = make_tls12_aad(seq, msg.typ, msg.version, msg.payload.len());

        // Wire format: [explicit_nonce:8][ciphertext][tag:TAG_LEN]
        payload.extend_from_slice(&nonce.0[4..]);
        payload.extend_from_chunks(&msg.payload);

        let ccm_nonce = Array::from(nonce.0);
        let tag = self
            .cipher
            .encrypt_inout_detached(
                &ccm_nonce,
                &aad,
                (&mut payload.as_mut()[EXPLICIT_NONCE_LEN..]).into(),
            )
            .map_err(|_| Error::EncryptError)?;
        payload.extend_from_slice(tag.as_slice());

        Ok(OutboundOpaqueMessage::new(msg.typ, msg.version, payload))
    }

    fn encrypted_payload_len(&self, payload_len: usize) -> usize {
        payload_len + EXPLICIT_NONCE_LEN + V::TAG_LEN
    }
}

struct Tls12CcmDecrypter<V: CcmVariant> {
    cipher: V::Cipher,
    salt: [u8; 4],
    _v: PhantomData<V>,
}

impl<V: CcmVariant> MessageDecrypter for Tls12CcmDecrypter<V> {
    fn decrypt<'a>(
        &mut self,
        mut msg: InboundOpaqueMessage<'a>,
        seq: u64,
    ) -> Result<InboundPlainMessage<'a>, Error> {
        let overhead = EXPLICIT_NONCE_LEN + V::TAG_LEN;
        let payload = &msg.payload;
        if payload.len() < overhead {
            return Err(Error::DecryptError);
        }

        let mut nonce = [0u8; 12];
        nonce[..4].copy_from_slice(&self.salt);
        nonce[4..].copy_from_slice(&payload[..EXPLICIT_NONCE_LEN]);

        let plain_len = payload.len() - overhead;
        let aad = make_tls12_aad(seq, msg.typ, msg.version, plain_len);

        let payload = &mut msg.payload;
        let ciphertext_end = EXPLICIT_NONCE_LEN + plain_len;

        let tag = Tag::<V::Cipher>::try_from(&payload[ciphertext_end..ciphertext_end + V::TAG_LEN])
            .map_err(|_| Error::DecryptError)?;

        let ccm_nonce = Array::from(nonce);
        self.cipher
            .decrypt_inout_detached(
                &ccm_nonce,
                &aad,
                (&mut payload[EXPLICIT_NONCE_LEN..ciphertext_end]).into(),
                &tag,
            )
            .map_err(|_| Error::DecryptError)?;

        if plain_len > MAX_FRAGMENT_LEN {
            return Err(Error::PeerSentOversizedRecord);
        }

        Ok(msg.into_plain_message_range(EXPLICIT_NONCE_LEN..ciphertext_end))
    }
}

fn build_iv(write_iv: &[u8], explicit: &[u8]) -> Iv {
    debug_assert_eq!(write_iv.len(), 4);
    debug_assert_eq!(explicit.len(), 8);
    let mut iv = [0u8; 12];
    iv[..4].copy_from_slice(write_iv);
    iv[4..].copy_from_slice(explicit);
    Iv::new(iv)
}