arcium-primitives 0.4.2

use aes::{cipher::KeyInit, Aes128Enc};
use blanket::blanket;
use hybrid_array::Array;

use crate::{
    algebra::{field::binary::Gf2_128, uniform_bytes::FromUniformBytes},
    ciphers::BlockCipher,
    types::SessionId,
};

#[blanket(derive(Mut))]
pub trait TweakableHasher {
    fn set_tweak(&mut self, tweak: u64);
    fn get_tweak(&self) -> u64;
    fn update_tweak(&mut self);
    fn hash(&self, val: &Gf2_128) -> Gf2_128;
    fn hash_to_bytes(&self, val: &Gf2_128, buf: &mut [u8]);
    fn hash_to_uniform_bytes<F: FromUniformBytes>(&self, val: &Gf2_128) -> F {
        let mut output = Array::<u8, F::UniformBytes>::default();
        self.hash_to_bytes(val, &mut output);
        F::from_uniform_bytes(&output)
    }
    fn hash_update_tweak(&mut self, val: &Gf2_128) -> Gf2_128 {
        let h = self.hash(val);
        self.update_tweak();
        h
    }
    fn hash_to_uniform_bytes_update_tweak<F: FromUniformBytes>(&mut self, val: &Gf2_128) -> F {
        let h = self.hash_to_uniform_bytes(val);
        self.update_tweak();
        h
    }
    fn refresh(&mut self, session_id: &SessionId);
}

/// Tweakable circular correlation-resistant hash function built from a block cipher.
///
/// ** Note: ** This function implements the TMMO hash from
/// Guo et al. "Efficient and Secure Multiparty Computation from Fixed-Key Block Ciphers"
pub struct HasherTccr {
    enc: Aes128Enc,
    tweak: u64,
}

impl HasherTccr {
    pub fn new(session_id: &SessionId) -> Self {
        Self {
            enc: Aes128Enc::new(&session_id.into()),
            tweak: 0,
        }
    }
}

impl TweakableHasher for HasherTccr {
    fn set_tweak(&mut self, tweak: u64) {
        self.tweak = tweak;
    }

    fn get_tweak(&self) -> u64 {
        self.tweak
    }

    fn update_tweak(&mut self) {
        self.tweak += 1;
    }

    fn hash(&self, val: &Gf2_128) -> Gf2_128 {
        // π(π(msg) ^ tweak) ^ π(val) where π is the AES permutation
        let t = self.enc.encrypt(val);
        self.enc.encrypt(&(t + Gf2_128::from_u64(self.tweak))) + t
    }

    fn hash_to_bytes(&self, val: &Gf2_128, buffer: &mut [u8]) {
        let t = self.enc.encrypt(val);
        buffer.chunks_mut(16).enumerate().for_each(|(k, chunk)| {
            let tweak = Gf2_128::from_limbs([self.tweak, k as u64]);
            let h = self.enc.encrypt(&(t + tweak)) + t;
            let tmp = h.into_le_block();
            chunk.copy_from_slice(&tmp.as_slice()[0..chunk.len()]);
        });
    }

    fn refresh(&mut self, session_id: &SessionId) {
        self.enc = Aes128Enc::new(&session_id.into());
        self.tweak = 0;
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        algebra::{
            elliptic_curve::{Curve25519Ristretto, ScalarField},
            field::binary::Gf2_128,
        },
        hashing::{HasherTccr, TweakableHasher},
        random::Random,
        types::SessionId,
    };

    type Fq = ScalarField<Curve25519Ristretto>;

    #[test]
    fn test_hash_tccr() {
        let mut rng = crate::random::test_rng();
        let mut hasher1 = HasherTccr::new(&SessionId::random(&mut rng));
        let hasher2 = HasherTccr::new(&SessionId::random(&mut rng));

        // Hash to Gf2_128
        assert_eq!(
            hasher1.hash(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash(&Gf2_128::from_limbs([42, 24])),
        );

        assert_ne!(
            hasher1.hash(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash(&Gf2_128::from_limbs([43, 24])),
        );

        assert_ne!(
            hasher1.hash(&Gf2_128::from_limbs([42, 24])),
            hasher2.hash(&Gf2_128::from_limbs([43, 24])),
        );

        assert_ne!(
            hasher1.hash_update_tweak(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash(&Gf2_128::from_limbs([42, 24])),
        );

        assert_eq!(
            hasher1.hash(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash_update_tweak(&Gf2_128::from_limbs([42, 24])),
        );

        // To field Gf2_128
        assert_eq!(
            hasher1.hash_to_uniform_bytes::<Gf2_128>(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash_to_uniform_bytes::<Gf2_128>(&Gf2_128::from_limbs([42, 24])),
        );

        assert_ne!(
            hasher1.hash_to_uniform_bytes::<Gf2_128>(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash_to_uniform_bytes::<Gf2_128>(&Gf2_128::from_limbs([43, 24])),
        );

        assert_ne!(
            hasher1.hash_to_uniform_bytes::<Gf2_128>(&Gf2_128::from_limbs([42, 24])),
            hasher2.hash_to_uniform_bytes::<Gf2_128>(&Gf2_128::from_limbs([43, 24])),
        );

        assert_ne!(
            hasher1.hash_to_uniform_bytes_update_tweak::<Gf2_128>(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash_to_uniform_bytes::<Gf2_128>(&Gf2_128::from_limbs([42, 24])),
        );

        assert_eq!(
            hasher1.hash_to_uniform_bytes::<Gf2_128>(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash_to_uniform_bytes_update_tweak::<Gf2_128>(&Gf2_128::from_limbs([42, 24])),
        );

        // To field Fq
        assert_eq!(
            hasher1.hash_to_uniform_bytes::<Fq>(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash_to_uniform_bytes::<Fq>(&Gf2_128::from_limbs([42, 24])),
        );

        assert_ne!(
            hasher1.hash_to_uniform_bytes::<Fq>(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash_to_uniform_bytes::<Fq>(&Gf2_128::from_limbs([43, 24])),
        );

        assert_ne!(
            hasher1.hash_to_uniform_bytes::<Fq>(&Gf2_128::from_limbs([42, 24])),
            hasher2.hash_to_uniform_bytes::<Fq>(&Gf2_128::from_limbs([43, 24])),
        );

        assert_ne!(
            hasher1.hash_to_uniform_bytes_update_tweak::<Fq>(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash_to_uniform_bytes::<Fq>(&Gf2_128::from_limbs([42, 24])),
        );

        assert_eq!(
            hasher1.hash_to_uniform_bytes::<Fq>(&Gf2_128::from_limbs([42, 24])),
            hasher1.hash_to_uniform_bytes_update_tweak::<Fq>(&Gf2_128::from_limbs([42, 24])),
        );
    }
}