philbin 1.0.1

use crate::{
  aligned_buffer::{AlignedBufHolder, AlignedBufRouter},
  arch::Simd,
  base::{aegis::Aegis, block::BlockDoubleRate},
  careful::Nonce128,
  easy::{AuthTag, AuthTag128, AuthTag256, Key128},
  utils::num_bits,
};
use zerocopy::FromZeros;

/// The state for the AEGIS-128L algorithm.
///
/// `StateBlock` represents the algo block storage and low-level block
/// manipulation functions, making the `Aegis128L` type generic across
/// architectures, SIMD sizes, CPU features and even AEGIS-128 variants.
///
/// AEGIS-128X2 and AEGIS-128X4 also use this type; only the `StateBlock` is
/// different!
///
/// See [`Block`] and [`BlockDoubleRate`].
pub struct Aegis128L<S: Simd, StateBlock, const OUTPUT_RATE_BYTES: usize> {
  s0: StateBlock,
  s1: StateBlock,
  s2: StateBlock,
  s3: StateBlock,
  s4: StateBlock,
  s5: StateBlock,
  s6: StateBlock,
  s7: StateBlock,
  simd: S,
}

// NOTE: There is no need for an assert (static or otherwise) checking that
// size_of::<StateBlock>() * 2 == OUTPUT_RATE_BYTES because we only impl
// BlockDoubleRate for *exact* values of OUTPUT_RATE_BYTES. In other words,
// we already get a compile-time error for all values of OUTPUT_RATE_BYTES
// other than the supported ones.
impl<S, StateBlock, const OUTPUT_RATE_BYTES: usize>
  Aegis128L<S, StateBlock, OUTPUT_RATE_BYTES>
where
  S: Simd,
  StateBlock: BlockDoubleRate<OUTPUT_RATE_BYTES>,
{
  /// Updates the internal state with the provided blocks.
  #[inline(always)]
  fn update(&mut self, m0: StateBlock, m1: StateBlock) {
    let orig_s7 = self.s7;
    self.s7 = StateBlock::aes_encrypt_round(self.s6, self.s7);
    self.s6 = StateBlock::aes_encrypt_round(self.s5, self.s6);
    self.s5 = StateBlock::aes_encrypt_round(self.s4, self.s5);
    self.s4 = StateBlock::aes_encrypt_round(self.s3, self.s4);
    self.s3 = StateBlock::aes_encrypt_round(self.s2, self.s3);
    self.s2 = StateBlock::aes_encrypt_round(self.s1, self.s2);
    self.s1 = StateBlock::aes_encrypt_round(self.s0, self.s1);
    self.s0 = StateBlock::aes_encrypt_round(orig_s7, self.s0);

    // PERF: Pulling out the XOR for s4 and s0 and doing them last *greatly*
    // improves perf for the 128x4 variant (2x!), while hurting 128x2 by a
    // ~couple of percent? Seems like a worthwhile trade.
    self.s0 = self.s0 ^ m0;
    self.s4 = self.s4 ^ m1;
  }
}

impl<S, StateBlock, const OUTPUT_RATE_BYTES: usize> Aegis<OUTPUT_RATE_BYTES>
  for Aegis128L<S, StateBlock, OUTPUT_RATE_BYTES>
where
  S: Simd,
  StateBlock: BlockDoubleRate<OUTPUT_RATE_BYTES, Simd = S>,
  AlignedBufHolder: AlignedBufRouter<OUTPUT_RATE_BYTES>,
{
  type AlignedBuf =
    <AlignedBufHolder as AlignedBufRouter<OUTPUT_RATE_BYTES>>::AlignedBuf;
  type Simd = S;
  type Key = Key128;
  type Nonce = Nonce128;

  #[inline(always)]
  fn init(simd: Self::Simd, key_raw: &Key128, nonce_raw: Nonce128) -> Self {
    let key = StateBlock::from_128_bits(simd, *key_raw.expose_secret());
    let nonce = StateBlock::from_128_bits(simd, nonce_raw.into_array());

    let mut state = Self {
      s0: key ^ nonce,
      s1: StateBlock::c1(simd),
      s2: StateBlock::c0(simd),
      s3: StateBlock::c1(simd),
      s4: key ^ nonce,
      s5: key ^ StateBlock::c0(simd),
      s6: key ^ StateBlock::c1(simd),
      s7: key ^ StateBlock::c0(simd),
      simd,
    };

    for _ in 0..10 {
      state.s3 = state.s3 ^ StateBlock::ctx(simd);
      state.s7 = state.s7 ^ StateBlock::ctx(simd);
      state.update(nonce, key);
    }
    state
  }

  #[inline(always)]
  fn absorb(&mut self, chunk: &[u8; OUTPUT_RATE_BYTES]) {
    let [t0, t1] = StateBlock::split(self.simd, chunk);
    self.update(t0, t1);
  }

  #[inline(always)]
  fn enc(
    &mut self,
    plaintext_chunk: &[u8; OUTPUT_RATE_BYTES],
    ciphertext_chunk: &mut [u8; OUTPUT_RATE_BYTES],
  ) {
    let [t0, t1] = StateBlock::split(self.simd, plaintext_chunk);
    // PERF: Interleaving operations for z0 and z1 is important.
    let mut z0 = t0 ^ self.s1;
    let mut z1 = t1 ^ self.s2;
    z0 = z0 ^ self.s6;
    z1 = z1 ^ self.s5;
    let tmp0 = self.s2 & self.s3;
    let tmp1 = self.s6 & self.s7;
    z0 = z0 ^ tmp0;
    z1 = z1 ^ tmp1;

    *ciphertext_chunk = StateBlock::concat(z0, z1);
    self.update(t0, t1); // PERF: Making Update() last is important.
  }

  #[inline(always)]
  fn dec(
    &mut self,
    ciphertext_chunk: &[u8; OUTPUT_RATE_BYTES],
    plaintext_chunk: &mut [u8; OUTPUT_RATE_BYTES],
  ) {
    let [t0, t1] = StateBlock::split(self.simd, ciphertext_chunk);

    // PERF: Interleaving operations for z0 and z1 is important.
    let mut z0 = t0 ^ self.s1;
    let mut z1 = t1 ^ self.s2;
    z0 = z0 ^ self.s6;
    z1 = z1 ^ self.s5;
    let tmp0 = self.s2 & self.s3;
    let tmp1 = self.s6 & self.s7;
    z0 = z0 ^ tmp0;
    z1 = z1 ^ tmp1;

    *plaintext_chunk = StateBlock::concat(z0, z1);
    self.update(z0, z1); // PERF: Making Update() last is important.
  }

  #[inline(always)]
  fn dec_partial(
    &mut self,
    ciphertext_chunk: &[u8; OUTPUT_RATE_BYTES],
    plaintext_remainder: &mut [u8],
  ) {
    let [t0, t1] = StateBlock::split(self.simd, ciphertext_chunk);

    // PERF: Interleaving operations for z0 and z1 is important.
    let mut z0 = t0 ^ self.s1;
    let mut z1 = t1 ^ self.s2;
    z0 = z0 ^ self.s6;
    z1 = z1 ^ self.s5;
    let tmp0 = self.s2 & self.s3;
    let tmp1 = self.s6 & self.s7;
    z0 = z0 ^ tmp0;
    z1 = z1 ^ tmp1;

    let mut out = StateBlock::concat(z0, z1);
    // SECURITY:
    // - The length of the ciphertext and plaintext are not secret.
    // - We don't need the fancy `zeroize()` here since the bytes we're throwing
    //   away are useless. `zero()` (from `zerocopy` crate) is good enough.
    out[plaintext_remainder.len()..].zero(); // ZeroPad(xn) in spec
    plaintext_remainder.copy_from_slice(&out[..plaintext_remainder.len()]);

    let [v0, v1] = StateBlock::split(self.simd, &out);
    self.update(v0, v1);
  }

  #[inline(always)]
  fn finalize<Tag: AuthTag>(
    &mut self,
    plaintext: &[u8],
    associated_data: &[u8],
  ) -> Tag {
    let lengths = StateBlock::from_ints(
      self.simd,
      num_bits(associated_data),
      num_bits(plaintext),
    );
    let t = self.s2 ^ lengths;
    for _ in 0..7 {
      self.update(t, t);
    }

    let mut tag = Tag::default();
    // NOTE: Due to:
    // - size_of::<Tag>() being a compile-time constant after
    //   monomorphization,
    // - the compiler using dead code elimination,
    // there will be no branching in the generated assembly.
    //
    // SECURITY: We do not depend on the absence of branching since the auth
    // tag size is not secret.
    //
    // While Rust likes to pretend that `size_of::<Tag>()` isn't known at
    // compile time (for complex technical reasons), the optimizer doesn't
    // care and eats those dead branches for breakfast.
    match size_of::<Tag>() {
      AuthTag128::BYTES => {
        // PERF: Adding parens or pulling out subexpressions doesn't even change
        // the generated assembly, let alone improve perf.
        let block =
          self.s0 ^ self.s1 ^ self.s2 ^ self.s3 ^ self.s4 ^ self.s5 ^ self.s6;
        // Can't panic as `tag` is guaranteed to be [u8; 16] here.
        tag.as_mut().copy_from_slice(&block.xor_down());
      },
      AuthTag256::BYTES => {
        let high = (self.s0 ^ self.s1 ^ self.s2 ^ self.s3).xor_down();
        let low = (self.s4 ^ self.s5 ^ self.s6 ^ self.s7).xor_down();
        // Can't panic as `tag` is guaranteed to be [u8; 32] here.
        tag.as_mut().copy_from_slice([high, low].as_flattened());
      },
      _ => unreachable!("Tag limited to 16 or 32 bytes at compile-time."),
    }
    tag
  }
}