arcium-core-utils 0.4.5

use std::fmt::Debug;

use primitives::{
    algebra::{
        field::{binary::Gf2_128, FieldExtension},
        BoxedUint,
    },
    correlated_randomness::{
        dabits::{DaBit, DaBits},
        pow::{PowPair, PowPairs},
        singlets::{Singlet, Singlets},
        triples::{Triple, Triples},
    },
    hashing::hash,
    random::{BaseRng, CryptoRngCore, Random, RandomWith, Seed, SeedableRng},
    sharing::GlobalFieldKey,
    types::Positive,
    utils::IntoExactSizeIterator,
};
use zeroize::Zeroize;

use super::{TestBatchSize, TrustedGenerator};

/// Trusted dealer for generating preprocessing elements for all peers (preprocessing shares must be
/// dealt to each peer).
#[derive(Clone, Debug)]
pub struct GlobalDealer<M = TestBatchSize> {
    n_parties: usize,
    alphas_seed: Seed,
    rng: BaseRng,

    _marker: std::marker::PhantomData<M>,
}

impl<M> Drop for GlobalDealer<M> {
    fn drop(&mut self) {
        self.alphas_seed.zeroize();
    }
}

impl<M> GlobalDealer<M> {
    /// Creates a global trusted dealer to generate correct preprocessing by seeding
    /// an Rng with the provided Seed.
    ///
    /// WARNING! The seed must be unique per session and disposed of after use
    /// to ensure security. Always prefer the `new` constructor unless you
    /// have a specific reason to manage the seed manually.
    pub fn new_with(n_parties: usize, alphas_seed: Seed) -> Self {
        let rng_seed = hash(&[alphas_seed.as_ref(), "global dealer".as_bytes()]).into();
        Self {
            n_parties,
            alphas_seed,
            rng: BaseRng::from_seed(rng_seed),
            _marker: std::marker::PhantomData,
        }
    }

    /// Creates a global trusted dealer to generate correct preprocessing.
    pub fn new(n_parties: usize) -> Self {
        let alphas_seed = Seed::random(&mut rand::thread_rng());
        Self::new_with(n_parties, alphas_seed)
    }

    pub fn n_parties(&self) -> usize {
        self.n_parties
    }
}

impl<M> RandomWith<usize> for GlobalDealer<M> {
    fn random_with(mut rng: impl CryptoRngCore, n_parties: usize) -> Self {
        Self {
            n_parties,
            alphas_seed: Seed::random(&mut rng),
            rng: BaseRng::from_rng(&mut rng).unwrap(),
            _marker: std::marker::PhantomData,
        }
    }
}

impl<M> GlobalDealer<M> {
    /// Returns the alphas for the current dealer. The alphas are generated
    /// using the session ID and the peer number.
    #[inline]
    pub fn get_alphas<F: FieldExtension>(&self) -> Vec<Vec<GlobalFieldKey<F>>> {
        #[cfg(not(feature = "compatibility"))]
        let mut alphas_rng = BaseRng::from_tagged_seed(self.alphas_seed, F::get_name());
        #[cfg(feature = "compatibility")]
        let mut alphas_rng = BaseRng::from_seed(self.alphas_seed);
        let alphas: Vec<Vec<GlobalFieldKey<F>>> = (0..self.n_parties)
            .map(|_| GlobalFieldKey::<F>::random_n(&mut alphas_rng, self.n_parties - 1))
            .collect();
        // DPF/SPFSS over binary global MAC keys requires the per-pair invariant
        //   lsb(Δ_ij) = 0 if i < j, else 1
        // where Δ_ij is the key party i holds for party j. This implies the
        // 2-party condition lsb(Δ_01 ⊕ Δ_10) = 1 and generalizes it for n > 2.
        #[cfg(not(feature = "compatibility"))]
        if F::get_name() == Gf2_128::get_name() {
            for (i, alphas_i) in alphas.iter_mut().enumerate() {
                for j in 0..self.n_parties {
                    if i == j {
                        continue;
                    }
                    let pos = if j < i { j } else { j - 1 };
                    let target_lsb = u8::from(i > j);
                    let current = alphas_i[pos].inner();
                    if current.to_le_bytes()[0] & 1 != target_lsb {
                        let tweaked = current + F::from(1u128);
                        alphas_i[pos] = GlobalFieldKey::new(FieldElement::new(tweaked));
                    }
                }
            }
        }
        alphas
    }
}

// ----------------------------- Trusted Dealer ----------------------------- //

impl<F: FieldExtension, M: Positive> TrustedGenerator<Singlet<F>> for GlobalDealer<M> {
    #[inline]
    fn generate_batch_for_each(
        &mut self,
        _associated_data: (),
    ) -> impl ExactSizeIterator<Item = impl IntoExactSizeIterator<Item = Singlet<F>>> {
        // Regenerate alphas upon each call
        let all_alphas = self.get_alphas::<F>();
        // Generate the singlet batch with fresh randomness
        let singlet_batch: Vec<_> = Singlets::<F, M>::random_n_with_each(&mut self.rng, all_alphas);
        singlet_batch.into_iter()
    }
}

impl<F: FieldExtension, M: Positive> TrustedGenerator<Triple<F>> for GlobalDealer<M> {
    #[inline]
    fn generate_batch_for_each(
        &mut self,
        _associated_data: (),
    ) -> impl ExactSizeIterator<Item = impl IntoExactSizeIterator<Item = Triple<F>>> {
        // Regenerate alphas upon each call
        let all_alphas = self.get_alphas::<F>();
        // Generate the triple batch with fresh randomness
        let triple_batch: Vec<_> = Triples::<F, M>::random_n_with_each(&mut self.rng, all_alphas);
        triple_batch.into_iter()
    }
}

impl<F: FieldExtension, M: Positive> TrustedGenerator<PowPair<F>> for GlobalDealer<M> {
    #[inline]
    fn generate_batch_for_each(
        &mut self,
        exponent: BoxedUint,
    ) -> impl ExactSizeIterator<Item = impl IntoExactSizeIterator<Item = PowPair<F>>> {
        // Regenerate alphas upon each call
        let all_alphas = self.get_alphas::<F>();
        // Generate the pow pair batch with fresh randomness
        let pow_pairs_batch: Vec<_> = PowPairs::<F, M>::random_n_with(
            &mut self.rng,
            self.n_parties,
            (exponent.clone(), all_alphas),
        );
        pow_pairs_batch.into_iter()
    }
}

impl<F: FieldExtension, M: Positive> TrustedGenerator<DaBit<F>> for GlobalDealer<M> {
    #[inline]
    fn generate_batch_for_each(
        &mut self,
        _associated_data: (),
    ) -> impl ExactSizeIterator<Item = impl IntoExactSizeIterator<Item = DaBit<F>>> {
        let alphas_f = self.get_alphas::<F>();
        let alphas_gf2 = self.get_alphas::<Gf2_128>();

        // Generate the daBit batch with fresh randomness
        let dabit_batch: Vec<_> =
            DaBits::<F, M>::random_n_with(&mut self.rng, self.n_parties, (alphas_f, alphas_gf2));
        dabit_batch.into_iter()
    }
}

#[cfg(test)]
mod tests {
    use primitives::{
        algebra::elliptic_curve::{Curve, Curve25519Ristretto as C},
        sharing::{Verifiable, VerifiableWith},
    };

    use super::*;
    use crate::preprocessing::dealer::mock::MockDealer;

    type M = typenum::U5;
    type F = <C as Curve>::Scalar;

    const N_PARTIES: usize = 3;

    pub fn generate_local_dealers<M: Positive>(n_parties: usize, seed: Seed) -> Vec<MockDealer<M>> {
        (0..n_parties)
            .map(|local_party_pos| MockDealer::new(n_parties, local_party_pos, seed))
            .collect()
    }

    #[tokio::test]
    async fn test_local_dealers_triples() {
        let seed = Seed::random(&mut primitives::random::test_rng());
        let n_triples = 11;

        // Create local dealers
        let mut dealers = generate_local_dealers::<M>(N_PARTIES, seed);

        // Generate triple batches asynchronously
        let mut triples: Vec<Vec<Triple<F>>> = vec![];
        for dealer in dealers.iter_mut() {
            triples.push(
                futures::future::try_join_all(dealer.request_n_triples(n_triples))
                    .await
                    .unwrap(),
            )
        }

        // Verify triples
        assert_eq!(triples.len(), N_PARTIES);
        triples.iter().for_each(|s| {
            assert_eq!(s.len(), n_triples);
        });
        for i in 0..n_triples {
            let triple = triples.iter().map(|s| s[i].clone()).collect::<Vec<_>>();
            Triple::verify_all(triple).unwrap();
        }
    }

    #[tokio::test]
    async fn test_local_dealers_pow_pairs() {
        let seed = Seed::random(&mut primitives::random::test_rng());
        let n_pow_pairs = 17;
        let exponent = BoxedUint::from(4u32);

        // Create local dealers
        let mut dealers = generate_local_dealers::<M>(N_PARTIES, seed);

        // Generate pow pair batches asynchronously
        let mut pow_pairs: Vec<Vec<PowPair<F>>> = vec![];
        for dealer in dealers.iter_mut() {
            pow_pairs.push(
                futures::future::try_join_all(
                    dealer.request_n_pow_pairs(n_pow_pairs, exponent.clone()),
                )
                .await
                .unwrap(),
            )
        }

        // Verify pow pairs
        assert_eq!(pow_pairs.len(), N_PARTIES);
        pow_pairs.iter().for_each(|s| {
            assert_eq!(s.len(), n_pow_pairs);
        });
        for i in 0..n_pow_pairs {
            let pow_pair_vec = pow_pairs.iter().map(|s| s[i].clone()).collect::<Vec<_>>();
            PowPair::verify_all_with(pow_pair_vec, exponent.clone()).unwrap();
        }
    }

    #[tokio::test]
    async fn test_global_queues_singlets() {
        let seed = Seed::random(&mut primitives::random::test_rng());
        let n_singlets = 13;

        // Create a global dealer
        let mut dealer = GlobalDealer::<M>::new_with(N_PARTIES, seed);

        // Generate single batch synchronously
        let singlets: Vec<Vec<Singlet<F>>> = dealer.generate_n_for_each(n_singlets, ());

        // Verify singlets
        assert_eq!(singlets.len(), N_PARTIES);
        singlets.iter().for_each(|s| {
            assert_eq!(s.len(), n_singlets);
        });
        for i in 0..n_singlets {
            let singlet = singlets.iter().map(|s| s[i].clone()).collect::<Vec<_>>();
            Singlet::verify_all(singlet).unwrap();
        }
    }

    #[tokio::test]
    async fn test_global_queues_triples() {
        let seed = Seed::random(&mut primitives::random::test_rng());
        let n_triples = 11;

        // Create local dealers
        let mut dealer = GlobalDealer::<M>::new_with(N_PARTIES, seed);

        // Generate triple batch synchronously
        let triples: Vec<Vec<Triple<F>>> = dealer.generate_n_for_each(n_triples, ());

        // Verify triples
        assert_eq!(triples.len(), N_PARTIES);
        triples.iter().for_each(|s| {
            assert_eq!(s.len(), n_triples);
        });
        for i in 0..n_triples {
            let triple = triples.iter().map(|s| s[i].clone()).collect::<Vec<_>>();
            Triple::verify_all(triple).unwrap();
        }
    }

    #[tokio::test]
    async fn test_global_dealer_singlets() {
        let seed = Seed::random(&mut primitives::random::test_rng());
        let n_singlets = 13;

        // Create a global dealer
        let mut dealer = GlobalDealer::<M>::new_with(N_PARTIES, seed);

        // Generate singlet batches
        let singlets: Vec<Vec<Singlet<F>>> = dealer.generate_n_for_each(n_singlets, ());

        // Verify singlets
        assert_eq!(singlets.len(), N_PARTIES);
        singlets.iter().for_each(|s| {
            assert_eq!(s.len(), n_singlets);
        });
        for i in 0..n_singlets {
            let singlet = singlets.iter().map(|s| s[i].clone()).collect::<Vec<_>>();
            Singlet::verify_all(singlet).unwrap();
        }
    }

    #[tokio::test]
    async fn test_global_dealer_triples() {
        let seed = Seed::random(&mut primitives::random::test_rng());
        let n_triples = 11;

        // Create local dealers
        let mut dealer = GlobalDealer::<M>::new_with(N_PARTIES, seed);
        // Generate triple batches
        let triples: Vec<Vec<Triple<F>>> = dealer.generate_n_for_each(n_triples, ());

        // Verify triples
        assert_eq!(triples.len(), N_PARTIES);
        triples.iter().for_each(|s| {
            assert_eq!(s.len(), n_triples);
        });
        for i in 0..n_triples {
            let triple = triples.iter().map(|s| s[i].clone()).collect::<Vec<_>>();
            Triple::verify_all(triple).unwrap();
        }
    }

    #[tokio::test]
    async fn test_global_dealer_pow_pairs() {
        let seed = Seed::random(&mut primitives::random::test_rng());
        let n_pow_pairs = 17;
        let exponent = BoxedUint::from(4u32);

        // Create local dealers
        let mut dealer = GlobalDealer::<M>::new_with(N_PARTIES, seed);

        // Generate pow pair batches
        let pow_pairs: Vec<Vec<PowPair<F>>> =
            dealer.generate_n_for_each(n_pow_pairs, exponent.clone());

        // Verify pow pairs
        assert_eq!(pow_pairs.len(), N_PARTIES);
        pow_pairs.iter().for_each(|s| {
            assert_eq!(s.len(), n_pow_pairs);
        });
        for i in 0..n_pow_pairs {
            let pow_pair_vec = pow_pairs.iter().map(|s| s[i].clone()).collect::<Vec<_>>();
            PowPair::verify_all_with(pow_pair_vec, exponent.clone()).unwrap();
        }
    }

    #[test]
    fn test_td_alphas_backwards_compatibility() {
        fn test_alphas_backwards_compatibility<const N: usize, F: FieldExtension>(
            v: Vec<Vec<GlobalFieldKey<F>>>,
            expected: [[u8; N]; 2],
        ) {
            assert_eq!(v.len(), expected.len());
            assert_eq!(v[0][0].to_le_bytes().as_slice(), expected[0]);
            assert_eq!(v[1][0].to_le_bytes().as_slice(), expected[1]);
        }
        let new = GlobalDealer::<M>::new_with(2, [0; 32]);
        let scalar_alphas = new.get_alphas::<F>();
        test_alphas_backwards_compatibility(
            scalar_alphas,
            [
                [
                    22, 33, 188, 127, 243, 114, 222, 165, 177, 158, 212, 131, 122, 34, 112, 164,
                    230, 48, 112, 90, 14, 78, 91, 42, 120, 206, 28, 215, 160, 190, 21, 0,
                ],
                [
                    207, 175, 72, 193, 159, 85, 167, 221, 77, 123, 109, 17, 179, 166, 249, 164, 54,
                    171, 24, 124, 120, 6, 218, 252, 132, 72, 92, 142, 19, 38, 161, 15,
                ],
            ],
        );
        let base_alphas = new.get_alphas::<<C as Curve>::BaseField>();
        test_alphas_backwards_compatibility(
            base_alphas,
            [
                [
                    36, 141, 17, 252, 235, 35, 94, 49, 239, 202, 72, 236, 42, 48, 10, 46, 101, 223,
                    11, 79, 69, 73, 151, 229, 218, 208, 23, 139, 95, 240, 195, 59,
                ],
                [
                    52, 191, 26, 57, 167, 171, 167, 220, 162, 248, 54, 82, 112, 87, 64, 232, 68,
                    36, 31, 28, 58, 201, 130, 209, 251, 161, 120, 21, 197, 121, 14, 106,
                ],
            ],
        );
        let binary_alphas = new.get_alphas::<Gf2_128>();
        test_alphas_backwards_compatibility(
            binary_alphas,
            [
                (6050961064690644123u128 + (15385182941806993281u128 << 64)).to_le_bytes(),
                (1474049585344358660u128 + (6851573923483025534u128 << 64)).to_le_bytes(),
            ],
        );
    }
}