vaea-ntt 0.1.0

// Copyright (C) 2024-2026 Vaea SAS
// SPDX-License-Identifier: AGPL-3.0-or-later
//
// This file is part of VaeaNTT.
//
// VaeaNTT is free software: you can redistribute it and/or modify it under
// the terms of the GNU Affero General Public License as published by the
// Free Software Foundation, either version 3 of the License, or (at your
// option) any later version.
//
// VaeaNTT is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public
// License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with VaeaNTT. If not, see <https://www.gnu.org/licenses/>.


// =============================================================================
// VaeaNTT — Security Exploit Test Suite (Day-0 Hunt)
// =============================================================================
// Attempts to break the NTT through:
// 1. Out-of-range inputs (values > q)
// 2. Malicious prime values
// 3. Timing oracle detection
// 4. Integer overflow triggers
// 5. Boundary condition exploits
// 6. Panic/DoS vectors

use vaea_ntt::ntt32::{Ntt32Context, generate_primes_28, is_prime_32};
use std::time::Instant;

fn main() {
    let mut pass = 0u32;
    let mut fail = 0u32;
    let mut vulns: Vec<String> = Vec::new();

    println!("╔══════════════════════════════════════════════════════════╗");
    println!("║  VaeaNTT — Security Exploit Suite (Day-0 Hunt)         ║");
    println!("╚══════════════════════════════════════════════════════════╝\n");

    // =========================================================================
    // Exploit 1: Out-of-range input values (data[i] >= q)
    // =========================================================================
    println!("── Exploit 1: Out-of-range inputs ────────────────────────");
    {
        let ctx = Ntt32Context::new(256, 12289);

        // Feed values > q — does it crash? Does roundtrip still work?
        // A secure implementation should handle this gracefully.
        let mut data = vec![12289u32; 256]; // exactly q
        let original = data.clone();
        ctx.forward(&mut data);
        ctx.inverse(&mut data);
        // After roundtrip, values should be normalized to [0, q)
        // 12289 mod 12289 = 0, so we expect all zeros
        let all_zero = data.iter().all(|&x| x == 0);
        if all_zero {
            println!("  ✓ Values == q → normalized to 0 after roundtrip");
            pass += 1;
        } else {
            // Not necessarily a vulnerability, but unexpected
            let in_range = data.iter().all(|&x| x < 12289);
            if in_range {
                println!("  ⚠ Values == q → roundtrip gives non-zero but in range");
                pass += 1;
            } else {
                println!("  ⚠ Values == q → OUTPUT OUT OF RANGE after roundtrip!");
                vulns.push("Out-of-range output when input == q".into());
                fail += 1;
            }
        }

        // Feed values >> q (e.g., 2^31 - 1)
        let mut data2 = vec![0x7FFF_FFFFu32; 256]; // max i32
        // This WILL break the NTT math, but should it crash?
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let mut d = data2.clone();
            ctx.forward(&mut d);
            d
        }));
        match result {
            Ok(output) => {
                println!("  ✓ Values = 2^31-1 → no crash (output may be garbage)");
                pass += 1;
            }
            Err(_) => {
                println!("  ✗ Values = 2^31-1 → PANIC!");
                vulns.push("Panic on large input values".into());
                fail += 1;
            }
        }

        // Feed value u32::MAX
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let mut d = vec![u32::MAX; 256];
            ctx.forward(&mut d);
            d
        }));
        match result {
            Ok(_) => {
                println!("  ✓ Values = u32::MAX → no crash");
                pass += 1;
            }
            Err(_) => {
                println!("  ✗ Values = u32::MAX → PANIC!");
                vulns.push("Panic on u32::MAX input".into());
                fail += 1;
            }
        }
    }
    println!();

    // =========================================================================
    // Exploit 2: Wrong-size data
    // =========================================================================
    println!("── Exploit 2: Wrong-size data (DoS via panic) ────────────");
    {
        let ctx = Ntt32Context::new(256, 12289);

        // Empty slice
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let mut d: Vec<u32> = vec![];
            ctx.forward(&mut d);
        }));
        if result.is_err() {
            println!("  ✓ Empty slice → panics (expected, assert_eq catches it)");
            pass += 1;
        } else {
            println!("  ✗ Empty slice → no panic (should have panicked!)");
            fail += 1;
        }

        // Wrong size (255 instead of 256)
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let mut d = vec![0u32; 255];
            ctx.forward(&mut d);
        }));
        if result.is_err() {
            println!("  ✓ Size 255 → panics (expected)");
            pass += 1;
        } else {
            println!("  ✗ Size 255 → no panic (BUFFER OVERFLOW RISK!)");
            vulns.push("No bounds check on wrong-size input".into());
            fail += 1;
        }

        // Too large (257)
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let mut d = vec![0u32; 257];
            ctx.forward(&mut d);
        }));
        if result.is_err() {
            println!("  ✓ Size 257 → panics (expected)");
            pass += 1;
        } else {
            println!("  ✗ Size 257 → no panic!");
            fail += 1;
        }
    }
    println!();

    // =========================================================================
    // Exploit 3: Timing side-channel probe
    // =========================================================================
    println!("── Exploit 3: Timing side-channel probe ──────────────────");
    {
        let ctx = Ntt32Context::new(256, 12289);
        let iterations = 10000;

        // Pre-allocate all buffers to avoid clone overhead in measurements
        let zeros: Vec<Vec<u32>> = (0..iterations).map(|_| vec![0u32; 256]).collect();
        let maxes: Vec<Vec<u32>> = (0..iterations).map(|_| vec![12288u32; 256]).collect();
        let mixed: Vec<Vec<u32>> = (0..iterations).map(|_| {
            (0..256).map(|i| ((i * 7 + 13) % 12289) as u32).collect()
        }).collect();

        // Warmup (5000 iterations to stabilize caches)
        for i in 0..5000 {
            let mut d = vec![0u32; 256];
            ctx.forward(&mut d);
        }

        let mut bufs_z = zeros;
        let t0 = Instant::now();
        for d in bufs_z.iter_mut() {
            ctx.forward(d);
        }
        let time_zeros = t0.elapsed().as_nanos() as f64 / iterations as f64;

        let mut bufs_m = maxes;
        let t0 = Instant::now();
        for d in bufs_m.iter_mut() {
            ctx.forward(d);
        }
        let time_max = t0.elapsed().as_nanos() as f64 / iterations as f64;

        let mut bufs_x = mixed;
        let t0 = Instant::now();
        for d in bufs_x.iter_mut() {
            ctx.forward(d);
        }
        let time_mixed = t0.elapsed().as_nanos() as f64 / iterations as f64;

        let max_diff_pct = ((time_max - time_zeros).abs() / time_zeros * 100.0).max(
            (time_mixed - time_zeros).abs() / time_zeros * 100.0
        );

        println!("  Forward NTT q=12289 (pre-allocated, post-warmup):");
        println!("    Timing (zeros):  {time_zeros:.1} ns");
        println!("    Timing (max):    {time_max:.1} ns");
        println!("    Timing (mixed):  {time_mixed:.1} ns");
        println!("    Max deviation:   {max_diff_pct:.2}%");

        if max_diff_pct < 10.0 {
            println!("  ✓ Constant-time within 10%");
            pass += 1;
        } else if max_diff_pct < 20.0 {
            println!("  ⚠ Timing varies {max_diff_pct:.1}% — borderline (cache effects likely)");
            pass += 1;
        } else {
            println!("  ✗ TIMING VARIES {max_diff_pct:.1}% — potential side-channel!");
            vulns.push(format!("Timing side-channel: {max_diff_pct:.1}% variation"));
            fail += 1;
        }

        // ML-DSA timing test
        let ctx_dsa = Ntt32Context::new(256, 8380417);

        let mut bufs_z2: Vec<Vec<u32>> = (0..iterations).map(|_| vec![0u32; 256]).collect();
        let t0 = Instant::now();
        for d in bufs_z2.iter_mut() {
            ctx_dsa.forward(d);
        }
        let time_z = t0.elapsed().as_nanos() as f64 / iterations as f64;

        let mut bufs_m2: Vec<Vec<u32>> = (0..iterations).map(|_| vec![8380416u32; 256]).collect();
        let t0 = Instant::now();
        for d in bufs_m2.iter_mut() {
            ctx_dsa.forward(d);
        }
        let time_m = t0.elapsed().as_nanos() as f64 / iterations as f64;

        let diff = (time_m - time_z).abs() / time_z * 100.0;
        println!("  ML-DSA: zeros={time_z:.1}ns, max={time_m:.1}ns, diff={diff:.2}%");
        if diff < 15.0 {
            println!("  ✓ ML-DSA constant-time within 15%");
            pass += 1;
        } else {
            println!("  ✗ ML-DSA timing varies {diff:.1}%!");
            vulns.push(format!("ML-DSA timing: {diff:.1}% variation"));
            fail += 1;
        }
    }
    println!();

    // =========================================================================
    // Exploit 4: Barrett reduction edge cases
    // =========================================================================
    println!("── Exploit 4: Barrett reduction edge cases ───────────────");
    {
        // Test with prime where Barrett constant is exact
        // bc = floor(2^32 / q). For q = 2^k, bc = 2^(32-k) exactly.
        // But q must be prime, so test near powers of 2.

        // q near 2^14: 16381 is prime (2^14 - 3)
        // Check: 16381 ≡ 1 mod 512? 16381-1=16380. 16380/512 = 31.99... No.
        // Try q = 12289 (known good)
        let ctx = Ntt32Context::new(256, 12289);

        // Forward with values that stress Barrett:
        // Values just below 2*q should trigger the vcgeq branch in Barrett
        let mut data: Vec<u32> = (0..256).map(|i| {
            if i % 2 == 0 { 12288 } else { 0 } // alternating max/zero
        }).collect();
        let original = data.clone();
        ctx.forward(&mut data);

        // Check all reduced
        let reduced = data.iter().all(|&x| x < 12289);
        if reduced {
            println!("  ✓ Barrett handles alternating max/zero correctly");
            pass += 1;
        } else {
            let bad: Vec<_> = data.iter().enumerate().filter(|(_, &x)| x >= 12289).collect();
            println!("  ✗ Barrett FAILED: {} values out of range", bad.len());
            vulns.push("Barrett fails on alternating max/zero".into());
            fail += 1;
        }

        ctx.inverse(&mut data);
        if data == original {
            pass += 1;
        } else {
            fail += 1;
        }
    }
    println!();

    // =========================================================================
    // Exploit 5: Integer wrap-around in butterfly
    // =========================================================================
    println!("── Exploit 5: Integer wrap-around in butterfly ───────────");
    {
        // The deferred-reduction butterfly does u + wv and u + 2q - wv
        // where wv ∈ [0, 2q). If u is large (after many stages), u + wv might wrap u32.
        // For q=12289: after 8 stages, max = 17*12289 = 208913. u + wv ≤ 208913 + 2*12289 = 233491.
        // 233491 << 2^32, so no wrap.
        // For q=268369921 (~28bit): after 3 stages, max = 7*q ≈ 1.88B. u + wv ≤ 1.88B + 2*268M = 2.42B.
        // 2.42B < 2^32 = 4.29B. OK.
        // But u + 2q: max u + 2q = 1.88B + 537M = 2.42B < 4.29B. OK.

        // Worst case: find q where the values get closest to 2^32
        let primes = generate_primes_28(256, 1);
        let q = primes[0]; // largest 28-bit prime for N=256
        let ctx = Ntt32Context::new(256, q);

        let mut data = vec![q - 1; 256]; // all max
        let original = data.clone();
        ctx.forward(&mut data);

        let reduced = data.iter().all(|&x| x < q);
        if reduced {
            println!("  ✓ Largest 28-bit prime (q={q}): max values survive forward");
            pass += 1;
        } else {
            println!("  ✗ Largest prime q={q}: values out of range after forward!");
            vulns.push(format!("Integer overflow with q={q}"));
            fail += 1;
        }

        ctx.inverse(&mut data);
        if data == original {
            println!("  ✓ Roundtrip correct with largest prime");
            pass += 1;
        } else {
            println!("  ✗ Roundtrip FAILED with largest prime!");
            fail += 1;
        }
    }
    println!();

    // =========================================================================
    // Exploit 6: Malicious context construction
    // =========================================================================
    println!("── Exploit 6: Malicious context construction ─────────────");
    {
        // q = 4 (not prime) — should be rejected
        let result = Ntt32Context::try_new(256, 4);
        if result.is_err() {
            println!("  ✓ q=4 (not prime) → rejected");
            pass += 1;
        } else {
            println!("  ✗ q=4 accepted! Non-prime modulus!");
            vulns.push("Non-prime modulus accepted".into());
            fail += 1;
        }

        // q = 1 — should be rejected
        let result = Ntt32Context::try_new(256, 1);
        if result.is_err() {
            println!("  ✓ q=1 → rejected");
            pass += 1;
        } else {
            println!("  ✗ q=1 accepted!");
            vulns.push("q=1 accepted".into());
            fail += 1;
        }

        // q = 0 — should be rejected
        let result = Ntt32Context::try_new(256, 0);
        if result.is_err() {
            println!("  ✓ q=0 → rejected");
            pass += 1;
        } else {
            println!("  ✗ q=0 accepted!");
            vulns.push("q=0 accepted".into());
            fail += 1;
        }

        // q = 2^28 exactly — should be rejected (too large)
        let result = Ntt32Context::try_new(256, 1 << 28);
        if result.is_err() {
            println!("  ✓ q=2^28 → rejected");
            pass += 1;
        } else {
            println!("  ✗ q=2^28 accepted!");
            fail += 1;
        }

        // N = 0
        let result = Ntt32Context::try_new(0, 12289);
        if result.is_err() {
            println!("  ✓ N=0 → rejected");
            pass += 1;
        } else {
            println!("  ✗ N=0 accepted!");
            vulns.push("N=0 accepted".into());
            fail += 1;
        }

        // N = 3 (not power of 2)
        let result = Ntt32Context::try_new(3, 12289);
        if result.is_err() {
            println!("  ✓ N=3 → rejected");
            pass += 1;
        } else {
            println!("  ✗ N=3 accepted!");
            fail += 1;
        }

        // N = 1
        let result = Ntt32Context::try_new(1, 12289);
        if result.is_err() {
            println!("  ✓ N=1 → rejected");
            pass += 1;
        } else {
            println!("  ✗ N=1 accepted!");
            fail += 1;
        }

        // q prime but NOT NTT-friendly for N
        // q=13 is prime, but 13-1=12, 2*256=512, 12 % 512 != 0
        let result = Ntt32Context::try_new(256, 13);
        if result.is_err() {
            println!("  ✓ q=13 (not NTT-friendly for N=256) → rejected");
            pass += 1;
        } else {
            println!("  ✗ q=13 accepted for N=256!");
            fail += 1;
        }
    }
    println!();

    // =========================================================================
    // Exploit 7: Double-free / use-after-move safety
    // =========================================================================
    println!("── Exploit 7: Memory safety (clone + reuse) ──────────────");
    {
        let ctx = Ntt32Context::new(256, 12289);
        let ctx2 = ctx.clone(); // Clone should work

        let mut d1 = vec![1u32; 256];
        let mut d2 = vec![1u32; 256];

        ctx.forward(&mut d1);
        ctx2.forward(&mut d2);

        if d1 == d2 {
            println!("  ✓ Clone produces identical results");
            pass += 1;
        } else {
            println!("  ✗ Clone gives different results!");
            vulns.push("Clone inconsistency".into());
            fail += 1;
        }

        // Reuse context many times
        for _ in 0..1000 {
            let mut d = vec![42u32; 256];
            ctx.forward(&mut d);
            ctx.inverse(&mut d);
            assert!(d.iter().all(|&x| x == 42));
        }
        println!("  ✓ Context reuse (1000 iterations) stable");
        pass += 1;
    }
    println!();

    // =========================================================================
    // Exploit 8: Concurrent safety (if applicable)
    // =========================================================================
    println!("── Exploit 8: Thread safety ──────────────────────────────");
    {
        use std::sync::Arc;
        use std::thread;

        let ctx = Arc::new(Ntt32Context::new(256, 12289));
        let mut handles = vec![];

        for tid in 0..8 {
            let ctx = ctx.clone();
            handles.push(thread::spawn(move || {
                let mut data: Vec<u32> = (0..256).map(|i| ((i + tid * 100) as u32) % 12289).collect();
                let original = data.clone();
                for _ in 0..100 {
                    ctx.forward(&mut data);
                    ctx.inverse(&mut data);
                    assert_eq!(data, original, "Thread {tid} roundtrip failed!");
                }
                true
            }));
        }

        let mut all_ok = true;
        for h in handles {
            if !h.join().unwrap() {
                all_ok = false;
            }
        }

        if all_ok {
            println!("  ✓ 8 threads × 100 roundtrips = stable");
            pass += 1;
        } else {
            println!("  ✗ Thread safety violation!");
            vulns.push("Thread safety issue".into());
            fail += 1;
        }
    }
    println!();

    // =========================================================================
    // Summary
    // =========================================================================
    let total = pass + fail;
    println!("╔══════════════════════════════════════════════════════════╗");
    println!("║  TOTAL: {pass:>4} pass | {fail:>3} fail | {total:>4} total              ║");
    if vulns.is_empty() {
        println!("║  🛡️  NO VULNERABILITIES FOUND                          ║");
    } else {
        println!("║  ⚠️  {} VULNERABILITIES FOUND:                          ║", vulns.len());
        for v in &vulns {
            println!("║  • {v:<53}║");
        }
    }
    println!("╚══════════════════════════════════════════════════════════╝");

    if fail > 0 {
        std::process::exit(1);
    }
}