arcium-primitives 0.4.4

#[cfg(any(test, feature = "dev"))]
pub use test::test_rng;
use typenum::Unsigned;

use crate::{constants::CollisionResistanceBytes, hashing::Digest, random::CryptoRngCore};

/// The default RNG type used in the library, set to `StdRng` (ChaCha12).
pub type BaseRng = rand::rngs::StdRng;

/// The generic type for all seeds across our protocols.
pub type Seed = [u8; CollisionResistanceBytes::USIZE];

/// A generic trait for secure pseudo-random generators (PRGs).
/// This trait is a combination of `RngCore`, `CryptoRng`, and `SeedableRng`.
pub trait SeedableRng:
    CryptoRngCore + rand::SeedableRng<Seed: AsRef<[u8]> + From<Digest>> + Clone + Send
{
    /// Creates a new random generator from the hash of a seed with a domain-separating tag.
    #[inline]
    fn from_tagged_seed(seed: Self::Seed, tag: impl AsRef<[u8]>) -> Self {
        let seed = crate::hashing::hash(&[seed.as_ref(), tag.as_ref()]);
        <Self as rand::SeedableRng>::from_seed(seed.into())
    }

    #[inline]
    /// Create a new PRNG using the given seed.
    fn from_seed(seed: Self::Seed) -> Self {
        <Self as rand::SeedableRng>::from_seed(seed)
    }

    #[inline]
    /// Create a new PRNG seeded from another `Rng`.
    fn from_rng<R: rand::RngCore>(mut rng: R) -> Result<Self, rand::Error> {
        <Self as rand::SeedableRng>::from_rng(&mut rng)
    }

    #[inline]
    /// Fills a buffer with random bytes.
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        rand::RngCore::fill_bytes(self, dest)
    }
}

impl<T> SeedableRng for T where
    T: CryptoRngCore + rand::SeedableRng<Seed = super::Seed> + Clone + Send
{
}

#[cfg(any(test, feature = "dev"))]
pub mod test {
    use std::{cell::UnsafeCell, rc::Rc};

    use rand::{CryptoRng, RngCore, SeedableRng};

    use crate::random::BaseRng;

    /// Returns a test RNG that can be used for testing purposes.
    ///
    /// If the environment variable `ASYNC_MPC_NON_DETERMINISTIC_TESTS` is set, it will use
    /// a random (non-deterministic) seed for the RNG. Otherwise, it will use the current
    /// thread's name as the seed (deterministic).
    pub fn test_rng() -> TestRng {
        let rng = TEST_RNG_KEY.with(|t| t.clone());
        TestRng { rng }
    }

    thread_local!(
        static TEST_RNG_KEY: Rc<UnsafeCell<BaseRng>> = {
        let rng = match std::env::var("ASYNC_MPC_NON_DETERMINISTIC_TESTS") {
            Ok(_) => BaseRng::from_rng(rand::thread_rng()).unwrap_or_else(|err|
                    panic!("could not initialize test_rng: {err}")),
            Err(_) => {
                let thread = std::thread::current();
                let seed = crate::hashing::hash(&[thread.name().unwrap_or("async_mpc_test").as_bytes()]);
                BaseRng::from_seed(seed.into())
            }
        };
        Rc::new(UnsafeCell::new(rng))
    });

    /// A reference to the thread-local random number generator used for testing.
    ///
    /// Based on [ThreadRng] from the rand crate, explicitly not Send or Sync.
    ///
    /// [`ThreadRng`]: rand::ThreadRng
    #[derive(Clone, Debug)]
    pub struct TestRng {
        // Rc is explicitly !Send and !Sync
        rng: Rc<UnsafeCell<BaseRng>>,
    }

    impl Default for TestRng {
        fn default() -> TestRng {
            test_rng()
        }
    }

    impl RngCore for TestRng {
        #[inline(always)]
        fn next_u32(&mut self) -> u32 {
            // SAFETY: We must make sure to stop using `rng` before anyone else
            // creates another mutable reference
            let rng = unsafe { &mut *self.rng.get() };
            rng.next_u32()
        }

        #[inline(always)]
        fn next_u64(&mut self) -> u64 {
            // SAFETY: We must make sure to stop using `rng` before anyone else
            // creates another mutable reference
            let rng = unsafe { &mut *self.rng.get() };
            rng.next_u64()
        }

        fn fill_bytes(&mut self, dest: &mut [u8]) {
            // SAFETY: We must make sure to stop using `rng` before anyone else
            // creates another mutable reference
            let rng = unsafe { &mut *self.rng.get() };
            rng.fill_bytes(dest)
        }

        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
            // SAFETY: We must make sure to stop using `rng` before anyone else
            // creates another mutable reference
            let rng = unsafe { &mut *self.rng.get() };
            rng.try_fill_bytes(dest)
        }
    }

    impl CryptoRng for TestRng {}
}