arcis-compiler 0.14.1

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        bounds::FieldBounds,
        circuits::{
            boolean::{
                boolean_value::{Boolean, BooleanValue},
                byte::Byte,
                sha2::Sha256,
                sha3::{Sha3_256, Sha3_512},
            },
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    utils::{
        crypto::{rescue_desc::RescueArg, rescue_prime_hash::RescuePrimeHash},
        field::BaseField,
    },
};
use hmac::Mac;

pub trait Hmac<T> {
    fn digest(&self, key: Vec<T>, message: Vec<T>) -> Vec<T>;
}

/// The Arcis Hmac-Rescue-Prime function.
#[derive(Debug)]
#[allow(non_camel_case_types)]
pub struct Hmac_RescuePrime<F: ActuallyUsedField, T: RescueArg<F>> {
    pub hasher: RescuePrimeHash<F, T>,
}

impl<F: ActuallyUsedField, T: RescueArg<F>> Hmac_RescuePrime<F, T> {
    pub fn new() -> Self {
        Self {
            hasher: RescuePrimeHash::new(),
        }
    }
}

impl<F: ActuallyUsedField, T: RescueArg<F>> Hmac<T> for Hmac_RescuePrime<F, T> {
    fn digest(&self, key: Vec<T>, message: Vec<T>) -> Vec<T> {
        // We follow https://datatracker.ietf.org/doc/html/rfc2104, though since Rescue-Prime is not based
        // on the Merkle-Damgard construction we cannot have an exact anology between the
        // parameters. For our purpose, we set B = hasher.rate and L = hasher.digest_len.
        // Note: since Rescue-Prime is based on the sponge construction, the above
        // mentioned HMAC could be replaced by the simpler keyed hash function
        // hasher.digest(key || message), see e.g. https://keccak.team/files/SpongeFunctions.pdf (Appendix B) or
        // https://keccak.team/keccak_strengths.html.
        // TODO: we can decide later if we want to go with the simpler MAC or not.
        assert!(
            self.hasher.digest_len <= key.len() && key.len() <= self.hasher.rate,
            "The length of the key is supposed to be at least the hash function's digest length and at most the hash function's rate (found key length: {}, digest length: {} and rate: {})",
            key.len(),
            self.hasher.digest_len,
            self.hasher.rate
        );
        let ipad = T::from(F::from_le_bytes([0x36; 32]));
        let opad = T::from(F::from_le_bytes([0x5c; 32]));
        // the key is first extended to length B
        let mut key_extended = key;
        key_extended.resize(self.hasher.rate, T::from(F::ZERO));
        // inner padding
        let mut key_plus_ipad = key_extended.iter().map(|k| *k + ipad).collect::<Vec<T>>();
        key_plus_ipad.extend(message);
        let inner_digest = self.hasher.digest(key_plus_ipad).to_vec();
        // outer padding
        let mut key_plus_opad = key_extended.iter().map(|k| *k + opad).collect::<Vec<T>>();
        key_plus_opad.extend(inner_digest);
        self.hasher.digest(key_plus_opad).to_vec()
    }
}

impl<F: ActuallyUsedField, T: RescueArg<F>> Default for Hmac_RescuePrime<F, T> {
    fn default() -> Self {
        Self::new()
    }
}

macro_rules! impl_hmac {
    ($t: ident, $hasher: ident) => {
        /// The Arcis Hmac-$hasher function.
        #[derive(Clone, Debug)]
        #[allow(non_camel_case_types)]
        pub struct $t {
            pub hasher: $hasher,
        }

        impl $t {
            pub fn new() -> Self {
                Self { hasher: $hasher::new() }
            }
        }

        impl<B: Boolean> Hmac<Byte<B>> for $t {
            fn digest(&self, key: Vec<Byte<B>>, message: Vec<Byte<B>>) -> Vec<Byte<B>> {
                // For our purpose, B = $hasher::BLOCK_BYTES and L = $hasher::DIGEST_BYTES.
                assert!(
                    $hasher::DIGEST_BYTES <= key.len() && key.len() <= $hasher::BLOCK_BYTES,
                    "The key length is supposed to be at least the hash function's digest length and at most the hash function's block length (found key length: {}, digest length: {} and block length: {})",
                    key.len(),
                    $hasher::DIGEST_BYTES,
                    $hasher::BLOCK_BYTES
                );
                let ipad = Byte::from(0x36);
                let opad = Byte::from(0x5C);
                // the key is first extended to length B
                let mut key_extended = key;
                key_extended.resize($hasher::BLOCK_BYTES, Byte::from(0u8));
                // inner padding
                let mut key_xor_ipad = key_extended.iter().map(|k| *k ^ ipad).collect::<Vec<Byte<B>>>();
                key_xor_ipad.extend(&message);
                let inner_digest = self.hasher.digest(key_xor_ipad).to_vec();
                // outer padding
                let mut key_xor_opad = key_extended.iter().map(|k| *k ^ opad).collect::<Vec<Byte<B>>>();
                key_xor_opad.extend(inner_digest);
                self.hasher.digest(key_xor_opad).to_vec()
            }
        }

        impl Default for $t {
            fn default() -> Self {
                Self::new()
            }
        }
    };
}

impl_hmac!(Hmac_Sha256, Sha256);
impl_hmac!(Hmac_Sha3_256, Sha3_256);
impl_hmac!(Hmac_Sha3_512, Sha3_512);

impl ArithmeticCircuit<BaseField> for Hmac_Sha256 {
    fn eval(&self, x: Vec<BaseField>) -> Result<Vec<BaseField>, EvalFailure> {
        // all inputs are expected to be bytes
        x.iter()
            .for_each(|byte| assert!(*byte <= BaseField::from(255)));
        let mut key = x
            .into_iter()
            .map(|val| val.to_le_bytes()[0])
            .collect::<Vec<u8>>();
        let message = key.split_off(Sha256::DIGEST_BYTES);

        let mut mac = hmac::Hmac::<sha2::Sha256>::new_from_slice(&key)
            .expect("HMAC can take key of any size");
        mac.update(&message);

        Ok(mac
            .finalize()
            .into_bytes()
            .iter()
            .map(|byte| BaseField::from(*byte as u64))
            .collect::<Vec<BaseField>>())
    }

    fn bounds(&self, _bounds: Vec<FieldBounds<BaseField>>) -> Vec<FieldBounds<BaseField>> {
        vec![FieldBounds::new(BaseField::from(0), BaseField::from(255)); Sha256::DIGEST_BYTES]
    }

    fn run(&self, vals: Vec<FieldValue<BaseField>>) -> Vec<FieldValue<BaseField>> {
        let mut key = vals
            .into_iter()
            .map(Byte::from)
            .collect::<Vec<Byte<BooleanValue>>>();
        let message = key.split_off(Sha256::DIGEST_BYTES);

        let hmac = Hmac_Sha256::new();
        let mac = hmac.digest(key, message);

        mac.into_iter()
            .map(FieldValue::<BaseField>::from)
            .collect::<Vec<FieldValue<BaseField>>>()
    }
}

impl ArithmeticCircuit<BaseField> for Hmac_Sha3_256 {
    fn eval(&self, x: Vec<BaseField>) -> Result<Vec<BaseField>, EvalFailure> {
        // all inputs are expected to be bytes
        x.iter()
            .for_each(|byte| assert!(*byte <= BaseField::from(255)));
        let mut key = x
            .into_iter()
            .map(|val| val.to_le_bytes()[0])
            .collect::<Vec<u8>>();
        let message = key.split_off(Sha3_256::DIGEST_BYTES);

        let mut mac = hmac::Hmac::<sha3::Sha3_256>::new_from_slice(&key)
            .expect("HMAC can take key of any size");
        mac.update(&message);

        Ok(mac
            .finalize()
            .into_bytes()
            .iter()
            .map(|byte| BaseField::from(*byte as u64))
            .collect::<Vec<BaseField>>())
    }

    fn bounds(&self, _bounds: Vec<FieldBounds<BaseField>>) -> Vec<FieldBounds<BaseField>> {
        vec![FieldBounds::new(BaseField::from(0), BaseField::from(255)); Sha3_256::DIGEST_BYTES]
    }

    fn run(&self, vals: Vec<FieldValue<BaseField>>) -> Vec<FieldValue<BaseField>> {
        let mut key = vals
            .into_iter()
            .map(Byte::from)
            .collect::<Vec<Byte<BooleanValue>>>();
        let message = key.split_off(Sha3_256::DIGEST_BYTES);

        let hmac = Hmac_Sha3_256::new();
        let mac = hmac.digest(key, message);

        mac.into_iter()
            .map(FieldValue::<BaseField>::from)
            .collect::<Vec<FieldValue<BaseField>>>()
    }
}

impl ArithmeticCircuit<BaseField> for Hmac_Sha3_512 {
    fn eval(&self, x: Vec<BaseField>) -> Result<Vec<BaseField>, EvalFailure> {
        // all inputs are expected to be bytes
        x.iter()
            .for_each(|byte| assert!(*byte <= BaseField::from(255)));
        let mut key = x
            .into_iter()
            .map(|val| val.to_le_bytes()[0])
            .collect::<Vec<u8>>();
        let message = key.split_off(Sha3_512::DIGEST_BYTES);

        let mut mac = hmac::Hmac::<sha3::Sha3_512>::new_from_slice(&key)
            .expect("HMAC can take key of any size");
        mac.update(&message);

        Ok(mac
            .finalize()
            .into_bytes()
            .iter()
            .map(|byte| BaseField::from(*byte as u64))
            .collect::<Vec<BaseField>>())
    }

    fn bounds(&self, _bounds: Vec<FieldBounds<BaseField>>) -> Vec<FieldBounds<BaseField>> {
        vec![FieldBounds::new(BaseField::from(0), BaseField::from(255)); Sha3_512::DIGEST_BYTES]
    }

    fn run(&self, vals: Vec<FieldValue<BaseField>>) -> Vec<FieldValue<BaseField>> {
        let mut key = vals
            .into_iter()
            .map(Byte::from)
            .collect::<Vec<Byte<BooleanValue>>>();
        let message = key.split_off(Sha3_512::DIGEST_BYTES);

        let hmac = Hmac_Sha3_512::new();
        let mac = hmac.digest(key, message);

        mac.into_iter()
            .map(FieldValue::<BaseField>::from)
            .collect::<Vec<FieldValue<BaseField>>>()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::circuits::traits::arithmetic_circuit::tests::TestedArithmeticCircuit;
    use rand::Rng;

    impl TestedArithmeticCircuit<BaseField> for Hmac_Sha256 {
        fn gen_desc<R: Rng + ?Sized>(_rng: &mut R) -> Self {
            Self::new()
        }

        fn gen_n_inputs<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
            Sha256::DIGEST_BYTES
                + TryInto::<usize>::try_into(rng.next_u32()).unwrap() % Sha256::BLOCK_BYTES
        }

        fn gen_input_bounds<R: Rng + ?Sized>(_rng: &mut R) -> FieldBounds<BaseField> {
            FieldBounds::new(BaseField::from(0), BaseField::from(255))
        }
    }

    impl TestedArithmeticCircuit<BaseField> for Hmac_Sha3_256 {
        fn gen_desc<R: Rng + ?Sized>(_rng: &mut R) -> Self {
            Self::new()
        }

        fn gen_n_inputs<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
            Sha3_256::DIGEST_BYTES
                + TryInto::<usize>::try_into(rng.next_u32()).unwrap() % Sha3_256::BLOCK_BYTES
        }

        fn gen_input_bounds<R: Rng + ?Sized>(_rng: &mut R) -> FieldBounds<BaseField> {
            FieldBounds::new(BaseField::from(0), BaseField::from(255))
        }
    }

    impl TestedArithmeticCircuit<BaseField> for Hmac_Sha3_512 {
        fn gen_desc<R: Rng + ?Sized>(_rng: &mut R) -> Self {
            Self::new()
        }

        fn gen_n_inputs<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
            Sha3_512::DIGEST_BYTES
                + TryInto::<usize>::try_into(rng.next_u32()).unwrap() % Sha3_512::BLOCK_BYTES
        }

        fn gen_input_bounds<R: Rng + ?Sized>(_rng: &mut R) -> FieldBounds<BaseField> {
            FieldBounds::new(BaseField::from(0), BaseField::from(255))
        }
    }

    #[test]
    fn tested_hmac_sha256() {
        Hmac_Sha256::test(1, 1)
    }

    #[test]
    fn tested_hmac_sha3_256() {
        Hmac_Sha3_256::test(1, 1)
    }

    #[test]
    fn tested_hmac_sha3_512() {
        Hmac_Sha3_512::test(1, 1)
    }
}