caprand 0.3.0

#[cfg(not(feature = "defmt"))]
#[allow(unused_imports)]
use log::{debug, error, info, trace, warn};

#[cfg(feature = "defmt")]
#[allow(unused_imports)]
use defmt::{debug, error, info, panic, trace, warn};

use core::ops::DerefMut;
use core::{cell::RefCell, convert::Infallible};

use critical_section::Mutex;
use rand_chacha::ChaCha20Rng;
use sha2::{Digest, Sha256};

use embassy_rp::{gpio::Pin, Peri};

use rand::Rng;
use rand_chacha::rand_core::SeedableRng;

static RNG: Mutex<RefCell<Option<CapRng>>> = Mutex::new(RefCell::new(None));

/// A random byte generator.
///
/// `buf` will be filled with random bytes.
/// [`setup()`](setup) must be called prior to using this function.
///
/// This may be used by an application's getrandom custom backend,
/// see getrandom documentation.
pub fn getrandom(buf: &mut [u8]) -> Result<(), ()> {
    critical_section::with(|cs| {
        let mut rng = RNG.borrow_ref_mut(cs);
        let rng = rng.deref_mut();
        if let Some(rng) = rng {
            rng.0.fill_bytes(buf);
            Ok(())
        } else {
            error!("setup() not called");
            Err(())
        }
    })
}

/// A random byte generator suitable for getrandom custom backends.
///
/// `dest` will be filled with `len` random bytes.
/// Call this from a `__getrandom_v03_custom()`, see getrandom crate
/// documentation.
///
/// # Safety
///
/// `dest` and `len` must be a valid destination to write. `dest`
/// does not need to contain initialised memory.
pub unsafe fn getrandom_raw(dest: *mut u8, len: usize) -> Result<(), ()> {
    let buf = unsafe {
        core::ptr::write_bytes(dest, 0, len);
        core::slice::from_raw_parts_mut(dest, len)
    };
    getrandom(buf)
}

/// Seed the random generator from a capacitor noise source.
///
/// Call this at early startup.
///
/// # Arguments
///
/// * pin - The GPIO pin with a capacitor attached. This will be driven low and pulled high,
/// with timing used as a random source. The `Pin` may be used for other purposes once
/// `setup()` completes.
///
/// # Examples
///
/// ```
/// let mut p = embassy_rp::init(Default::default());
///
/// caprand::setup(&mut p.PIN_10).unwrap();
///
/// #[unsafe(no_mangle)]
/// unsafe extern "Rust" fn __getrandom_v03_custom(
///     dest: *mut u8,
///     len: usize,
/// ) -> Result<(), getrandom::Error> {
///     caprand::getrandom_raw(dest, len).map_err(|_| getrandom::Error::new_custom(123))
/// }
///
/// let mut mystery = [0u8; 10];
/// getrandom::getrandom(&mut mystery).unwrap();
/// ```
/// `getrandom` custom backend requires building with `--cfg getrandom_backend="custom",
/// see [`getrandom`] documentation.
pub fn setup(pin: Peri<impl Pin>) -> Result<(), ()> {
    let r = CapRng::new(pin)?;

    critical_section::with(|cs| {
        let mut rng = RNG.borrow_ref_mut(cs);
        let _ = rng.insert(r);
    });
    Ok(())
}

// TODO: this is another impl of chacha20, can it use chacha20 crate instead? Is the size much?
// TODO: have some kind of fast erasure RNG instead?
/// A cryptographic PRNG seeded by the capacitor noise source.
pub struct CapRng(ChaCha20Rng);

impl CapRng {
    /// The number of noise samples to use for seeding.
    ///
    /// We need to produce a 256 bit output seed.
    pub const SEED_SAMPLES: usize = 256 * 100;

    const MAX_FAILURES: usize = 3;

    pub fn new(pin: Peri<impl Pin>) -> Result<Self, ()> {
        let low_cycles = 1;
        let mut noise = crate::cap::RawNoise::new(pin, low_cycles);

        let mut valid_samples = 0;
        let mut h = Sha256::new();

        let mut health = crate::health::TotalHealth::new();
        let mut failures = 0;

        while valid_samples < Self::SEED_SAMPLES {
            let (v, valid) = noise
                .next()
                // OK unwrap, iterator doesn't end
                .unwrap();
            if valid {
                valid_samples += 1;
                if health.test(v).is_err() {
                    valid_samples = 0;
                    failures += 1;
                    if failures > Self::MAX_FAILURES {
                        error!(
                            "Health tests failed after {} retries",
                            Self::MAX_FAILURES
                        );
                        return Err(());
                    }
                }
            }

            // even "invalid" samples are included in the hash
            h.update([v]);
        }

        let seed: [u8; 32] = h.finalize().into();
        Ok(Self(ChaCha20Rng::from_seed(seed)))
    }
}

impl rand::TryCryptoRng for CapRng {}

impl rand::TryRng for CapRng {
    type Error = Infallible;
    fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
        self.0.try_next_u32()
    }

    fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
        self.0.try_next_u64()
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {
        self.0.try_fill_bytes(dest)
    }
}