lez-fuzzing/fuzz_props/src/invariants.rs

use common::transaction::NSSATransaction;
use nssa::V03State;
use nssa_core::account::Nonce;

/// Snapshot of public account balances used for conservation checks.
#[derive(Clone, Debug)]
pub struct BalanceSnapshot(pub std::collections::HashMap<nssa::AccountId, u128>);

impl BalanceSnapshot {
    /// Capture current total balance over all known accounts.
    pub fn total(&self) -> u128 {
        self.0.values().copied().fold(0u128, u128::saturating_add)
    }
}

/// Snapshot of account nonces captured before a transaction is applied.
///
/// Mirrors [`BalanceSnapshot`]: one entry per account that should remain
/// stable on a failed transaction (`FailedTxNonceStability` invariant).
#[derive(Clone, Debug)]
pub struct NonceSnapshot(pub std::collections::HashMap<nssa::AccountId, Nonce>);

/// Shared context threaded through every per-transaction invariant check.
///
/// Build this **after** calling `execute_check_on_state` so that `state_after`
/// reflects the post-execution state and `execution_succeeded` matches the
/// actual outcome.
pub struct InvariantCtx<'a> {
    /// State snapshot captured **before** applying the transaction.
    pub state_before: &'a V03State,
    /// Live state **after** applying (or attempting) the transaction.
    pub state_after: &'a V03State,
    /// `true` when `execute_check_on_state` returned `Ok`, `false` on `Err`.
    pub execution_succeeded: bool,
    /// Per-account balances captured before the transaction.
    pub balances_before: BalanceSnapshot,
    /// Per-account nonces captured before the transaction.
    pub nonces_before: NonceSnapshot,
}

/// A named invariant violation with an actionable diagnostic message.
#[derive(Debug)]
pub struct InvariantViolation {
    pub invariant: &'static str,
    pub message: String,
}

/// A protocol rule that can be checked against a single transaction's context.
pub trait ProtocolInvariant {
    fn name(&self) -> &'static str;
    fn check(&self, ctx: &InvariantCtx<'_>) -> Option<InvariantViolation>;
}

// ── Concrete invariants ───────────────────────────────────────────────────────

/// Sum of all public account balances must never change when a transaction is rejected.
///
/// A balance mutation on failure means the protocol leaks state on error paths —
/// e.g., a debit that is not rolled back after a later validation failure.
pub struct StateIsolationOnFailure;

impl ProtocolInvariant for StateIsolationOnFailure {
    fn name(&self) -> &'static str {
        "StateIsolationOnFailure"
    }

    fn check(&self, ctx: &InvariantCtx<'_>) -> Option<InvariantViolation> {
        if !ctx.execution_succeeded {
            for (acc_id, &expected_balance) in &ctx.balances_before.0 {
                let actual_balance = ctx.state_after.get_account_by_id(*acc_id).balance;
                if actual_balance != expected_balance {
                    return Some(InvariantViolation {
                        invariant: self.name(),
                        message: format!(
                            "balance changed despite tx rejection: account {:?} had \
                             {expected_balance} before, {actual_balance} after",
                            acc_id,
                        ),
                    });
                }
            }
        }
        None
    }
}

/// Total balance of all known accounts must be conserved when a transaction succeeds.
///
/// Catches double-credit and token-inflation bugs: a transfer path that credits the
/// recipient without debiting the sender would inflate the sum of all known balances.
/// The check uses the same account set captured in [`BalanceSnapshot`] so that new
/// accounts silently created by execution are NOT included (see known limitation
/// in `fuzz_validate_execute_consistency`).
pub struct BalanceConservation;

impl ProtocolInvariant for BalanceConservation {
    fn name(&self) -> &'static str {
        "BalanceConservation"
    }

    fn check(&self, ctx: &InvariantCtx<'_>) -> Option<InvariantViolation> {
        if ctx.execution_succeeded {
            let total_before = ctx.balances_before.total();
            let total_after: u128 = ctx
                .balances_before
                .0
                .keys()
                .map(|&id| ctx.state_after.get_account_by_id(id).balance)
                .fold(0u128, u128::saturating_add);
            if total_before != total_after {
                return Some(InvariantViolation {
                    invariant: self.name(),
                    message: format!(
                        "total balance of known accounts changed after successful transaction: \
                         before={total_before}, after={total_after} \
                         (possible double-credit or token-inflation bug)",
                    ),
                });
            }
        }
        None
    }
}

/// When a transaction is rejected, every account's nonce must remain unchanged.
///
/// A nonce mutation on a failed transaction constitutes a griefing attack: an
/// adversary could force arbitrary transactions to fail and permanently burn the
/// victim's nonce, rendering their account unusable.
pub struct FailedTxNonceStability;

impl ProtocolInvariant for FailedTxNonceStability {
    fn name(&self) -> &'static str {
        "FailedTxNonceStability"
    }

    fn check(&self, ctx: &InvariantCtx<'_>) -> Option<InvariantViolation> {
        if !ctx.execution_succeeded {
            for (&acc_id, expected_nonce) in &ctx.nonces_before.0 {
                let actual_nonce = ctx.state_after.get_account_by_id(acc_id).nonce;
                if actual_nonce != *expected_nonce {
                    return Some(InvariantViolation {
                        invariant: self.name(),
                        message: format!(
                            "nonce changed despite tx rejection: account {:?} nonce was \
                             {:?} before, {:?} after \
                             (griefing attack — victim nonce permanently burned on failed tx)",
                            acc_id, expected_nonce, actual_nonce,
                        ),
                    });
                }
            }
        }
        None
    }
}

/// A successfully accepted transaction must be rejected when replayed.
///
/// # Note
///
/// This invariant **cannot** be exercised through [`InvariantCtx`] alone because
/// the replay check requires re-applying the `NSSATransaction` that was consumed
/// by `execute_check_on_state`.  The `ProtocolInvariant` impl here is a registry
/// placeholder only; it always returns `None`.
///
/// Use the standalone [`assert_replay_rejection`] function instead, which accepts
/// the `NSSATransaction` returned on success and performs the replay inline.
pub struct ReplayRejection;

impl ProtocolInvariant for ReplayRejection {
    fn name(&self) -> &'static str {
        "ReplayRejection"
    }

    fn check(&self, _ctx: &InvariantCtx<'_>) -> Option<InvariantViolation> {
        // ReplayRejection cannot be fully exercised through InvariantCtx alone.
        // Use `assert_replay_rejection(applied_tx, state, next_block_id, next_ts)` instead.
        None
    }
}

/// A successfully applied transaction must increment the nonce of every signer account
/// by exactly one.
///
/// # Note
///
/// This invariant **cannot** be exercised through [`InvariantCtx`] alone because
/// `InvariantCtx` does not carry a signer-ID list — that information is private to the
/// `nssa` crate and is consumed by `apply_state_diff` before it returns.  The
/// `ProtocolInvariant` impl here is a registry placeholder only; it always returns `None`.
///
/// Use the standalone [`assert_nonce_increment_correctness`] function instead, passing
/// the signer IDs derived from the transaction's witness set, the [`NonceSnapshot`]
/// captured before execution, and the post-execution state.
pub struct NonceIncrementCorrectness;

impl ProtocolInvariant for NonceIncrementCorrectness {
    fn name(&self) -> &'static str {
        "NonceIncrementCorrectness"
    }

    fn check(&self, _ctx: &InvariantCtx<'_>) -> Option<InvariantViolation> {
        // NonceIncrementCorrectness requires explicit signer_ids not available in InvariantCtx.
        // Use `assert_nonce_increment_correctness(signer_ids, nonces_before, state_after)` instead.
        None
    }
}

// ── Standalone helpers ────────────────────────────────────────────────────────

/// Assert that a successfully-applied transaction is **rejected** when replayed.
///
/// Call this immediately after `execute_check_on_state` returns `Ok(applied_tx)`,
/// passing `applied_tx` as the first argument.  The transaction is re-applied to
/// `state` at `next_block_id` / `next_timestamp`; if it is accepted a second time
/// the function panics with a structured `INVARIANT VIOLATION [ReplayRejection]`
/// message.
///
/// # Why a standalone function?
///
/// `execute_check_on_state` consumes the `NSSATransaction` and returns it on `Ok`,
/// so the transaction is not available as a shared reference inside [`InvariantCtx`].
/// This function accepts ownership of the returned transaction and performs the
/// replay in-place.
///
/// # Example
///
/// ```rust,ignore
/// let result = tx.execute_check_on_state(&mut state, block_id, timestamp);
/// if let Ok(applied_tx) = result {
///     assert_replay_rejection(applied_tx, &mut state, block_id + 1, timestamp + 1);
/// }
/// ```
pub fn assert_replay_rejection(
    applied_tx: NSSATransaction,
    state: &mut V03State,
    next_block_id: u64,
    next_timestamp: u64,
) {
    let replay = applied_tx.execute_check_on_state(state, next_block_id, next_timestamp);
    assert!(
        replay.is_err(),
        "INVARIANT VIOLATION [ReplayRejection]: transaction accepted a second time — \
         nonce replay not prevented (replay block_id={next_block_id}, \
         replay timestamp={next_timestamp})",
    );
}

/// Assert that every signer account's nonce was incremented by exactly one after a
/// successfully applied transaction.
///
/// Call this immediately after `apply_state_diff` (or `execute_check_on_state`) succeeds,
/// passing the signer IDs derived from the transaction's witness set, the [`NonceSnapshot`]
/// captured **before** execution, and the post-execution state.
///
/// For a `NSSATransaction::Public(tx)`, derive signer IDs as:
///
/// ```rust,ignore
/// let signer_ids: Vec<nssa::AccountId> = tx
///     .witness_set()
///     .signatures_and_public_keys()
///     .iter()
///     .map(|(_, pk)| nssa::AccountId::from(pk))
///     .collect();
/// ```
///
/// For `NSSATransaction::ProgramDeployment`, there are no signers; pass an empty slice.
///
/// # Why a standalone function?
///
/// `apply_state_diff` consumes the `ValidatedStateDiff`, whose `signer_account_ids` field
/// is private to the `nssa` crate.  The caller must therefore derive signer IDs from the
/// transaction's witness set before consuming the diff, and thread them into this helper.
///
/// # Example
///
/// ```rust,ignore
/// let signer_ids = /* derived from tx.witness_set() */;
/// let nonces_before = NonceSnapshot(
///     signer_ids.iter().map(|&id| (id, state.get_account_by_id(id).nonce)).collect(),
/// );
/// state.apply_state_diff(diff);
/// assert_nonce_increment_correctness(&signer_ids, &nonces_before, &state);
/// ```
pub fn assert_nonce_increment_correctness(
    signer_ids: &[nssa::AccountId],
    nonces_before: &NonceSnapshot,
    state_after: &V03State,
) {
    for &id in signer_ids {
        let nonce_before = match nonces_before.0.get(&id) {
            Some(n) => *n,
            None => continue, // Account not in snapshot (e.g. newly created); skip.
        };
        let nonce_after = state_after.get_account_by_id(id).nonce;
        let expected = Nonce(
            nonce_before
                .0
                .checked_add(1)
                .expect("nonce overflow — signer nonce at u128::MAX"),
        );
        assert_eq!(
            nonce_after, expected,
            "INVARIANT VIOLATION [NonceIncrementCorrectness]: signer account {:?} nonce \
             not incremented by 1 after successful transaction \
             — before={:?}, expected={:?}, got={:?} \
             (apply_state_diff failed to increment nonce exactly once)",
            id, nonce_before, expected, nonce_after,
        );
    }
}

// ── Dispatcher ───────────────────────────────────────────────────────────────

/// Run every registered [`ProtocolInvariant`] and panic with a structured message
/// on the first violation.
///
/// Invariants checked:
/// - [`StateIsolationOnFailure`] — balances unchanged on rejection
/// - [`BalanceConservation`] — total balance conserved on success
/// - [`FailedTxNonceStability`] — nonces unchanged on rejection
/// - [`ReplayRejection`] — stub only; use [`assert_replay_rejection`] directly
/// - [`NonceIncrementCorrectness`] — stub only; use [`assert_nonce_increment_correctness`] directly
pub fn assert_invariants(ctx: &InvariantCtx<'_>) {
    let invariants: &[&dyn ProtocolInvariant] = &[
        &StateIsolationOnFailure,
        &BalanceConservation,
        &FailedTxNonceStability,
        &ReplayRejection,
        &NonceIncrementCorrectness,
    ];
    for inv in invariants {
        if let Some(violation) = inv.check(ctx) {
            panic!(
                "INVARIANT VIOLATION [{inv}]: {msg}",
                inv = violation.invariant,
                msg = violation.message,
            );
        }
    }
}

// ── Unit tests ────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use nssa::V03State;

    fn make_empty_state() -> V03State {
        V03State::new_with_genesis_accounts(&[], vec![], 0)
    }

    fn make_empty_snapshot() -> BalanceSnapshot {
        BalanceSnapshot(std::collections::HashMap::new())
    }

    fn make_empty_nonce_snapshot() -> NonceSnapshot {
        NonceSnapshot(std::collections::HashMap::new())
    }

    #[test]
    fn invariant_state_isolation_on_failure_does_not_panic_on_error() {
        let state = make_empty_state();
        let ctx = InvariantCtx {
            state_before: &state,
            state_after: &state,
            execution_succeeded: false,
            balances_before: make_empty_snapshot(),
            nonces_before: make_empty_nonce_snapshot(),
        };
        assert_invariants(&ctx);
    }

    #[test]
    fn invariant_replay_rejection_does_not_panic() {
        let state = make_empty_state();
        let ctx = InvariantCtx {
            state_before: &state,
            state_after: &state,
            execution_succeeded: true,
            balances_before: make_empty_snapshot(),
            nonces_before: make_empty_nonce_snapshot(),
        };
        assert_invariants(&ctx);
    }

    #[test]
    fn balance_conservation_catches_inflation_on_success() {
        // Arrange: one account with balance 100.
        let acc_id = nssa::AccountId::new([1u8; 32]);
        let state_before = V03State::new_with_genesis_accounts(&[(acc_id, 100)], vec![], 0);
        // Simulate execution that inflated the balance to 200.
        let state_after = V03State::new_with_genesis_accounts(&[(acc_id, 200)], vec![], 0);

        let mut balances = std::collections::HashMap::new();
        balances.insert(acc_id, 100u128);

        let ctx = InvariantCtx {
            state_before: &state_before,
            state_after: &state_after,
            execution_succeeded: true,
            balances_before: BalanceSnapshot(balances),
            nonces_before: make_empty_nonce_snapshot(),
        };

        let result = std::panic::catch_unwind(|| assert_invariants(&ctx));
        assert!(result.is_err(), "expected panic for balance inflation");
    }

    #[test]
    fn nonce_increment_correctness_passes_with_no_signers() {
        // Empty signer list — no accounts to check; trivially satisfies the invariant.
        let state = make_empty_state();
        assert_nonce_increment_correctness(&[], &make_empty_nonce_snapshot(), &state);
    }

    #[test]
    fn nonce_increment_correctness_passes_when_signer_not_in_snapshot() {
        // Signer ID is present in the list but absent from the snapshot — skipped.
        let acc_id = nssa::AccountId::new([9u8; 32]);
        let state = make_empty_state();
        // Empty snapshot → `continue` branch fires; no assertion is made.
        assert_nonce_increment_correctness(&[acc_id], &make_empty_nonce_snapshot(), &state);
    }

    #[test]
    fn nonce_increment_correctness_catches_unchanged_nonce() {
        // Arrange: signer has nonce 5 in the snapshot; the state returns Nonce(0) for the
        // same account (genesis default).  expected = Nonce(6), actual = Nonce(0) → VIOLATION.
        let acc_id = nssa::AccountId::new([3u8; 32]);
        let state = V03State::new_with_genesis_accounts(&[], vec![], 0);

        let mut nonces = std::collections::HashMap::new();
        nonces.insert(acc_id, Nonce(5));

        let result = std::panic::catch_unwind(|| {
            assert_nonce_increment_correctness(&[acc_id], &NonceSnapshot(nonces), &state);
        });
        assert!(result.is_err(), "expected panic for unchanged nonce");
    }

    #[test]
    fn failed_tx_nonce_stability_catches_nonce_mutation() {
        let acc_id = nssa::AccountId::new([2u8; 32]);
        // before: nonce 5; after: nonce 6 (should not happen on failure)
        let state_before = V03State::new_with_genesis_accounts(&[(acc_id, 100)], vec![], 0);
        let state_after = V03State::new_with_genesis_accounts(&[(acc_id, 100)], vec![], 0);

        // We check the nonce snapshot directly; the states both return default nonce (0).
        // Fake a discrepancy by inserting nonce=1 in the snapshot while state_after has nonce=0.
        let mut nonces = std::collections::HashMap::new();
        // Nonce(1) in snapshot, but state_after will return Nonce(0).
        nonces.insert(acc_id, Nonce(1));

        let mut balances = std::collections::HashMap::new();
        balances.insert(acc_id, 100u128);

        let ctx = InvariantCtx {
            state_before: &state_before,
            state_after: &state_after,
            execution_succeeded: false,
            balances_before: BalanceSnapshot(balances),
            nonces_before: NonceSnapshot(nonces),
        };

        let result = std::panic::catch_unwind(|| assert_invariants(&ctx));
        assert!(
            result.is_err(),
            "expected panic for nonce mutation on failure"
        );
    }
}

// ── ReplayRejection proptest suite ───────────────────────────────────────────
//
// This suite constitutes the formal, reproducible exercise of the ReplayRejection
// invariant.  It generates a realistic initial state and a correctly-signed
// native-transfer transaction, applies it once, and asserts that a second
// application is rejected.
//
// Run with: cargo test -p fuzz_props replay_rejection
#[cfg(test)]
mod replay_proptest {
    use crate::generators::{arb_native_transfer_tx, test_accounts};
    use nssa::V03State;
    use proptest::prelude::*;

    /// Build a `V03State` from the testnet accounts, assigning each a fixed
    /// balance large enough for any reasonable transfer amount.
    fn make_test_state() -> V03State {
        let accounts = test_accounts();
        let init_accs: Vec<(nssa::AccountId, u128)> = accounts
            .iter()
            .map(|(id, _)| (*id, 1_000_000u128))
            .collect();
        V03State::new_with_genesis_accounts(&init_accs, vec![], 0)
    }

    proptest! {
        /// **ReplayRejection** — a transaction accepted in block N must be
        /// rejected when replayed in block N+1, because the nonce is consumed
        /// on first acceptance.
        #[test]
        fn replay_rejection_proptest(tx in arb_native_transfer_tx(test_accounts())) {
            let mut state = make_test_state();

            // Stateless gate — skip structurally invalid transactions (e.g. those
            // whose public key does not match the declared sender).
            let validated_tx = match tx.transaction_stateless_check() {
                Ok(v) => v,
                Err(_) => return Ok(()),
            };

            // First application — may fail for state-level reasons (e.g. sender
            // has insufficient balance, wrong nonce).  In that case there is
            // nothing to replay.
            let first_result = validated_tx.execute_check_on_state(&mut state, 1, 0);

            if let Ok(validated_tx) = first_result {
                // Use the shared framework function.  assert_replay_rejection uses
                // assert!() rather than prop_assert!(); for structured proptest
                // inputs the framework-level panic is equivalent.
                super::assert_replay_rejection(validated_tx, &mut state, 2, 1);
            }
        }
    }
}