//! Hash-based Key Derivation Function
//! We just repackage other crates

use ::sha3::Sha3_256;
use ::zeroize::Zeroize;

use crate::enc::sym::Secret;

/// Kind of HKDF
#[derive(Debug, Copy, Clone, PartialEq, ::num_derive::FromPrimitive)]
#[non_exhaustive]
#[repr(u8)]
pub enum HkdfKind {
    /// Sha3
    Sha3 = 0,
}

/// Generic wrapper on Hkdfs
#[derive(Clone)]
pub enum Hkdf {
    /// Sha3 based
    Sha3(HkdfSha3),
}

// Fake debug implementation to avoid leaking secrets
impl ::core::fmt::Debug for Hkdf {
    fn fmt(
        &self,
        f: &mut core::fmt::Formatter<'_>,
    ) -> Result<(), ::std::fmt::Error> {
        ::core::fmt::Debug::fmt("[hidden hkdf]", f)
    }
}

impl Hkdf {
    /// New Hkdf
    pub fn new(kind: HkdfKind, salt: &[u8], key: Secret) -> Self {
        match kind {
            HkdfKind::Sha3 => Self::Sha3(HkdfSha3::new(salt, key)),
        }
    }
    /// Get a secret generated from the key and a given context
    pub fn get_secret(&self, context: &[u8]) -> Secret {
        match self {
            Hkdf::Sha3(sha3) => sha3.get_secret(context),
        }
    }
}

// Hack & tricks:
// HKDF are pretty important, but this lib don't zero out the data.
// we can't use #[derive(Zeroing)] either.
// So we craete a union with a Zeroing object, and drop both manually.

#[derive(Zeroize)]
#[zeroize(drop)]
struct Zeroable([u8; ::core::mem::size_of::<::hkdf::Hkdf<Sha3_256>>()]);

union HkdfInner {
    hkdf: ::core::mem::ManuallyDrop<::hkdf::Hkdf<Sha3_256>>,
    zeroable: ::core::mem::ManuallyDrop<Zeroable>,
}

impl Drop for HkdfInner {
    fn drop(&mut self) {
        #[allow(unsafe_code)]
        unsafe {
            drop(&mut self.hkdf);
            drop(&mut self.zeroable);
        }
    }
}
impl Clone for HkdfInner {
    fn clone(&self) -> Self {
        #[allow(unsafe_code)]
        unsafe {
            Self {
                hkdf: self.hkdf.clone(),
            }
        }
    }
}

/// Sha3 based HKDF
#[derive(Clone)]
pub struct HkdfSha3 {
    inner: HkdfInner,
}

impl HkdfSha3 {
    /// Instantiate a new HKDF with Sha3-256
    pub fn new(salt: &[u8], key: Secret) -> Self {
        let hkdf = ::hkdf::Hkdf::<Sha3_256>::new(Some(salt), key.as_ref());
        Self {
            inner: HkdfInner {
                hkdf: ::core::mem::ManuallyDrop::new(hkdf),
            },
        }
    }
    /// Get a secret generated from the key and a given context
    pub fn get_secret(&self, context: &[u8]) -> Secret {
        let mut out: [u8; 32] = [0; 32];
        #[allow(unsafe_code)]
        unsafe {
            self.inner.hkdf.expand(context, &mut out);
        }
        out.into()
    }
}

// Fake debug implementation to avoid leaking secrets
impl ::core::fmt::Debug for HkdfSha3 {
    fn fmt(
        &self,
        f: &mut core::fmt::Formatter<'_>,
    ) -> Result<(), ::std::fmt::Error> {
        ::core::fmt::Debug::fmt("[hidden hkdf]", f)
    }
}