// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Hardware RNG for x86_64 TEEs.
//!
//! Policy for randomness in Rialo TEEs:
//! - Call RDSEED (full-entropy seed material), then RDRAND (CPU DRBG).
//! - No OS fallback in production: inside a TEE we strictly require HW RNG.
//! - Retries: Both instructions can temporarily return "not ready" (CF=0).
//!   Vendors require software to retry briefly. We cap retries to bound tail
//!   latency; hitting the cap is a health signal.
//!
//! Test mode:
//! - On x86_64 hosts: try hardware RNG first, fallback to `getrandom` if hardware fails.
//! - On non-x86_64 hosts: use `getrandom` so CI can run on any architecture.
//! -----------------------------------------------------------------------------
//! TODO (Tracking): Re-evaluate RNG policy for TEE binaries.
//! We may switch to relying solely on the kernel CSPRNG (`getrandom(2)`)
//! (e.g., build with `--cfg getrandom_backend="linux_getrandom"`) and possibly
//! drop this module. Investigate version skew, TLS/deps behavior, and audit needs.
//! Ticket: <https://linear.app/subzero-labs/issue/SUB-1176/tee-builds-standardize-rng-policy-for-dependencies-use-kernel>
//! -----------------------------------------------------------------------------

/// Errors that can occur during hardware RNG operations.
#[derive(Debug, thiserror::Error)]
pub enum HwRngError {
    /// Neither RDSEED nor RDRAND instructions are available on this CPU.
    #[error("Hardware RNG not supported: neither RDSEED nor RDRAND available")]
    Unsupported,

    /// Hardware RNG instruction didn't succeed within retry bounds.
    #[error("Hardware RNG failed: instruction didn't succeed within retry bounds")]
    Failed,

    /// OS RNG error occurred (only used in test configurations).
    #[error("OS RNG error in test mode")]
    OsError,
}

// ===== x86_64 helpers ========================================================

#[cfg(target_arch = "x86_64")]
#[allow(unsafe_code)]
#[allow(dead_code)]
#[inline]
fn rdseed64_once() -> Option<u64> {
    let mut out = 0u64;
    // SAFETY: guarded by is_x86_feature_detected! below
    let ok = unsafe { core::arch::x86_64::_rdseed64_step(&mut out) };
    (ok == 1).then_some(out)
}

#[cfg(target_arch = "x86_64")]
#[allow(unsafe_code)]
#[allow(dead_code)]
#[inline]
fn rdrand64_once() -> Option<u64> {
    let mut out = 0u64;
    // SAFETY: guarded by is_x86_feature_detected! below
    let ok = unsafe { core::arch::x86_64::_rdrand64_step(&mut out) };
    (ok == 1).then_some(out)
}

#[cfg(target_arch = "x86_64")]
#[allow(unsafe_code)]
#[allow(dead_code)]
fn try_fill_rdseed(dst: &mut [u8]) -> Result<(), HwRngError> {
    if !std::is_x86_feature_detected!("rdseed") {
        return Err(HwRngError::Unsupported);
    }
    const MAX: usize = 16; // conservative; typically succeeds on first try
    let mut i = 0;
    while i < dst.len() {
        let mut got = None;
        for attempt in 0..MAX {
            if let Some(v) = rdseed64_once() {
                got = Some(v);
                break;
            }
            if attempt >= 3 {
                // Brief backoff after a few misses to reduce contention
                core::arch::x86_64::_mm_pause();
            } else {
                core::hint::spin_loop();
            }
        }
        let v = got.ok_or(HwRngError::Failed)?;
        let take = core::cmp::min(8, dst.len() - i);
        dst[i..i + take].copy_from_slice(&v.to_ne_bytes()[..take]);
        i += take;
    }
    Ok(())
}

#[cfg(target_arch = "x86_64")]
#[allow(unsafe_code)]
#[allow(dead_code)]
fn try_fill_rdrand(dst: &mut [u8]) -> Result<(), HwRngError> {
    if !std::is_x86_feature_detected!("rdrand") {
        return Err(HwRngError::Unsupported);
    }
    const MAX: usize = 8; // common guidance ~8–10
    let mut i = 0;
    while i < dst.len() {
        let mut got = None;
        for attempt in 0..MAX {
            if let Some(v) = rdrand64_once() {
                got = Some(v);
                break;
            }
            if attempt >= 3 {
                core::arch::x86_64::_mm_pause();
            } else {
                core::hint::spin_loop();
            }
        }
        let v = got.ok_or(HwRngError::Failed)?;
        let take = core::cmp::min(8, dst.len() - i);
        dst[i..i + take].copy_from_slice(&v.to_ne_bytes()[..take]);
        i += take;
    }
    Ok(())
}

// Non-x86_64 stubs to keep the code simple elsewhere.
#[cfg(not(target_arch = "x86_64"))]
#[allow(dead_code)]
fn try_fill_rdseed(_dst: &mut [u8]) -> Result<(), HwRngError> {
    Err(HwRngError::Unsupported)
}
#[cfg(not(target_arch = "x86_64"))]
#[allow(dead_code)]
fn try_fill_rdrand(_dst: &mut [u8]) -> Result<(), HwRngError> {
    Err(HwRngError::Unsupported)
}

// ===== Public API ============================================================

/// Fill `dst` with random bytes from hardware RNG (RDSEED→RDRAND).
/// - Production: **no OS fallback**. Returns `Unsupported` or `Failed` on error.
/// - Tests: try hardware RNG first on x86_64, fallback to getrandom for portability.
#[cfg(any(test, feature = "testing"))]
pub fn try_fill_bytes(dst: &mut [u8]) -> Result<(), HwRngError> {
    // Try hardware first on x86_64 hosts
    #[cfg(target_arch = "x86_64")]
    {
        if try_fill_rdseed(dst).is_ok() {
            return Ok(());
        }
        if try_fill_rdrand(dst).is_ok() {
            return Ok(());
        }
    }

    // Fallback to getrandom for non-x86_64 or when HW RNG fails
    getrandom::getrandom(dst).map_err(|_| HwRngError::OsError)
}
#[cfg(not(any(test, feature = "testing")))]
pub fn try_fill_bytes(dst: &mut [u8]) -> Result<(), HwRngError> {
    if try_fill_rdseed(dst).is_ok() {
        return Ok(());
    }
    if try_fill_rdrand(dst).is_ok() {
        return Ok(());
    }
    Err(HwRngError::Unsupported)
}

/// Convenience: return `N` hardware-random bytes.
pub fn bytes<const N: usize>() -> Result<[u8; N], HwRngError> {
    let mut out = [0u8; N];
    try_fill_bytes(&mut out)?;
    Ok(out)
}

/// Panicking helper for call-sites that *must* have HW RNG.
pub fn fill_bytes(dst: &mut [u8]) {
    try_fill_bytes(dst).expect("HW RNG required in TEE");
}

// ===== RNG adapter for APIs that expect rand_core::RngCore ===================

pub struct HwRng;
impl Default for HwRng {
    fn default() -> Self {
        Self
    }
}

impl rand_core::RngCore for HwRng {
    fn next_u32(&mut self) -> u32 {
        let a = bytes::<4>().expect("HW RNG");
        u32::from_ne_bytes(a)
    }
    fn next_u64(&mut self) -> u64 {
        let a = bytes::<8>().expect("HW RNG");
        u64::from_ne_bytes(a)
    }
    fn fill_bytes(&mut self, dst: &mut [u8]) {
        try_fill_bytes(dst).expect("HW RNG")
    }
}
impl rand_core::CryptoRng for HwRng {}

#[cfg(test)]
mod tests {
    #[cfg(not(target_arch = "x86_64"))]
    use assert_matches::assert_matches;
    use rand_core::RngCore;

    use super::*;

    #[test]
    fn test_try_fill_bytes_basic() {
        let mut buffer = [0u8; 32];
        let result = try_fill_bytes(&mut buffer);
        assert!(result.is_ok(), "try_fill_bytes should succeed");

        // Basic check that we got some data (not all zeros)
        assert_ne!(buffer, [0u8; 32], "Buffer should not be all zeros");
    }

    #[test]
    fn test_try_fill_bytes_different_sizes() {
        // Test various buffer sizes
        let sizes = [1, 4, 8, 16, 32, 64, 100];

        for size in sizes {
            let mut buffer = vec![0u8; size];
            let result = try_fill_bytes(&mut buffer);
            assert!(
                result.is_ok(),
                "try_fill_bytes should succeed for size {}",
                size
            );
        }
    }

    #[test]
    fn test_try_fill_bytes_empty_buffer() {
        let mut buffer = [];
        let result = try_fill_bytes(&mut buffer);
        assert!(
            result.is_ok(),
            "try_fill_bytes should succeed with empty buffer"
        );
    }

    #[test]
    fn test_bytes_convenience_function() {
        let result: Result<[u8; 16], _> = bytes();
        assert!(result.is_ok(), "bytes() should succeed");

        let data = result.unwrap();
        assert_ne!(data, [0u8; 16], "bytes() should not return all zeros");
    }

    #[test]
    fn test_bytes_different_sizes() {
        // Test different const generic sizes
        assert!(bytes::<1>().is_ok());
        assert!(bytes::<4>().is_ok());
        assert!(bytes::<8>().is_ok());
        assert!(bytes::<16>().is_ok());
        assert!(bytes::<32>().is_ok());
        assert!(bytes::<64>().is_ok());
    }

    #[test]
    fn test_repeated_calls_produce_different_output() {
        let data1: [u8; 32] = bytes().expect("First call should succeed");
        let data2: [u8; 32] = bytes().expect("Second call should succeed");

        // It's extremely unlikely that two random 32-byte arrays are identical
        assert_ne!(
            data1, data2,
            "Repeated calls should produce different output"
        );
    }

    #[test]
    fn test_fill_bytes_panicking_wrapper() {
        let mut buffer = [0u8; 16];

        // This should not panic in test mode since we have getrandom fallback
        fill_bytes(&mut buffer);

        // Check that buffer was filled
        assert_ne!(buffer, [0u8; 16], "fill_bytes should fill the buffer");
    }

    #[test]
    fn test_hw_rng_next_u32() {
        let mut rng = HwRng;
        let val1 = rng.next_u32();
        let val2 = rng.next_u32();

        // Very unlikely to get the same value twice
        assert_ne!(val1, val2, "next_u32() should produce different values");
    }

    #[test]
    fn test_hw_rng_next_u64() {
        let mut rng = HwRng;
        let val1 = rng.next_u64();
        let val2 = rng.next_u64();

        // Very unlikely to get the same value twice
        assert_ne!(val1, val2, "next_u64() should produce different values");
    }

    #[test]
    fn test_hw_rng_fill_bytes() {
        let mut rng = HwRng;
        let mut buffer = [0u8; 64];

        rng.fill_bytes(&mut buffer);

        assert_ne!(
            buffer, [0u8; 64],
            "HwRng::fill_bytes should fill the buffer"
        );
    }

    #[test]
    fn test_hw_rng_implements_crypto_rng() {
        fn accepts_crypto_rng<R: rand_core::CryptoRng>(_rng: R) {}

        let rng = HwRng;
        accepts_crypto_rng(rng); // Should compile
    }

    #[test]
    fn test_randomness_quality_basic() {
        let data: [u8; 256] = bytes().expect("Should get random bytes");

        // Check that we don't have obvious patterns
        // Count unique bytes
        let mut seen = std::collections::HashSet::new();
        for &byte in &data {
            seen.insert(byte);
        }

        // We should see a reasonable variety of byte values
        // With 256 random bytes, we expect to see most possible byte values
        assert!(
            seen.len() > 100,
            "Should see reasonable variety in random bytes, got {} unique values",
            seen.len()
        );

        // Check we don't have too many repeated consecutive bytes
        let mut consecutive_same = 0;
        let mut max_consecutive = 0;
        for i in 1..data.len() {
            if data[i] == data[i - 1] {
                consecutive_same += 1;
                max_consecutive = max_consecutive.max(consecutive_same);
            } else {
                consecutive_same = 0;
            }
        }

        // With good randomness, we shouldn't see many consecutive identical bytes
        assert!(
            max_consecutive < 10,
            "Too many consecutive identical bytes: {}",
            max_consecutive
        );
    }

    // Architecture-specific tests
    #[cfg(target_arch = "x86_64")]
    #[test]
    fn test_x86_64_hardware_functions_available() {
        // These functions should be available in test mode on x86_64
        let mut buffer = [0u8; 8];

        // Note: These might fail if hardware doesn't support the instructions,
        // but the functions should at least be callable
        let _rdseed_result = try_fill_rdseed(&mut buffer);
        let _rdrand_result = try_fill_rdrand(&mut buffer);
    }

    #[cfg(not(target_arch = "x86_64"))]
    #[test]
    fn test_non_x86_64_hardware_functions_unsupported() {
        let mut buffer = [0u8; 8];

        // On non-x86_64, hardware functions should return Unsupported
        assert_matches!(try_fill_rdseed(&mut buffer), Err(HwRngError::Unsupported));
        assert_matches!(try_fill_rdrand(&mut buffer), Err(HwRngError::Unsupported));
    }

    #[test]
    fn test_large_buffer_handling() {
        // Test with a larger buffer to ensure chunking works correctly
        let mut buffer = vec![0u8; 1024];
        let result = try_fill_bytes(&mut buffer);
        assert!(result.is_ok(), "Should handle large buffers");

        // Check that the entire buffer was filled
        assert_ne!(buffer, vec![0u8; 1024], "Large buffer should be filled");

        // Check different regions of the buffer to ensure it's all filled
        let quarter = buffer.len() / 4;
        assert_ne!(
            &buffer[0..quarter],
            &vec![0u8; quarter][..],
            "First quarter should be filled"
        );
        assert_ne!(
            &buffer[quarter..2 * quarter],
            &vec![0u8; quarter][..],
            "Second quarter should be filled"
        );
        assert_ne!(
            &buffer[2 * quarter..3 * quarter],
            &vec![0u8; quarter][..],
            "Third quarter should be filled"
        );
        assert_ne!(
            &buffer[3 * quarter..],
            &vec![0u8; buffer.len() - 3 * quarter][..],
            "Last quarter should be filled"
        );
    }
}
