// Copyright 2018-2024 Brian Smith.

//

// Permission to use, copy, modify, and/or distribute this software for any

// purpose with or without fee is hereby granted, provided that the above

// copyright notice and this permission notice appear in all copies.

//

// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES

// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF

// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY

// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES

// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION

// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN

// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

use super::{Block, KeyBytes, Overlapping, BLOCK_LEN};

use crate::{bits::BitLength, c, error};

use core::{

    ffi::{c_int, c_uint},

    num::{NonZeroU32, NonZeroUsize},

};

/// nonce || big-endian counter.

#[repr(transparent)]

pub(in super::super) struct Counter(pub(super) [u8; BLOCK_LEN]);

// Keep this in sync with AES_KEY in aes.h.

#[repr(C)]

#[derive(Clone)]

pub(in super::super) struct AES_KEY {

    pub rd_key: [u32; 4 * (MAX_ROUNDS + 1)],

    pub rounds: c_uint,

}

// Keep this in sync with `AES_MAXNR` in aes.h.

const MAX_ROUNDS: usize = 14;

impl AES_KEY {

    #[inline]

    pub(super) unsafe fn new(

        f: unsafe extern "C" fn(*const u8, BitLength<c_int>, *mut AES_KEY) -> c_int,

        bytes: KeyBytes<'_>,

    ) -> Result<Self, error::Unspecified> {

        let mut key = Self {

            rd_key: [0; 4 * (MAX_ROUNDS + 1)],

            rounds: 0,

        };

        let (bytes, key_bits) = match bytes {

            KeyBytes::AES_128(bytes) => (&bytes[..], BitLength::from_bits(128)),

            KeyBytes::AES_256(bytes) => (&bytes[..], BitLength::from_bits(256)),

        };

        // Unusually, in this case zero means success and non-zero means failure.

        if 0 == unsafe { f(bytes.as_ptr(), key_bits, &mut key) } {

            debug_assert_ne!(key.rounds, 0); // Sanity check initialization.

            Ok(key)

        } else {

            Err(error::Unspecified)

        }

    }

}

#[cfg(all(target_arch = "arm", target_endian = "little"))]

impl AES_KEY {

    pub(super) unsafe fn derive(

        f: for<'a> unsafe extern "C" fn(*mut AES_KEY, &'a AES_KEY),

        src: &Self,

    ) -> Self {

        let mut r = AES_KEY {

            rd_key: [0u32; 4 * (MAX_ROUNDS + 1)],

            rounds: 0,

        };

        unsafe { f(&mut r, src) };

        r

    }

    pub(super) fn rounds(&self) -> u32 {

        self.rounds

    }

}

// SAFETY:

//  * The function `$name` must read `bits` bits from `user_key`; `bits` will

//    always be a valid AES key length, i.e. a whole number of bytes.

//  * `$name` must set `key.rounds` to the value expected by the corresponding

//    encryption/decryption functions and return 0, or otherwise must return

//    non-zero to indicate failure.

//  * `$name` may inspect CPU features.

//

// In BoringSSL, the C prototypes for these are in

// crypto/fipsmodule/aes/internal.h.

macro_rules! set_encrypt_key {

    ( $name:ident, $key_bytes:expr $(,)? ) => {{

        use crate::bits::BitLength;

        use core::ffi::c_int;

        prefixed_extern! {

            fn $name(user_key: *const u8, bits: BitLength<c_int>, key: *mut AES_KEY) -> c_int;

        }

        $crate::aead::aes::ffi::AES_KEY::new($name, $key_bytes)

    }};

}

macro_rules! encrypt_block {

    ($name:ident, $block:expr, $key:expr) => {{

        use crate::aead::aes::{ffi::AES_KEY, Block};

        prefixed_extern! {

            fn $name(a: &Block, r: *mut Block, key: &AES_KEY);

        }

        $key.encrypt_block($name, $block)

    }};

}

impl AES_KEY {

    #[inline]

    pub(super) unsafe fn encrypt_block(

        &self,

        f: unsafe extern "C" fn(&Block, *mut Block, &AES_KEY),

        a: Block,

    ) -> Block {

        let mut result = core::mem::MaybeUninit::uninit();

        unsafe {

            f(&a, result.as_mut_ptr(), self);

            result.assume_init()

        }

    }

}

/// SAFETY:

///   * The caller must ensure that `$key` was initialized with the

///     `set_encrypt_key!` invocation that `$name` requires.

///   * The caller must ensure that fhe function `$name` satisfies the conditions

///     for the `f` parameter to `ctr32_encrypt_blocks`.

macro_rules! ctr32_encrypt_blocks {

    ($name:ident, $in_out:expr, $key:expr, $ctr:expr $(,)? ) => {{

        use crate::{

            aead::aes::{ffi::AES_KEY, Counter, BLOCK_LEN},

            c,

        };

        prefixed_extern! {

            fn $name(

                input: *const [u8; BLOCK_LEN],

                output: *mut [u8; BLOCK_LEN],

                blocks: c::NonZero_size_t,

                key: &AES_KEY,

                ivec: &Counter,

            );

        }

        $key.ctr32_encrypt_blocks($name, $in_out, $ctr)

    }};

}

impl AES_KEY {

    /// SAFETY:

    ///   * `f` must not read more than `blocks` blocks from `input`.

    ///   * `f` must write exactly `block` blocks to `output`.

    ///   * In particular, `f` must handle blocks == 0 without reading from `input`

    ///     or writing to `output`.

    ///   * `f` must support the input overlapping with the output exactly or

    ///     with any nonnegative offset `n` (i.e. `input == output.add(n)`);

    ///     `f` does NOT need to support the cases where input < output.

    ///   * `key` must have been initialized with the `set_encrypt_key!` invocation

    ///      that corresponds to `f`.

    ///   * `f` may inspect CPU features.

    #[inline]

    pub(super) unsafe fn ctr32_encrypt_blocks(

        &self,

        f: unsafe extern "C" fn(

            input: *const [u8; BLOCK_LEN],

            output: *mut [u8; BLOCK_LEN],

            blocks: c::NonZero_size_t,

            key: &AES_KEY,

            ivec: &Counter,

        ),

        in_out: Overlapping<'_>,

        ctr: &mut Counter,

    ) {

        in_out.with_input_output_len(|input, output, len| {

            debug_assert_eq!(len % BLOCK_LEN, 0);

            let blocks = match NonZeroUsize::new(len / BLOCK_LEN) {

                Some(blocks) => blocks,

                None => {

                    return;

                }

            };

            let input: *const [u8; BLOCK_LEN] = input.cast();

            let output: *mut [u8; BLOCK_LEN] = output.cast();

            let blocks_u32: NonZeroU32 = blocks.try_into().unwrap();

            // SAFETY:

            //  * `input` points to `blocks` blocks.

            //  * `output` points to space for `blocks` blocks to be written.

            //  * input == output.add(n), where n == src.start, and the caller is

            //    responsible for ensuing this sufficient for `f` to work correctly.

            //  * `blocks` is non-zero so `f` doesn't have to work for empty slices.

            //  * The caller is responsible for ensuring `key` was initialized by the

            //    `set_encrypt_key!` invocation required by `f`.

            unsafe {

                f(input, output, blocks, self, ctr);

            }

            ctr.increment_by_less_safe(blocks_u32);

        });

    }

}