sframe 2.0.0

use crate::{
    crypto::{
        aead::AeadEncrypt,
        buffer::{AadData, encryption::EncryptionBuffer},
        key_derivation::KeyDerivation,
    },
    error::{Result, SframeError},
    header::{Counter, SframeHeader},
    key::GenericEncryptionKey,
};

use super::{
    FrameBuffer, FrameCounter,
    encrypted_frame::{EncryptedFrame, EncryptedFrameView},
};
use std::convert::Infallible;

fn next_counter(frame_counter: &mut impl FrameCounter) -> Result<Counter> {
    frame_counter
        .try_next()
        .map_err(|err| SframeError::CounterCreationFailed(Box::new(err)))
}

/// A view on a buffer (as a continuous slice of memory), representing a media frame.
/// Can optionally have meta data (e.g. a header) associated to it, which is considered for authentication.
#[derive(Debug, PartialEq, Eq)]
pub struct MediaFrameView<'buf> {
    counter: Counter,
    meta_data: &'buf [u8],
    payload: &'buf [u8],
}

impl<'ibuf> MediaFrameView<'ibuf> {
    /// Creates a new view on a payload buffer and assigns it the next counter (CTR) value.
    /// Only available for a [`FrameCounter`] which cannot fail.
    pub fn new<C, P>(frame_counter: &mut C, payload: &'ibuf P) -> Self
    where
        C: FrameCounter<Error = Infallible>,
        P: AsRef<[u8]> + ?Sized,
    {
        Self::with_meta_data(frame_counter, payload, &[])
    }

    /// Creates a new view on a payload buffer, assigns it the next counter (CTR) value and associates it with the meta data.
    /// Only available for a [`FrameCounter`] which cannot fail.
    pub fn with_meta_data<C, P, M>(
        frame_counter: &mut C,
        payload: &'ibuf P,
        meta_data: &'ibuf M,
    ) -> Self
    where
        C: FrameCounter<Error = Infallible>,
        P: AsRef<[u8]> + ?Sized,
        M: AsRef<[u8]> + ?Sized,
    {
        let Ok(counter) = frame_counter.try_next();
        Self::with_meta_data_and_ctr(counter, payload, meta_data)
    }

    /// Creates a new view on a payload buffer and assigns it the next counter (CTR) value.
    /// Returns an [`SframeError::CounterCreationFailed`] wrapping the error of the
    /// [`FrameCounter`] if no counter value could be created.
    pub fn try_new<C, P>(frame_counter: &mut C, payload: &'ibuf P) -> Result<Self>
    where
        C: FrameCounter,
        P: AsRef<[u8]> + ?Sized,
    {
        Self::try_with_meta_data(frame_counter, payload, &[])
    }

    /// Creates a new view on a payload buffer, assigns it the next counter (CTR) value and associates it with the meta data.
    /// Returns an [`SframeError::CounterCreationFailed`] wrapping the error of the
    /// [`FrameCounter`] if no counter value could be created.
    pub fn try_with_meta_data<C, P, M>(
        frame_counter: &mut C,
        payload: &'ibuf P,
        meta_data: &'ibuf M,
    ) -> Result<Self>
    where
        C: FrameCounter,
        P: AsRef<[u8]> + ?Sized,
        M: AsRef<[u8]> + ?Sized,
    {
        let counter = next_counter(frame_counter)?;
        Ok(Self::with_meta_data_and_ctr(counter, payload, meta_data))
    }

    pub(super) fn with_meta_data_and_ctr<P, M>(
        counter: Counter,
        payload: &'ibuf P,
        meta_data: &'ibuf M,
    ) -> Self
    where
        P: AsRef<[u8]> + ?Sized,
        M: AsRef<[u8]> + ?Sized,
    {
        let payload = payload.as_ref();
        let meta_data = meta_data.as_ref();
        log::trace!(
            "Creating MediaFrame # {counter} with payload size {} using meta data of size {}",
            payload.len(),
            meta_data.len(),
        );
        Self {
            counter,
            meta_data,
            payload,
        }
    }

    /// Meta data associated with this media frame
    pub fn meta_data(&self) -> &'ibuf [u8] {
        self.meta_data
    }

    /// Payload of this media frame
    pub fn payload(&self) -> &'ibuf [u8] {
        self.payload
    }

    /// Frame count for the Sframe scheme associated to this media frame
    pub fn counter(&self) -> Counter {
        self.counter
    }

    /// Encrypts the media frame with the sframe key according to [RFC 9605 4.4.3](https://www.rfc-editor.org/rfc/rfc9605.html#name-encryption). Dynamically allocates memory for the resulting [`EncryptedFrame`].
    /// The associated meta data is not encrypted but considered for the authentication tag.
    /// Returns an [`crate::error::SframeError`] when encryption fails.
    pub fn encrypt<A, D>(&self, key: &GenericEncryptionKey<A, D>) -> Result<EncryptedFrame>
    where
        A: AeadEncrypt<Secret = D::Secret>,
        D: KeyDerivation,
    {
        let mut buffer = Vec::new();
        let view = self.encrypt_into(key, &mut buffer)?;

        let header = *view.header();
        let meta_len = view.meta_data().len();

        let encrypted_frame = EncryptedFrame::from_buffer(buffer, header, meta_len);

        Ok(encrypted_frame)
    }

    /// Encrypts the media frame with the sframe key according to [RFC 9605 4.4.3](https://www.rfc-editor.org/rfc/rfc9605.html#name-encryption) and stores the result, an [`EncryptedFrameView`], into the provided buffer.
    /// The associated meta data is not encrypted but considered for the authentication tag.
    /// Returns an [`crate::error::SframeError`] when encryption fails.
    pub fn encrypt_into<'obuf, A, D>(
        &self,
        key: &GenericEncryptionKey<A, D>,
        buffer: &'obuf mut impl FrameBuffer,
    ) -> Result<EncryptedFrameView<'obuf>>
    where
        A: AeadEncrypt<Secret = D::Secret>,
        D: KeyDerivation,
    {
        let key_id = key.key_id();
        log::trace!(
            "Encrypting MediaFrame # {} using KeyId {key_id} and CipherSuite {}",
            self.counter,
            key.cipher_suite()
        );

        let header = SframeHeader::new(key_id, self.counter);
        log::trace!("MediaFrame # {} using header {}", self.counter, header);

        let aad = Aad {
            meta_data: self.meta_data,
            header: &header,
        };
        let mut crypto_buffer =
            EncryptionBuffer::try_allocate(buffer, key.cipher_suite(), &aad, self.payload)?;

        log::trace!("MediaFrame # {} trying to encrypt", self.counter);
        key.encrypt(&mut crypto_buffer, self.counter)?;

        let meta_len = self.meta_data().len();
        let buffer: &mut [u8] = crypto_buffer.into();
        let encrypted =
            EncryptedFrameView::with_header(header, &buffer[meta_len..], &buffer[..meta_len]);

        Ok(encrypted)
    }
}

struct Aad<'a> {
    meta_data: &'a [u8],
    header: &'a SframeHeader,
}

impl AadData for Aad<'_> {
    fn len(&self) -> usize {
        self.meta_data.len() + self.header.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);
        self.header.serialize(header)?;
        Ok(())
    }
}

/// A an abstraction of a media frame owning an internal buffer.
/// Can optionally have meta data (e.g. a header) associated to it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MediaFrame {
    counter: Counter,
    buffer: Vec<u8>,
    meta_len: usize,
}

impl MediaFrame {
    /// Creates a new media frame by copying the data of a payload buffer and assigning it the next counter (CTR) value.
    /// Only available for a [`FrameCounter`] which cannot fail.
    pub fn new<C, P>(frame_counter: &mut C, payload: P) -> Self
    where
        C: FrameCounter<Error = Infallible>,
        P: AsRef<[u8]>,
    {
        Self::with_meta_data(frame_counter, payload, [])
    }

    /// Creates a new media frame and assigns it the next counter (CTR) value.
    /// Payload and meta data are copied into an internal buffer.
    /// Only available for a [`FrameCounter`] which cannot fail.
    pub fn with_meta_data<C, P, M>(frame_counter: &mut C, payload: P, meta_data: M) -> Self
    where
        C: FrameCounter<Error = Infallible>,
        P: AsRef<[u8]>,
        M: AsRef<[u8]>,
    {
        let Ok(counter) = frame_counter.try_next();
        Self::with_meta_data_and_ctr(counter, payload, meta_data)
    }

    /// Creates a new media frame by copying the data of a payload buffer and assigning it the next counter (CTR) value.
    /// Returns an [`SframeError::CounterCreationFailed`] wrapping the error of the
    /// [`FrameCounter`] if no counter value could be created.
    pub fn try_new<C, P>(frame_counter: &mut C, payload: P) -> Result<Self>
    where
        C: FrameCounter,
        P: AsRef<[u8]>,
    {
        Self::try_with_meta_data(frame_counter, payload, [])
    }

    /// Creates a new media frame and assigns it the next counter (CTR) value.
    /// Payload and meta data are copied into an internal buffer.
    /// Returns an [`SframeError::CounterCreationFailed`] wrapping the error of the
    /// [`FrameCounter`] if no counter value could be created.
    pub fn try_with_meta_data<C, P, M>(
        frame_counter: &mut C,
        payload: P,
        meta_data: M,
    ) -> Result<Self>
    where
        C: FrameCounter,
        P: AsRef<[u8]>,
        M: AsRef<[u8]>,
    {
        let counter = next_counter(frame_counter)?;
        Ok(Self::with_meta_data_and_ctr(counter, payload, meta_data))
    }

    pub(super) fn with_meta_data_and_ctr<P, M>(counter: Counter, payload: P, meta_data: M) -> Self
    where
        P: AsRef<[u8]>,
        M: AsRef<[u8]>,
    {
        let payload = payload.as_ref();
        let meta_data = meta_data.as_ref();
        let meta_len = meta_data.len();
        let payload_len = payload.len();
        log::trace!(
            "Creating MediaFrame # {counter} with payload of size {payload_len} using meta data of size {meta_len}"
        );

        let mut buffer = Vec::with_capacity(payload_len + meta_len);
        buffer.extend(meta_data);
        buffer.extend(payload);

        Self {
            counter,
            buffer,
            meta_len,
        }
    }

    pub(super) fn with_buffer(counter: Counter, buffer: Vec<u8>, meta_len: usize) -> Self {
        Self {
            counter,
            buffer,
            meta_len,
        }
    }

    /// Meta data associated with this media frame
    pub fn meta_data(&self) -> &[u8] {
        &self.buffer[..self.meta_len]
    }

    /// Payload of this media frame
    pub fn payload(&self) -> &[u8] {
        &self.buffer[self.meta_len..]
    }

    /// Frame count for the Sframe scheme associated to this media frame
    pub fn counter(&self) -> Counter {
        self.counter
    }

    /// Encrypts the media frame with the sframe key according to [RFC 9605 4.4.3](https://www.rfc-editor.org/rfc/rfc9605.html#name-encryption).
    /// Dynamically allocates memory for the resulting [`EncryptedFrame`].
    /// The associated meta data is not encrypted but considered for the authentication tag.
    /// Returns an [`crate::error::SframeError`] when encryption fails
    pub fn encrypt<A, D>(&self, key: &GenericEncryptionKey<A, D>) -> Result<EncryptedFrame>
    where
        A: AeadEncrypt<Secret = D::Secret>,
        D: KeyDerivation,
    {
        let view =
            MediaFrameView::with_meta_data_and_ctr(self.counter, self.payload(), self.meta_data());
        view.encrypt(key)
    }

    /// Encrypts the media frame with the sframe key according to [RFC 9605 4.4.3](https://www.rfc-editor.org/rfc/rfc9605.html#name-encryption)
    /// and stores the result into the provided buffer. An [`EncryptedFrameView`] on the buffer is returned on success.
    /// The associated meta data is not encrypted but considered for the authentication tag.
    /// Returns an [`crate::error::SframeError`] when encryption fails.
    pub fn encrypt_into<'obuf, A, D>(
        &self,
        key: &GenericEncryptionKey<A, D>,
        buffer: &'obuf mut impl FrameBuffer,
    ) -> Result<EncryptedFrameView<'obuf>>
    where
        A: AeadEncrypt<Secret = D::Secret>,
        D: KeyDerivation,
    {
        let view =
            MediaFrameView::with_meta_data_and_ctr(self.counter, self.payload(), self.meta_data());
        view.encrypt_into(key, buffer)
    }
}

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

#[cfg(all(test, crypto_backend))]
mod test {
    use crate::{
        CipherSuite,
        frame::{
            FrameCounter,
            media_frame::{MediaFrame, MediaFrameView},
        },
        header::Counter,
        key::EncryptionKey,
        util::test::assert_bytes_eq,
    };
    use pretty_assertions::assert_eq;
    use std::convert::Infallible;

    const COUNTER: u64 = 42;
    const PAYLOAD: &[u8] = &[6, 6, 6, 6, 6, 6];
    const KEY_ID: u64 = 666;
    const META_DATA: &[u8] = b"META";

    /// A counter which cannot fail, so that the infallible API can be used
    struct InfallibleCounter;
    impl FrameCounter for InfallibleCounter {
        type Error = Infallible;
        fn try_next(&mut self) -> Result<Counter, Self::Error> {
            Ok(COUNTER)
        }
    }

    #[derive(Debug, thiserror::Error)]
    #[error("no counter left")]
    struct NoCounter;

    /// A counter which always fails, to observe how its error is propagated
    struct FailingCounter;
    impl FrameCounter for FailingCounter {
        type Error = NoCounter;
        fn try_next(&mut self) -> Result<Counter, Self::Error> {
            Err(NoCounter)
        }
    }

    #[test]
    fn create_media_frame_view_with_an_infallible_counter() {
        let frame_view = MediaFrameView::new(&mut InfallibleCounter, &PAYLOAD);

        assert_eq!(frame_view.counter(), COUNTER);
    }

    #[test]
    fn create_media_frame_with_an_infallible_counter() {
        let frame = MediaFrame::new(&mut InfallibleCounter, PAYLOAD);

        assert_eq!(frame.counter(), COUNTER);
    }

    #[test]
    fn propagate_the_counter_error_creating_a_media_frame_view() {
        let error = MediaFrameView::try_new(&mut FailingCounter, &PAYLOAD).unwrap_err();

        assert!(error.source_as::<NoCounter>().is_some());
    }

    #[test]
    fn propagate_the_counter_error_creating_a_media_frame() {
        let error = MediaFrame::try_new(&mut FailingCounter, PAYLOAD).unwrap_err();

        assert!(error.source_as::<NoCounter>().is_some());
    }

    #[test]
    fn create_media_frame_with_meta_data() {
        let frame_view = MediaFrameView::with_meta_data_and_ctr(COUNTER, &PAYLOAD, &META_DATA);

        assert_eq!(frame_view.payload(), PAYLOAD);
        assert_eq!(frame_view.counter(), COUNTER);
        assert_eq!(frame_view.meta_data(), META_DATA);
    }

    #[test]
    fn encrypt_media_frame_view() {
        let key =
            EncryptionKey::derive_from(CipherSuite::AesGcm256Sha512, KEY_ID, "SECRET").unwrap();
        let mut encrypt_buffer = Vec::new();

        let media_frame = MediaFrameView::with_meta_data_and_ctr(COUNTER, &PAYLOAD, META_DATA);
        let encrypted_frame = media_frame.encrypt_into(&key, &mut encrypt_buffer).unwrap();

        assert_eq!(encrypted_frame.header().key_id(), KEY_ID);
        assert_eq!(encrypted_frame.header().counter(), COUNTER);
        assert_bytes_eq(encrypted_frame.meta_data(), META_DATA);
        assert!(!encrypted_frame.cipher_text().is_empty());
    }

    #[test]
    fn encrypt_media_frame() {
        let key =
            EncryptionKey::derive_from(CipherSuite::AesGcm256Sha512, KEY_ID, "SECRET").unwrap();

        let media_frame = MediaFrame::with_meta_data_and_ctr(COUNTER, PAYLOAD, META_DATA);
        let encrypted_frame = media_frame.encrypt(&key).unwrap();

        assert_eq!(encrypted_frame.header().key_id(), KEY_ID);
        assert_eq!(encrypted_frame.header().counter(), COUNTER);
        assert_bytes_eq(encrypted_frame.meta_data(), META_DATA);
        assert!(!encrypted_frame.cipher_text().is_empty());
    }
}