philbin 1.0.1

use constant_time_eq::constant_time_eq;
use std::iter::zip;
use zerocopy::IntoBytes;
use zeroize::Zeroize;

use crate::{
  aligned_buffer::AlignedBuf,
  arch::Simd,
  easy::{
    AssociatedData, AuthTag, Ciphertext, CiphertextMut, Error, Plaintext,
    PlaintextMut, Result,
  },
};

/// The core AEGIS algo logic that is common across all variants of both
/// AEGIS-256 and AEGIS-128.
pub trait Aegis<const OUTPUT_RATE_BYTES: usize>
where
  Self: Sized,
{
  type AlignedBuf: AlignedBuf<OUTPUT_RATE_BYTES>;
  type Simd: Simd;

  // AEGIS-128 variants use Key128 / Nonce128.
  // AEGIS-256 variants use Key256 / Nonce256.
  type Key;
  type Nonce;

  // (See doc comments in lib.rs)
  #[inline(always)]
  fn encrypt_to_slice_detached<Tag: AuthTag>(
    simd: Self::Simd,
    plaintext: Plaintext,
    associated_data: AssociatedData,
    key: &Self::Key,
    nonce: Self::Nonce,
    mut ciphertext: CiphertextMut,
  ) -> Result<Tag> {
    // The output buffer the user provided as `ciphertext` can be bigger than
    // necessary, so we slice it down to the exact size we need because it
    // makes the code simpler.
    //
    // This also checks that `ciphertext.len() >= plaintext.len().
    //
    // SECURITY: The (indirect) IF check is safe WRT timing attacks because:
    // - The _length_ of the plaintext is considered public because the length
    //   of an AEGIS ciphertext equals that of the plaintext.
    let ciphertext_mut_slice = ciphertext
      .as_bytes_mut()
      .get_mut(..plaintext.as_bytes().len())
      .ok_or(Error::OutputBufferTooSmall)?;

    // PERF: `with_target_features` lets the compiler choose if `init` should be
    // inlined. Speeds up Fallback backend 10x.
    let mut state =
      simd.with_target_features(|simd| Self::init(simd, key, nonce));
    state.absorb_associated_data(associated_data.as_bytes());
    state.encrypt_to_ciphertext(plaintext.as_bytes(), ciphertext_mut_slice);

    Ok(state.finalize(plaintext.as_bytes(), associated_data.as_bytes()))
  }

  // (See doc comments in lib.rs)
  #[inline(always)]
  fn encrypt_in_place_detached<Tag: AuthTag>(
    simd: Self::Simd,
    mut buffer: PlaintextMut,
    associated_data: AssociatedData,
    key: &Self::Key,
    nonce: Self::Nonce,
  ) -> Result<Tag> {
    // PERF: `with_target_features` lets the compiler choose if `init` should be
    // inlined. Speeds up Fallback backend 10x.
    let mut state =
      simd.with_target_features(|simd| Self::init(simd, key, nonce));
    state.absorb_associated_data(associated_data.as_bytes());
    state.encrypt_buffer(buffer.as_bytes_mut());
    Ok(state.finalize(buffer.as_bytes(), associated_data.as_bytes()))
  }

  // (See doc comments in lib.rs)
  #[inline(always)]
  fn decrypt_to_slice_detached<Tag: AuthTag>(
    simd: Self::Simd,
    ciphertext: Ciphertext,
    auth_tag: &Tag,
    associated_data: AssociatedData,
    key: &Self::Key,
    nonce: Self::Nonce,
    mut plaintext: PlaintextMut,
  ) -> Result<()> {
    // The output buffer the user provided as `plaintext` can be bigger than
    // necessary, so we slice it down to the exact size we need because it
    // makes the code simpler.
    //
    // This also checks that `plaintext.len() >= ciphertext.len().
    //
    // SECURITY: The (indirect) IF check is safe WRT timing attacks because:
    // - The _length_ of the plaintext is considered public because the length
    //   of an AEGIS ciphertext equals that of the plaintext.
    let plaintext_mut_slice = plaintext
      .as_bytes_mut()
      .get_mut(..ciphertext.as_bytes().len())
      .ok_or(Error::OutputBufferTooSmall)?;

    // PERF: `with_target_features` lets the compiler choose if `init` should be
    // inlined. Speeds up Fallback backend 10x.
    let mut state =
      simd.with_target_features(|simd| Self::init(simd, key, nonce));
    state.absorb_associated_data(associated_data.as_bytes());
    state.decrypt_to_plaintext(ciphertext.as_bytes(), plaintext_mut_slice);

    let mut expected_tag = state.finalize::<Tag>(
      plaintext_mut_slice.as_bytes(),
      associated_data.as_bytes(),
    );

    // SECURITY: The IF check is safe WRT timing attacks because:
    // - We use a constant-time comparison to validate `auth_tag`.
    if !constant_time_eq(auth_tag.as_ref(), expected_tag.as_ref()) {
      plaintext_mut_slice.zeroize();
      expected_tag.as_mut().zeroize();
      Err(Error::AuthTagInvalid)
    } else {
      Ok(())
    }
  }

  // (See doc comments in lib.rs)
  #[inline(always)]
  fn decrypt_in_place_detached<Tag: AuthTag>(
    simd: Self::Simd,
    mut buffer: CiphertextMut,
    auth_tag: &Tag,
    associated_data: AssociatedData,
    key: &Self::Key,
    nonce: Self::Nonce,
  ) -> Result<()> {
    // PERF: `with_target_features` lets the compiler choose if `init` should be
    // inlined. Speeds up Fallback backend 10x.
    let mut state =
      simd.with_target_features(|simd| Self::init(simd, key, nonce));
    state.absorb_associated_data(associated_data.as_bytes());
    state.decrypt_buffer(buffer.as_bytes_mut());

    let mut expected_tag =
      state.finalize::<Tag>(buffer.as_bytes(), associated_data.as_bytes());

    // SECURITY: The IF check is safe WRT timing attacks because:
    // - We use a constant-time comparison to validate `auth_tag`.
    if !constant_time_eq(auth_tag.as_ref(), expected_tag.as_ref()) {
      buffer.as_bytes_mut().zeroize();
      expected_tag.as_mut().zeroize();
      Err(Error::AuthTagInvalid)
    } else {
      Ok(())
    }
  }

  /// Constructs the initial AEGIS state using the given `key` and `nonce`.
  ///
  /// The function is called `init` instead of `new` to more closely follow
  /// RFC 10032.
  fn init(simd: Self::Simd, key: &Self::Key, nonce: Self::Nonce) -> Self;

  /// Absorbs a chunk of associated data into the state.
  ///
  /// Takes `chunk` by-ref so we can pass a correctly aligned reference.
  fn absorb(&mut self, chunk: &[u8; OUTPUT_RATE_BYTES]);

  /// Absorbs the whole `associated_data` into the state.
  #[inline(always)]
  fn absorb_associated_data(&mut self, associated_data: &[u8]) {
    let (chunks, remainder) = associated_data.as_chunks();
    for chunk in chunks {
      self.absorb(chunk);
    }

    // The relevant code from the spec:
    //   ad_blocks = Split(ZeroPad(ad, R), R)
    //
    // The spec's definition of ZeroPad:
    //
    // ZeroPad(x, n): returns x after appending zeros until its length is a
    // multiple of n bits. No padding is added if the length of x is already a
    // multiple of n, including when x is empty.
    //
    // Thus we MUST NOT absorb an all-zero chunk when `remainder` is
    // empty.
    //
    // SECURITY: This IF check is safe WRT timing attacks because:
    // - The associated data is NOT secret and is assumed to be public.
    // - This function can never be constant-time because associated data is
    //   variable-length.
    if !remainder.is_empty() {
      let buf = Self::zero_padded(remainder);
      self.absorb(buf.as_ref());
    }
  }

  /// The Enc function encrypts a `[u8; OUTPUT_RATE_BYTES]` input block using
  /// the state.
  fn enc(
    &mut self,
    plaintext_chunk: &[u8; OUTPUT_RATE_BYTES],
    ciphertext_chunk: &mut [u8; OUTPUT_RATE_BYTES],
  );

  /// Encrypts the `plaintext` into the `ciphertext` buffer.
  ///
  /// # Panics
  ///
  /// Panics if `plaintext.len() != ciphertext.len()`.
  #[inline(always)]
  fn encrypt_to_ciphertext(&mut self, plaintext: &[u8], ciphertext: &mut [u8]) {
    // We assert not just for correctness, but also to help the optimizer.
    // Calling code verifies lengths so this won't actually ever panic (modulo
    // bugs).
    //
    // SECURITY: This IF check is safe WRT timing attacks because:
    // - Calling code ensures these lengths are equal, so this check will
    //   always be true.
    // - Lengths of ciphertext and plaintext are not secret.
    assert!(
      ciphertext.len() == plaintext.len(),
      "ciphertext len: {}, plaintext len: {}",
      ciphertext.len(),
      plaintext.len()
    );

    let (pt_chunks, pt_remainder) = plaintext.as_chunks::<OUTPUT_RATE_BYTES>();
    let (ct_chunks, ct_remainder) =
      ciphertext.as_chunks_mut::<OUTPUT_RATE_BYTES>();

    for (pt_chunk, ct_chunk) in zip(pt_chunks, ct_chunks) {
      self.enc(pt_chunk, ct_chunk);
    }

    // We MUST NOT encrypt an all-zero chunk when `remainder` is
    // empty. See note on the same IF check in `absorb_associated_data`.
    //
    // SECURITY: This IF check is safe WRT timing attacks because:
    // - This function can never be constant-time because plaintext is
    //   variable-length.
    // - The _length_ of the plaintext is considered public because the length
    //   of an AEGIS ciphertext equals that of the plaintext.
    if !pt_remainder.is_empty() {
      let pt_chunk = Self::zero_padded(pt_remainder);
      let mut ct_chunk = Self::AlignedBuf::new();
      self.enc(pt_chunk.as_ref(), ct_chunk.as_mut());
      ct_remainder.copy_from_slice(&ct_chunk.as_ref()[..pt_remainder.len()]);
    }
  }

  /// Encrypts the `buffer` in place.
  #[inline(always)]
  fn encrypt_buffer(&mut self, buffer: &mut [u8]) {
    let (chunks, remainder) = buffer.as_chunks_mut::<OUTPUT_RATE_BYTES>();

    for chunk in chunks {
      let pt_chunk = Self::AlignedBuf::from_array(*chunk);
      self.enc(pt_chunk.as_ref(), chunk);
    }

    // We MUST NOT encrypt an all-zero chunk when `remainder` is
    // empty. See note on the same IF check in `absorb_associated_data`.
    //
    // SECURITY: This IF check is safe WRT timing attacks because:
    // - This function can never be constant-time because plaintext is
    //   variable-length.
    // - The _length_ of the plaintext is considered public because the length
    //   of an AEGIS ciphertext equals that of the plaintext.
    // - The alignment of plaintext/ciphertext references is not secret
    //   information.
    if !remainder.is_empty() {
      let pt_chunk = Self::zero_padded(remainder);
      let mut ct_chunk = Self::AlignedBuf::new();
      self.enc(pt_chunk.as_ref(), ct_chunk.as_mut());
      remainder.copy_from_slice(&ct_chunk.as_ref()[..remainder.len()]);
    }
  }

  /// Decrypts a `[u8; OUTPUT_RATE_BYTES]` input block using the state.
  ///
  /// Takes `ciphertext_chunk` by-ref so we can pass a correctly aligned
  /// reference.
  fn dec(
    &mut self,
    ciphertext_chunk: &[u8; OUTPUT_RATE_BYTES],
    plaintext_chunk: &mut [u8; OUTPUT_RATE_BYTES],
  );

  /// Decrypts the remainder of the ciphertext that does not fill a
  /// `[u8; OUTPUT_RATE_BYTES]` bit chunk.
  ///
  /// The remainder should be padded with zero and provided as
  /// `ciphertext_chunk`.
  ///
  /// Takes `ciphertext_chunk` by-ref so we can pass a correctly aligned
  /// reference.
  fn dec_partial(
    &mut self,
    ciphertext_chunk: &[u8; OUTPUT_RATE_BYTES],
    plaintext_remainder: &mut [u8],
  );

  /// Decrypts the `ciphertext` into the `plaintext` buffer.
  ///
  /// # Panics
  ///
  /// Panics if `ciphertext.len() != plaintext.len()`.
  #[inline(always)]
  fn decrypt_to_plaintext(&mut self, ciphertext: &[u8], plaintext: &mut [u8]) {
    // We assert not just for correctness, but also to help the optimizer.
    // Calling code verifies lengths so this won't actually ever panic (modulo
    // bugs).
    //
    // SECURITY: This IF check is safe WRT timing attacks because:
    // - Calling code ensures these lengths are equal, so this check will
    //   always be true.
    // - Lengths of ciphertext and plaintext are not secret.
    assert!(
      ciphertext.len() == plaintext.len(),
      "ciphertext len: {}, plaintext len: {}",
      ciphertext.len(),
      plaintext.len()
    );

    let (pt_chunks, pt_remainder) = plaintext.as_chunks_mut();
    let (ct_chunks, ct_remainder) = ciphertext.as_chunks();

    for (pt_chunk, ct_chunk) in zip(pt_chunks, ct_chunks) {
      self.dec(ct_chunk, pt_chunk);
    }

    // SECURITY: This IF check is safe WRT timing attacks because:
    // - The ciphertext is not secret.
    if !ct_remainder.is_empty() {
      let ct_chunk = Self::zero_padded(ct_remainder);
      self.dec_partial(ct_chunk.as_ref(), pt_remainder);
    }
  }

  /// Decrypts the `buffer` in place.
  #[inline(always)]
  fn decrypt_buffer(&mut self, buffer: &mut [u8]) {
    let (chunks, remainder) = buffer.as_chunks_mut();

    for chunk in chunks {
      let ct_chunk = Self::AlignedBuf::from_array(*chunk);
      self.dec(ct_chunk.as_ref(), chunk);
    }

    // SECURITY: This IF check is safe WRT timing attacks because:
    // - The ciphertext is not secret.
    if !remainder.is_empty() {
      let ct_chunk = Self::zero_padded(remainder);
      self.dec_partial(ct_chunk.as_ref(), remainder);
    }
  }

  /// Produces the authorization tag.
  fn finalize<Tag: AuthTag>(
    &mut self,
    plaintext: &[u8],
    associated_data: &[u8],
  ) -> Tag;

  /// Takes a byte slice shorter than `BYTES`.
  /// Returns a version of it padded with zero (on the "right") so that it
  /// fills a `Self::AlignedBuf`.
  #[inline(always)]
  fn zero_padded(undersized: &[u8]) -> Self::AlignedBuf {
    // SECURITY: This IF check is safe WRT timing attacks because:
    // - The length of the provided associated data, plaintext and ciphertext
    //   are not secret.
    //
    // NOTE: If this ever panics, there's a bug in this crate.
    assert!(!undersized.is_empty() && undersized.len() < OUTPUT_RATE_BYTES);

    let mut padded = Self::AlignedBuf::from_array([0; OUTPUT_RATE_BYTES]);
    padded.as_mut()[..undersized.len()].copy_from_slice(undersized);
    padded
  }
}