// Copyright Open Logistics Foundation
//
// Licensed under the Open Logistics Foundation License 1.3.
// For details on the licensing terms, see the LICENSE file.
// SPDX-License-Identifier: OLFL-1.3

//! Random number generation module which defines the [`CtrDrbg`] RNG type

// Additionally, it defines the [`rng_try_fill_bytes_callback_fn`] used as callback function to
// `mbedtls_ssl_conf_rng`.

#[cfg(feature = "alloc")]
use alloc::boxed::Box;

use core::{marker::PhantomData, ops::DerefMut};

use cty::c_void;
use embedded_mbedtls_sys::{
    mbedtls_ctr_drbg_context, mbedtls_ctr_drbg_init, mbedtls_ctr_drbg_seed,
    MBEDTLS_ERR_CTR_DRBG_REQUEST_TOO_BIG, MBEDTLS_ERR_ENTROPY_SOURCE_FAILED,
};
use rand_core::{CryptoRng, RngCore};

use crate::error::Error;

/// This function is used as callback function to `mbedtls_ssl_conf_rng`. It can be used for
/// generic RNGs as well as for the [`CtrDrbg`] wrapper defined here.
pub(crate) unsafe extern "C" fn rng_try_fill_bytes_callback_fn<RNG: RngCore>(
    entropy_context: *mut cty::c_void,
    buf: *mut cty::c_uchar,
    buf_len: usize,
) -> cty::c_int {
    let rng: &mut RNG = &mut *(entropy_context as *mut RNG);
    let bytes = core::slice::from_raw_parts_mut(buf, buf_len);
    if rng.try_fill_bytes(bytes).is_err() {
        return embedded_mbedtls_sys::MBEDTLS_ERR_SSL_NO_RNG;
    }
    0
}

/// CTR_DRBG pseudorandom generator
///
/// Cryptographically secure pseudorandom number generator (CSPRNG).
/// The Mbed TLS implementation of CTR_DRBG which uses AES-256 (default) or AES-128.
///
/// Use this if:
/// - your RNG source isn't suitable for generating sufficient amounts of cryptographically secure
/// random numbers (e.g. your RNG isn't fast enough)
/// - if you do not have a true hardware RNG and want to use it as an RNG initialized with a
/// user-supplied (fixed) seed (use the `personalization_string` then); it should be obvious that
/// a non-random seed does not yield true random numbers and that reusing the same seed undermines
/// security
/// - if you want to augment your RNG with another source of entropy (again, use the
/// `personalization_string`)
///
/// This type is a wrapper around the `mbedtls_ctr_drbg_context`
pub struct CtrDrbg<'a, RNG: RngCore + CryptoRng, D: DerefMut<Target = RNG>> {
    context: mbedtls_ctr_drbg_context,
    entropy_source: D,
    _custom: PhantomData<&'a [u8]>,
}

impl<'a, RNG: RngCore + CryptoRng> CtrDrbg<'a, RNG, &'a mut RNG> {
    /// Initialize and seed a CtrDrbg context
    ///
    /// Might fail when no entropy can be collected.
    ///
    /// When the `alloc` feature is activated, the `new_with_heap_rng` constructor can be used to
    /// pass an owned entropy source which will be moved to the heap.
    /// This enables you to freely move the `CtrDrbg` together with the context.
    ///
    /// # Example
    /// ```
    /// use embedded_mbedtls::rng::CtrDrbg;
    /// use rand_core::RngCore;
    ///
    /// # let rng = rand::thread_rng();
    /// let mut entropy_src = rng;
    ///
    /// let mut ctr_drbg = CtrDrbg::new(&mut entropy_src, None).unwrap();
    ///
    /// let _random = ctr_drbg.next_u32();
    /// ```
    pub fn new(
        entropy_source: &'a mut RNG,
        personalization_string: Option<&'a [u8]>,
    ) -> Result<Self, Error> {
        Self::new_generic(entropy_source, personalization_string)
    }
}

#[cfg(feature = "alloc")]
impl<'a, RNG: RngCore + CryptoRng> CtrDrbg<'a, RNG, Box<RNG>> {
    /// Initialize and seed a CtrDrbg context, moving the entropy source into a `Box`
    ///
    /// This allows to move the [`CtrDrbg`] instance freely. Especially, it allows to return from
    /// an initializer function in which the entropy source was set up.
    ///
    /// Might fail when no entropy can be collected.
    ///
    /// Example:
    /// ```
    /// use embedded_mbedtls::rng::CtrDrbg;
    /// use rand_core::RngCore;
    ///
    /// # let rng = rand::thread_rng();
    /// let mut ctr_drbg = CtrDrbg::new_with_heap_rng(rng, None).unwrap();
    /// // The ctr_drbg can now be moved around freely
    ///
    /// let _random = ctr_drbg.next_u32();
    /// ```
    pub fn new_with_heap_rng(
        entropy_source: RNG,
        personalization_string: Option<&'a [u8]>,
    ) -> Result<Self, Error> {
        Self::new_generic(Box::new(entropy_source), personalization_string)
    }
}

impl<'a, RNG: RngCore + CryptoRng, D: DerefMut<Target = RNG>> CtrDrbg<'a, RNG, D> {
    /// Initialize and seed a CtrDrbg context
    ///
    /// _Note_: Internal function which should __not__ be made public.
    /// With this function the `entropy_source` is only constrained by `DerefMut<Target = RNG>`,
    /// which makes it possible to pass a type which implements DerefMut itself and won't guarantee that the `entropy_source` is never moved.
    /// This is crucial since the C context of the CTR_DRBG holds a raw pointer to it.
    fn new_generic(
        entropy_source: D,
        personalization_string: Option<&'a [u8]>,
    ) -> Result<Self, Error> {
        let context = mbedtls_ctr_drbg_context::default();
        let mut this = Self {
            context,
            entropy_source,
            _custom: PhantomData,
        };
        unsafe { mbedtls_ctr_drbg_init(&mut this.context) };
        if let Some(custom) = personalization_string {
            let ret = unsafe {
                mbedtls_ctr_drbg_seed(
                    &mut this.context,
                    Some(rng_try_fill_bytes_callback_fn::<RNG>),
                    this.entropy_source.deref_mut() as *mut RNG as *mut cty::c_void,
                    custom.as_ptr(),
                    custom.len(),
                )
            };
            if ret != 0 {
                return Err(ret.into());
            }
        } else {
            let ret = unsafe {
                mbedtls_ctr_drbg_seed(
                    &mut this.context,
                    Some(rng_try_fill_bytes_callback_fn::<RNG>),
                    this.entropy_source.deref_mut() as *mut RNG as *mut cty::c_void,
                    core::ptr::null(),
                    0,
                )
            };
            if ret != 0 {
                return Err(ret.into());
            }
        }

        Ok(this)
    }
}

impl<R: RngCore + CryptoRng, D: DerefMut<Target = R>> Drop for CtrDrbg<'_, R, D> {
    fn drop(&mut self) {
        unsafe {
            embedded_mbedtls_sys::mbedtls_ctr_drbg_free(&mut self.context);
        }
    }
}

impl<RNG: RngCore + CryptoRng, D: DerefMut<Target = RNG>> CryptoRng for CtrDrbg<'_, RNG, D> {}

impl<RNG: RngCore + CryptoRng, D: DerefMut<Target = RNG>> RngCore for CtrDrbg<'_, RNG, D> {
    fn next_u32(&mut self) -> u32 {
        rand_core::impls::next_u32_via_fill(self)
    }

    fn next_u64(&mut self) -> u64 {
        rand_core::impls::next_u64_via_fill(self)
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        if dest.len() > embedded_mbedtls_sys::MBEDTLS_CTR_DRBG_MAX_REQUEST as usize {
            log::error!("Failed to generate random data: Request too big!");
            panic!("Failed to generate random data: Request too big!");
        }
        if self.try_fill_bytes(dest).is_err() {
            panic!("Failed to generate random data: Entropy source failed!");
        }
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand_core::Error> {
        let ret = unsafe {
            embedded_mbedtls_sys::mbedtls_ctr_drbg_random(
                &mut self.context as *mut mbedtls_ctr_drbg_context as *mut c_void,
                dest.as_mut_ptr(),
                dest.len(),
            )
        };

        if ret == MBEDTLS_ERR_CTR_DRBG_REQUEST_TOO_BIG {
            log::error!("Failed to generate random data: Request to too big!");
            use core::num::NonZeroU32;
            use rand_core::Error;
            return Err(Error::from(unsafe {
                NonZeroU32::new_unchecked(Error::CUSTOM_START)
            }));
        }
        if ret == MBEDTLS_ERR_ENTROPY_SOURCE_FAILED {
            log::error!("Failed to generate random data: Entropy source failed!");
            use core::num::NonZeroU32;
            use rand_core::Error;
            return Err(Error::from(unsafe {
                NonZeroU32::new_unchecked(Error::CUSTOM_START + 1)
            }));
        }
        if ret < 0 {
            log::error!("Failed to generate random data: mbedtls error {ret}");
            use core::num::NonZeroU32;
            use rand_core::Error;
            return Err(Error::from(unsafe {
                NonZeroU32::new_unchecked(Error::CUSTOM_START + 2)
            }));
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use rand_core::RngCore;

    use super::CtrDrbg;

    #[test]
    fn stack_entropy_drbg() {
        let mut entropy_source = rand::thread_rng();
        let mut ctr_drbg = CtrDrbg::new(&mut entropy_source, None).unwrap();

        let _random = ctr_drbg.next_u32();
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn boxed_entropy_drbg() {
        let mut ctr_drbg = CtrDrbg::new_with_heap_rng(Box::new(rand::thread_rng()), None).unwrap();

        let _random = ctr_drbg.next_u32();
    }
}