lez-fuzzing/fuzz_props/src/generators.rs

use arbitrary::{Arbitrary, Unstructured};
use common::{block::HashableBlockData, transaction::NSSATransaction};
use nssa::{AccountId, PrivateKey};
use proptest::prelude::*;
use testnet_initial_state::initial_pub_accounts_private_keys;

// ── Arbitrary (for libFuzzer targets) ────────────────────────────────────────

/// A best-effort attempt to create a structurally plausible `NSSATransaction`
/// from unstructured bytes. Falls back to raw borsh decoding.
pub fn arbitrary_transaction(u: &mut Unstructured<'_>) -> arbitrary::Result<NSSATransaction> {
    // Prefer structured generation; raw decode as fallback
    if bool::arbitrary(u)? {
        let raw = Vec::<u8>::arbitrary(u)?;
        borsh::from_slice::<NSSATransaction>(&raw)
            .map_err(|_| arbitrary::Error::IncorrectFormat)
    } else {
        // Generate a minimal empty public tx using known test keys
        let signing_key = PrivateKey::try_new([u8::arbitrary(u)?; 32])
            .map_err(|_| arbitrary::Error::IncorrectFormat)?;
        let program_id = nssa::program::Program::authenticated_transfer_program().id();
        let message = nssa::public_transaction::Message::try_new(
            program_id,
            vec![],
            vec![],
            u128::arbitrary(u)?,
        )
        .map_err(|_| arbitrary::Error::IncorrectFormat)?;
        let witness = nssa::public_transaction::WitnessSet::for_message(&message, &[&signing_key]);
        Ok(NSSATransaction::Public(
            nssa::PublicTransaction::new(message, witness),
        ))
    }
}

// ── proptest strategies ───────────────────────────────────────────────────────

prop_compose! {
    /// Strategy: a valid native-transfer public transaction between two known accounts.
    pub fn arb_native_transfer_tx(
        accounts: Vec<(AccountId, PrivateKey)>,
    )(
        from_idx in 0..accounts.len(),
        to_idx   in 0..accounts.len(),
        nonce    in 0u128..1_000u128,
        amount   in 0u128..10_000u128,
    ) -> NSSATransaction {
        let (from_id, from_key) = &accounts[from_idx];
        let (to_id, _)          = &accounts[to_idx];
        common::test_utils::create_transaction_native_token_transfer(
            *from_id, nonce, *to_id, amount, from_key,
        )
    }
}

/// Return the test accounts from `testnet_initial_state` as `(AccountId, PrivateKey)` pairs.
pub fn test_accounts() -> Vec<(AccountId, PrivateKey)> {
    initial_pub_accounts_private_keys()
        .into_iter()
        .map(|k| (k.account_id, k.pub_sign_key))
        .collect()
}

/// Strategy: raw bytes that are valid borsh encodings of `NSSATransaction`.
pub fn arb_borsh_transaction_bytes() -> impl Strategy<Value = Vec<u8>> {
    any::<Vec<u8>>().prop_map(|bytes| {
        // Either pass through raw bytes OR encode a known dummy transaction
        if borsh::from_slice::<NSSATransaction>(&bytes).is_ok() {
            bytes
        } else {
            borsh::to_vec(&common::test_utils::produce_dummy_empty_transaction()).unwrap()
        }
    })
}

/// Strategy: a `HashableBlockData` with 0–8 transactions.
pub fn arb_hashable_block_data() -> impl Strategy<Value = HashableBlockData> {
    let accounts = test_accounts();
    proptest::collection::vec(arb_native_transfer_tx(accounts), 0..8).prop_map(|txs| {
        HashableBlockData {
            block_id:        1,
            prev_block_hash: common::HashType([0; 32]),
            timestamp:       0,
            transactions:    txs,
        }
    })
}

// ── IS-3: Invalid account / state combinations ────────────────────────────────

prop_compose! {
    /// Strategy: a transfer from an account that does not exist in the genesis state,
    /// or a transfer whose amount exceeds the sender's balance (invalid state combo).
    /// These inputs are expected to be rejected; the invariant being tested is that
    /// the state is left unchanged on rejection (StateIsolationOnFailure).
    pub fn arb_invalid_account_state_tx()(
        // Use a random 32-byte seed as a "phantom" account id not in genesis
        phantom_id_bytes in proptest::array::uniform32(0u8..),
        amount in (u128::MAX / 2)..u128::MAX,   // overflow-inducing amount
        nonce  in 0u128..10u128,
    ) -> NSSATransaction {
        let phantom_id = nssa::AccountId::new(phantom_id_bytes);
        // Attempt to sign with a key that has no matching on-chain account
        let signing_key = nssa::PrivateKey::try_new(phantom_id_bytes)
            .expect("phantom signing key");
        let (valid_to_id, _) = test_accounts()
            .into_iter()
            .next()
            .expect("at least one account");
        common::test_utils::create_transaction_native_token_transfer(
            phantom_id, nonce, valid_to_id, amount, &signing_key,
        )
    }
}

// ── IS-4: Re-ordered / duplicated inputs ─────────────────────────────────────

/// Strategy: a sequence of transactions where some are exact duplicates (replay
/// attack candidates) and some are re-ordered permutations of a valid sequence.
/// Used in proptest-level tests and as a seed generator for the state-transition
/// fuzz target.
pub fn arb_duplicate_tx_sequence() -> impl Strategy<Value = Vec<NSSATransaction>> {
    let accounts = test_accounts();
    proptest::collection::vec(arb_native_transfer_tx(accounts), 1..5_usize).prop_flat_map(
        |txs| {
            // Build a sequence that: original | duplicates | reversed
            let duped: Vec<NSSATransaction> = txs
                .iter()
                .cloned()
                .chain(txs.iter().cloned()) // append exact duplicates
                .chain(txs.iter().rev().cloned()) // append reversed order
                .collect();
            Just(duped)
        },
    )
}

// ── IS-5: Pathological sequences intended to violate protocol rules ───────────

/// Strategy: sequences designed to probe boundary conditions and protocol rules:
/// - zero-value transfers (no-op drain),
/// - self-transfers (sender == recipient),
/// - max-nonce wrapping,
/// - alternating valid / invalid transactions to test partial-batch isolation.
pub fn arb_pathological_sequence() -> impl Strategy<Value = Vec<NSSATransaction>> {
    let accounts = test_accounts();
    let n = accounts.len();
    proptest::collection::vec(
        (0..n, 0..n, 0u128..5u128, any::<bool>()),
        1..8_usize,
    )
    .prop_map(move |params| {
        params
            .into_iter()
            .map(|(from_idx, to_idx, nonce, zero_amount)| {
                let (from_id, from_key) = &accounts[from_idx];
                let (to_id, _) = &accounts[to_idx];
                let amount = if zero_amount { 0u128 } else { u128::MAX }; // 0 or overflow
                common::test_utils::create_transaction_native_token_transfer(
                    *from_id, nonce, *to_id, amount, from_key,
                )
            })
            .collect()
    })
}