sframe 2.0.0

use crate::{
    crypto::{
        aead::AeadDecrypt,
        buffer::{AadData, decryption::DecryptionBuffer},
        key_derivation::KeyDerivation,
    },
    error::{Result, SframeError},
    header::{KeyId, SframeHeader},
    key::KeyStore,
};

use super::{
    FrameBuffer,
    media_frame::{MediaFrame, MediaFrameView},
    validation::{FrameValidation, UnvalidatedFrame},
};
/// A view on a buffer which contains an encrypted frame in the format as of [RFC 9605 4.2](https://www.rfc-editor.org/rfc/rfc9605.html#section-4.2).
/// The frame is assumed to be stored in the buffer as follows:
/// ```txt
/// | Meta Data | Sframe Header | Encrypted Data | Auth Tag |
/// ````
/// where the  meta data is optional and can e.g. be a media header
#[derive(Clone, Copy, Debug)]
pub struct EncryptedFrameView<'buf> {
    header: SframeHeader,
    meta_data: &'buf [u8],
    header_buf: &'buf [u8],
    cipher_text: &'buf [u8],
}

impl<'ibuf> EncryptedFrameView<'ibuf> {
    /// Tries to create a new view on a buffer deserializing the contained [`SframeHeader`].
    /// Fails with an [`crate::error::SframeError`] if the buffer/header is invalid
    pub fn try_new<D>(data: &'ibuf D) -> Result<Self>
    where
        D: AsRef<[u8]> + ?Sized,
    {
        EncryptedFrameView::try_with_meta_data(data, &[])
    }

    /// Tries to create a new view on a buffer deserializing the contained [`SframeHeader`].
    /// Associates the provided meta data with the frame.
    /// Fails with an [`crate::error::SframeError`] if the buffer is invalid
    pub fn try_with_meta_data<D, M>(data: &'ibuf D, meta_data: &'ibuf M) -> Result<Self>
    where
        D: AsRef<[u8]> + ?Sized,
        M: AsRef<[u8]> + ?Sized,
    {
        let header = SframeHeader::deserialize(data)?;
        log::trace!("EncryptedFrame # {} with header {header}", header.counter());

        let (header_buf, cipher_text) = data.as_ref().split_at(header.len());
        Ok(Self {
            header,
            meta_data: meta_data.as_ref(),
            header_buf,
            cipher_text,
        })
    }

    pub(super) fn with_header<D, M>(
        header: SframeHeader,
        data: &'ibuf D,
        meta_data: &'ibuf M,
    ) -> Self
    where
        D: AsRef<[u8]> + ?Sized,
        M: AsRef<[u8]> + ?Sized,
    {
        let (header_buf, cipher_text) = data.as_ref().split_at(header.len());
        Self {
            header,
            meta_data: meta_data.as_ref(),
            header_buf,
            cipher_text,
        }
    }

    /// the header of the encrypted frame
    pub fn header(&self) -> &SframeHeader {
        &self.header
    }

    /// the associated meta data of the encrypted frame
    pub fn meta_data(&self) -> &[u8] {
        self.meta_data
    }

    /// the cipher text (Sframe Header + Encrypted Data + Authentication Tag) of the encrypted frame
    pub fn cipher_text(&self) -> &[u8] {
        self.cipher_text
    }

    /// Tries to decrypt the encrypted frame with a key from the provided key store.
    /// As [`crate::key::GenericDecryptionKey`] implements [`KeyStore`] this can also be a single key.
    /// Dynamically allocates memory for the resulting [`MediaFrame`]
    /// returns an [`crate::error::SframeError`] if no matching key with the key id in this [`SframeHeader`] is available
    /// or if decryption has failed in general.
    pub fn decrypt<A, D>(&self, key_store: impl KeyStore<A, D>) -> Result<MediaFrame>
    where
        A: AeadDecrypt<Secret = D::Secret>,
        D: KeyDerivation,
    {
        let mut buffer = Vec::new();
        let view = self.decrypt_into(key_store, &mut buffer)?;

        Ok(MediaFrame::with_buffer(
            view.counter(),
            buffer,
            self.meta_data.len(),
        ))
    }

    /// Tries to validate the encrypted frame with the provided [`FrameValidation`] and to decrypt
    /// it with a key from the provided key store.
    /// As [`crate::key::GenericDecryptionKey`] implements [`KeyStore`] this can also be a single key.
    /// Dynamically allocates memory for the resulting [`MediaFrame`]
    ///
    /// The frame is screened BEFORE it is decrypted and only recorded once decryption
    /// authenticated it, see [`FrameValidation`] on why this is split in two steps.
    /// A frame which does not decrypt is not recorded, thus it may be received again.
    ///
    /// returns an [`crate::error::SframeError::FrameValidationFailed`] if the validator rejected
    /// the frame, or the errors of [`decrypt`](Self::decrypt) if the decryption has failed.
    pub fn validated_decrypt<A, D, V>(
        &self,
        key_store: impl KeyStore<A, D>,
        validator: &mut V,
    ) -> Result<MediaFrame>
    where
        A: AeadDecrypt<Secret = D::Secret>,
        D: KeyDerivation,
        V: FrameValidation,
    {
        self.validated(validator, || self.decrypt(key_store))
    }

    /// Tries to decrypt the encrypted frame with a key from the provided key store and stores the result
    /// into the provided buffer. On success an [`MediaFrameView`] on the buffer is returned.
    /// As [`crate::key::GenericDecryptionKey`] implements [`KeyStore`] this can also be a single key.
    /// returns an [`crate::error::SframeError`] if no matching key with the key id in this [`SframeHeader`] is available
    /// or if decryption has failed in general.
    pub fn decrypt_into<'obuf, A, D>(
        &self,
        mut key_store: impl KeyStore<A, D>,
        buffer: &'obuf mut impl FrameBuffer,
    ) -> Result<MediaFrameView<'obuf>>
    where
        A: AeadDecrypt<Secret = D::Secret>,
        D: KeyDerivation,
    {
        let counter = self.header().counter();
        let key_id = self.header.key_id();

        log::trace!("Trying to decrypt EncryptedFrame # {counter} with KeyId {key_id}");

        let key = key_store.lookup(key_id).map_err(missing_key(key_id))?;
        let dec_key = key.as_ref();

        let mut decryption_buffer = DecryptionBuffer::try_allocate(buffer, self, self.cipher_text)?;

        dec_key.decrypt(&mut decryption_buffer, counter)?;

        let meta_len = self.meta_data.len();
        decryption_buffer.truncate(dec_key.cipher_suite(), meta_len);

        // the frame authenticated, so the store may keep the key it handed out
        key_store.record(key);

        let buffer_slice: &mut [u8] = decryption_buffer.into();
        let (meta_data, payload) = buffer_slice.split_at(meta_len);

        let media_frame = MediaFrameView::with_meta_data_and_ctr(counter, payload, meta_data);
        Ok(media_frame)
    }

    /// Tries to validate the encrypted frame with the provided [`FrameValidation`] and to decrypt
    /// it with a key from the provided key store, storing the result into the provided buffer.
    /// On success an [`MediaFrameView`] on the buffer is returned.
    /// As [`crate::key::GenericDecryptionKey`] implements [`KeyStore`] this can also be a single key.
    ///
    /// The frame is screened BEFORE it is decrypted and only recorded once decryption
    /// authenticated it, see [`FrameValidation`] on why this is split in two steps.
    /// A frame which does not decrypt is not recorded, thus it may be received again.
    ///
    /// returns an [`crate::error::SframeError::FrameValidationFailed`] if the validator rejected
    /// the frame, or the errors of [`decrypt_into`](Self::decrypt_into) if the decryption has failed.
    pub fn validated_decrypt_into<'obuf, A, D, V>(
        &self,
        key_store: impl KeyStore<A, D>,
        buffer: &'obuf mut impl FrameBuffer,
        validator: &mut V,
    ) -> Result<MediaFrameView<'obuf>>
    where
        A: AeadDecrypt<Secret = D::Secret>,
        D: KeyDerivation,
        V: FrameValidation,
    {
        self.validated(validator, || self.decrypt_into(key_store, buffer))
    }

    /// Screens the frame, decrypts it and records it only if that succeeded.
    fn validated<V, T>(&self, validator: &mut V, decrypt: impl FnOnce() -> Result<T>) -> Result<T>
    where
        V: FrameValidation,
    {
        log::trace!("Validating EncryptedFrame # {}", self.header.counter());
        let token = validator
            .screen(UnvalidatedFrame::from(self))
            .map_err(validation_failed)?;
        let decrypted = decrypt()?;
        validator.record(token);

        Ok(decrypted)
    }
}

impl AadData for EncryptedFrameView<'_> {
    fn len(&self) -> usize {
        self.header.len() + self.meta_data.len()
    }

    fn serialize(&self, buffer: &mut [u8]) -> Result<()> {
        let (meta_data, header) = buffer.split_at_mut(self.meta_data.len());
        meta_data.copy_from_slice(self.meta_data);
        header.copy_from_slice(self.header_buf);
        Ok(())
    }
}

impl<'buf> TryFrom<&'buf [u8]> for EncryptedFrameView<'buf> {
    type Error = SframeError;

    fn try_from(data: &'buf [u8]) -> Result<Self> {
        EncryptedFrameView::try_new(data)
    }
}

impl<'buf> TryFrom<&'buf Vec<u8>> for EncryptedFrameView<'buf> {
    type Error = SframeError;

    fn try_from(data: &'buf Vec<u8>) -> Result<Self> {
        EncryptedFrameView::try_new(data)
    }
}

impl<'frame> From<&'frame EncryptedFrameView<'_>> for UnvalidatedFrame<'frame> {
    fn from(frame: &'frame EncryptedFrameView<'_>) -> Self {
        Self::new(frame.header(), frame.meta_data())
    }
}

/// An abstraction of an encrypted frame in the format as of [RFC 9605 4.2](https://www.rfc-editor.org/rfc/rfc9605.html#section-4.2),
/// owing an internal buffer containing the cipher text and optionally associated meta data (e.g. be a media header).
pub struct EncryptedFrame {
    buffer: Vec<u8>,
    header: SframeHeader,
    meta_len: usize,
}

impl EncryptedFrame {
    /// Tries to create a new encrypted frame, copying the data of the buffer and deserializing the contained [`SframeHeader`].
    /// Fails with an [`crate::error::SframeError`] if the buffer/header is invalid.
    pub fn try_new<D>(data: D) -> Result<Self>
    where
        D: AsRef<[u8]>,
    {
        EncryptedFrame::try_with_meta_data(data, [])
    }

    /// Tries to create a new encrypted frame, copying the data of the data and meta data buffer and deserializing the contained [`SframeHeader`].
    /// Fails with an [`crate::error::SframeError`] if the buffer/header is invalid.
    pub fn try_with_meta_data<D, M>(data: D, meta_data: M) -> Result<Self>
    where
        D: AsRef<[u8]>,
        M: AsRef<[u8]>,
    {
        let data = data.as_ref();
        let meta_data = meta_data.as_ref();

        let header = SframeHeader::deserialize(data)?;
        log::trace!(
            "EncryptedFrame # {} with header {}",
            header.counter(),
            header
        );

        let meta_len = meta_data.len();
        let mut buffer = Vec::with_capacity(data.len() + meta_len);
        buffer.extend(meta_data);
        buffer.extend(data);

        Ok(Self {
            buffer,
            header,
            meta_len,
        })
    }
    pub(super) fn from_buffer(buffer: Vec<u8>, header: SframeHeader, meta_len: usize) -> Self {
        EncryptedFrame {
            buffer,
            header,
            meta_len,
        }
    }

    /// the header of the encrypted frame
    pub fn header(&self) -> &SframeHeader {
        &self.header
    }

    /// the associated meta data of the encrypted frame
    pub fn meta_data(&self) -> &[u8] {
        &self.buffer[..self.meta_len]
    }

    /// the cipher text (Sframe Header + Encrypted Data + Authentication Tag) of the encrypted frame
    pub fn cipher_text(&self) -> &[u8] {
        &self.buffer[self.meta_len + self.header.len()..]
    }

    /// Tries to decrypt the encrypted frame with a key from the provided key store.
    /// As [`crate::key::GenericDecryptionKey`] implements [`KeyStore`] this can also be a single key.
    /// Dynamically allocats memory for the resulting [`MediaFrame`]
    /// returns an [`crate::error::SframeError`] if no matching key with the key id in this [`SframeHeader`] is available
    /// or if decryption has failed in general.
    pub fn decrypt<A, D>(&self, key_store: impl KeyStore<A, D>) -> Result<MediaFrame>
    where
        A: AeadDecrypt<Secret = D::Secret>,
        D: KeyDerivation,
    {
        let view = EncryptedFrameView::from(self);
        view.decrypt(key_store)
    }

    /// Tries to validate the encrypted frame with the provided [`FrameValidation`] and to decrypt
    /// it with a key from the provided key store.
    /// As [`crate::key::GenericDecryptionKey`] implements [`KeyStore`] this can also be a single key.
    /// Dynamically allocates memory for the resulting [`MediaFrame`]
    ///
    /// The frame is screened BEFORE it is decrypted and only recorded once decryption
    /// authenticated it, see [`FrameValidation`] on why this is split in two steps.
    /// A frame which does not decrypt is not recorded, thus it may be received again.
    ///
    /// returns an [`crate::error::SframeError::FrameValidationFailed`] if the validator rejected
    /// the frame, or the errors of [`decrypt`](Self::decrypt) if the decryption has failed.
    pub fn validated_decrypt<A, D, V>(
        &self,
        key_store: impl KeyStore<A, D>,
        validator: &mut V,
    ) -> Result<MediaFrame>
    where
        A: AeadDecrypt<Secret = D::Secret>,
        D: KeyDerivation,
        V: FrameValidation,
    {
        let view = EncryptedFrameView::from(self);
        view.validated_decrypt(key_store, validator)
    }

    /// Tries to decrypt the encrypted frame with a key from the provided key store and stores the result
    /// into the provided buffer. On success an [`MediaFrameView`] on the buffer is returned.
    /// As [`crate::key::GenericDecryptionKey`] implements [`KeyStore`] this can also be a single key.
    /// returns an [`crate::error::SframeError`] if no matching key with the key id in this [`SframeHeader`] is available
    /// or if decryption has failed in general.
    pub fn decrypt_into<'obuf, A, D>(
        &self,
        key_store: impl KeyStore<A, D>,
        buffer: &'obuf mut impl FrameBuffer,
    ) -> Result<MediaFrameView<'obuf>>
    where
        A: AeadDecrypt<Secret = D::Secret>,
        D: KeyDerivation,
    {
        let view = EncryptedFrameView::from(self);
        view.decrypt_into(key_store, buffer)
    }

    /// Tries to validate the encrypted frame with the provided [`FrameValidation`] and to decrypt
    /// it with a key from the provided key store, storing the result into the provided buffer.
    /// On success an [`MediaFrameView`] on the buffer is returned.
    /// As [`crate::key::GenericDecryptionKey`] implements [`KeyStore`] this can also be a single key.
    ///
    /// The frame is screened BEFORE it is decrypted and only recorded once decryption
    /// authenticated it, see [`FrameValidation`] on why this is split in two steps.
    /// A frame which does not decrypt is not recorded, thus it may be received again.
    ///
    /// returns an [`crate::error::SframeError::FrameValidationFailed`] if the validator rejected
    /// the frame, or the errors of [`decrypt_into`](Self::decrypt_into) if the decryption has failed.
    pub fn validated_decrypt_into<'obuf, A, D, V>(
        &self,
        key_store: impl KeyStore<A, D>,
        buffer: &'obuf mut impl FrameBuffer,
        validator: &mut V,
    ) -> Result<MediaFrameView<'obuf>>
    where
        A: AeadDecrypt<Secret = D::Secret>,
        D: KeyDerivation,
        V: FrameValidation,
    {
        let view = EncryptedFrameView::from(self);
        view.validated_decrypt_into(key_store, buffer, validator)
    }
}

impl<'buf> From<&'buf EncryptedFrame> for EncryptedFrameView<'buf> {
    fn from(frame: &'buf EncryptedFrame) -> Self {
        EncryptedFrameView::with_header(
            frame.header,
            &frame.buffer[frame.meta_len..],
            frame.meta_data(),
        )
    }
}

impl AsRef<[u8]> for EncryptedFrame {
    fn as_ref(&self) -> &[u8] {
        self.buffer.as_slice()
    }
}

impl TryFrom<&[u8]> for EncryptedFrame {
    type Error = SframeError;

    fn try_from(data: &[u8]) -> Result<Self> {
        EncryptedFrame::try_new(data)
    }
}

impl TryFrom<&Vec<u8>> for EncryptedFrame {
    type Error = SframeError;

    fn try_from(data: &Vec<u8>) -> Result<Self> {
        EncryptedFrame::try_new(data)
    }
}

fn validation_failed<E>(err: E) -> SframeError
where
    E: std::error::Error + Send + Sync + 'static,
{
    SframeError::FrameValidationFailed(Box::new(err))
}

fn missing_key<E>(key_id: KeyId) -> impl FnOnce(E) -> SframeError
where
    E: std::error::Error + Send + Sync + 'static,
{
    move |err| SframeError::MissingDecryptionKey {
        key_id,
        source: Box::new(err),
    }
}

#[cfg(test)]
mod test {
    use super::EncryptedFrameView;
    use crate::{frame::encrypted_frame::EncryptedFrame, header::SframeHeader};
    use pretty_assertions::assert_eq;

    #[test]
    fn new_encrypted_frame_view_with_meta_data() {
        let meta_data = [42u8, 43u8];
        let header = SframeHeader::new(42, 666);
        let header_buf = Vec::from(&header);
        let cipher_text = vec![6u8; 3];
        let data = [header_buf.clone(), cipher_text.clone()].concat();

        let frame_view = EncryptedFrameView::try_with_meta_data(&data, &meta_data).unwrap();

        assert_eq!(frame_view.header(), &header);
        assert_eq!(frame_view.cipher_text(), &cipher_text);
        assert_eq!(frame_view.meta_data(), meta_data);
    }

    #[test]
    fn new_encrypted_frame_with_meta_data() {
        let meta_data = [42u8, 43u8];
        let header = SframeHeader::new(42, 666);
        let header_buf = Vec::from(&header);
        let cipher_text = vec![6u8; 3];
        let data = [header_buf.clone(), cipher_text.clone()].concat();

        let frame = EncryptedFrame::try_with_meta_data(data, meta_data).unwrap();

        assert_eq!(frame.header(), &header);
        assert_eq!(frame.cipher_text(), &cipher_text);
        assert_eq!(frame.meta_data(), meta_data);
    }
}