#![deny(unsafe_code)]
#![warn(missing_docs, rust_2018_idioms)]

pub use aead::{self, AeadCore, AeadInPlace, Error, Key, KeyInit, KeySizeUser};
pub use aes;

use cipher::{
  BlockCipher, BlockEncrypt, BlockSizeUser, InnerIvInit, StreamCipherCore,
  consts::{U0, U16},
  generic_array::{ArrayLength, GenericArray},
};
use core::marker::PhantomData;
use ghash::{GHash, universal_hash::UniversalHash};

use aes::{Aes128, Aes256, cipher::consts::U12};

/// Maximum length of associated data.
pub const A_MAX: u64 = 1 << 36;

/// Maximum length of plaintext.
pub const P_MAX: u64 = 1 << 36;

/// Maximum length of ciphertext.
pub const C_MAX: u64 = (1 << 36) + 16;

/// AES-GCM nonces.
pub type Nonce<NonceSize> = GenericArray<u8, NonceSize>;

/// AES-GCM tags.
pub type Tag<TagSize = U16> = GenericArray<u8, TagSize>;

/// Trait implemented for valid tag sizes, i.e.
/// [`U12`][consts::U12], [`U13`][consts::U13], [`U14`][consts::U14],
/// [`U15`][consts::U15] and [`U16`][consts::U16].
pub trait TagSize: private::SealedTagSize {}

impl<T: private::SealedTagSize> TagSize for T {}

mod private {
  use aead::generic_array::ArrayLength;
  use cipher::{Unsigned, consts};

  // Sealed traits stop other crates from implementing any traits that use it.
  pub trait SealedTagSize: ArrayLength<u8> + Unsigned {}

  impl SealedTagSize for consts::U8 {}
  impl SealedTagSize for consts::U12 {}
  impl SealedTagSize for consts::U13 {}
  impl SealedTagSize for consts::U14 {}
  impl SealedTagSize for consts::U15 {}
  impl SealedTagSize for consts::U16 {}
}

/// AES-GCM with a 128-bit key and 96-bit nonce.
pub type Aes128Gcm = AesGcm<Aes128, U12>;

/// AES-GCM with a 256-bit key and 96-bit nonce.
pub type Aes256Gcm = AesGcm<Aes256, U12>;

/// AES block.
type Block = GenericArray<u8, U16>;

/// Counter mode with a 32-bit big endian counter.
type Ctr32BE<Aes> = ctr::CtrCore<Aes, ctr::flavors::Ctr32BE>;

/// AES-GCM: generic over an underlying AES implementation and nonce size.
///
/// This type is generic to support substituting alternative AES implementations
/// (e.g. embedded hardware implementations)
///
/// It is NOT intended to be instantiated with any block cipher besides AES!
/// Doing so runs the risk of unintended cryptographic properties!
///
/// The `NonceSize` generic parameter can be used to instantiate AES-GCM with other
/// nonce sizes, however it's recommended to use it with `typenum::U12`,
/// the default of 96-bits.
///
/// The `TagSize` generic parameter can be used to instantiate AES-GCM with other
/// authorization tag sizes, however it's recommended to use it with `typenum::U16`,
/// the default of 128-bits.
///
/// If in doubt, use the built-in [`Aes128Gcm`] and [`Aes256Gcm`] type aliases.
#[derive(Clone)]
pub struct AesGcm<Aes, NonceSize, TagSize = U16>
where
  TagSize: self::TagSize,
{
  /// Encryption cipher.
  cipher: Aes,

  /// GHASH authenticator.
  ghash: GHash,

  /// Length of the nonce.
  nonce_size: PhantomData<NonceSize>,

  /// Length of the tag.
  tag_size: PhantomData<TagSize>,
}

impl<Aes, NonceSize, TagSize> KeySizeUser for AesGcm<Aes, NonceSize, TagSize>
where
  Aes: KeySizeUser,
  TagSize: self::TagSize,
{
  type KeySize = Aes::KeySize;
}

impl<Aes, NonceSize, TagSize> KeyInit for AesGcm<Aes, NonceSize, TagSize>
where
  Aes: BlockSizeUser<BlockSize = U16> + BlockEncrypt + KeyInit,
  TagSize: self::TagSize,
{
  fn new(key: &Key<Self>) -> Self {
    Aes::new(key).into()
  }
}

impl<Aes, NonceSize, TagSize> From<Aes> for AesGcm<Aes, NonceSize, TagSize>
where
  Aes: BlockSizeUser<BlockSize = U16> + BlockEncrypt,
  TagSize: self::TagSize,
{
  fn from(cipher: Aes) -> Self {
    let mut ghash_key = ghash::Key::default();
    cipher.encrypt_block(&mut ghash_key);

    let ghash = GHash::new(&ghash_key);

    Self {
      cipher,
      ghash,
      nonce_size: PhantomData,
      tag_size: PhantomData,
    }
  }
}

impl<Aes, NonceSize, TagSize> AeadCore for AesGcm<Aes, NonceSize, TagSize>
where
  NonceSize: ArrayLength<u8>,
  TagSize: self::TagSize,
{
  type NonceSize = NonceSize;
  type TagSize = TagSize;
  type CiphertextOverhead = U0;
}

impl<Aes, NonceSize, TagSize> AeadInPlace for AesGcm<Aes, NonceSize, TagSize>
where
  Aes: BlockCipher + BlockSizeUser<BlockSize = U16> + BlockEncrypt,
  NonceSize: ArrayLength<u8>,
  TagSize: self::TagSize,
{
  fn encrypt_in_place_detached(
    &self,
    nonce: &Nonce<NonceSize>,
    associated_data: &[u8],
    buffer: &mut [u8],
  ) -> Result<Tag<TagSize>, Error> {
    if buffer.len() as u64 > P_MAX || associated_data.len() as u64 > A_MAX {
      return Err(Error);
    }

    let (ctr, mask) = self.init_ctr(nonce);

    ctr.apply_keystream_partial(buffer.into());

    let full_tag = self.compute_tag(mask, associated_data, buffer);
    Ok(Tag::clone_from_slice(&full_tag[..TagSize::to_usize()]))
  }

  fn decrypt_in_place_detached(
    &self,
    nonce: &Nonce<NonceSize>,
    associated_data: &[u8],
    buffer: &mut [u8],
    tag: &Tag<TagSize>,
  ) -> Result<(), Error> {
    if buffer.len() as u64 > C_MAX || associated_data.len() as u64 > A_MAX {
      return Err(Error);
    }

    let (ctr, mask) = self.init_ctr(nonce);

    let expected_tag = self.compute_tag(mask, associated_data, buffer);

    use subtle::ConstantTimeEq;
    if expected_tag[..TagSize::to_usize()].ct_eq(tag).into() {
      ctr.apply_keystream_partial(buffer.into());
      Ok(())
    } else {
      Err(Error)
    }
  }
}

impl<Aes, NonceSize, TagSize> AesGcm<Aes, NonceSize, TagSize>
where
  Aes: BlockCipher + BlockSizeUser<BlockSize = U16> + BlockEncrypt,
  NonceSize: ArrayLength<u8>,
  TagSize: self::TagSize,
{
  /// Initialize counter mode.
  ///
  /// See algorithm described in Section 7.2 of NIST SP800-38D:
  /// <https://nvlpubs.nist.gov/nistpubs/Legacy/SP/nistspecialpublication800-38d.pdf>
  ///
  /// > Define a block, J0, as follows:
  /// > If len(IV)=96, then J0 = IV || 0{31} || 1.
  /// > If len(IV) ≠ 96, then let s = 128 ⎡len(IV)/128⎤-len(IV), and
  /// >     J0=GHASH(IV||0s+64||[len(IV)]64).
  fn init_ctr(&self, nonce: &Nonce<NonceSize>) -> (Ctr32BE<&Aes>, Block) {
    let j0 = if NonceSize::to_usize() == 12 {
      let mut block = ghash::Block::default();
      block[..12].copy_from_slice(nonce);
      block[15] = 1;
      block
    } else {
      let mut ghash = self.ghash.clone();
      ghash.update_padded(nonce);

      let mut block = ghash::Block::default();
      let nonce_bits = (NonceSize::to_usize() as u64) * 8;
      block[8..].copy_from_slice(&nonce_bits.to_be_bytes());
      ghash.update(&[block]);
      ghash.finalize()
    };

    let mut ctr = Ctr32BE::inner_iv_init(&self.cipher, &j0);
    let mut tag_mask = Block::default();
    ctr.write_keystream_block(&mut tag_mask);
    (ctr, tag_mask)
  }

  /// Authenticate the given plaintext and associated data using GHASH.
  fn compute_tag(&self, mask: Block, associated_data: &[u8], buffer: &[u8]) -> Tag {
    let mut ghash = self.ghash.clone();
    ghash.update_padded(associated_data);
    ghash.update_padded(buffer);

    let associated_data_bits = (associated_data.len() as u64) * 8;
    let buffer_bits = (buffer.len() as u64) * 8;

    let mut block = ghash::Block::default();
    block[..8].copy_from_slice(&associated_data_bits.to_be_bytes());
    block[8..].copy_from_slice(&buffer_bits.to_be_bytes());
    ghash.update(&[block]);

    let mut tag = ghash.finalize();
    for (a, b) in tag.as_mut_slice().iter_mut().zip(mask.as_slice()) {
      *a ^= *b;
    }

    tag
  }
}
