Merge 81d9c2a925a656070b0379a812a04ba9b44b9396 into 42a2f04cd52779b109750e86dcd2f5486ec3ec4f

nssa/core/src/program.rs

@@ -376,8 +376,8 @@ impl ProgramOutput {
 }
 
 /// Representation of a number as `lo + hi * 2^128`.
-#[derive(PartialEq, Eq)]
-struct WrappedBalanceSum {
+#[derive(Debug, PartialEq, Eq)]
+pub struct WrappedBalanceSum {
     lo: u128,
     hi: u128,
 }
@@ -387,7 +387,7 @@ impl WrappedBalanceSum {
     ///
     /// Returns [`None`] if balance sum overflows `lo + hi * 2^128` representation, which is not
     /// expected in practical scenarios.
-    fn from_balances(balances: impl Iterator<Item = u128>) -> Option<Self> {
+    pub fn from_balances(balances: impl Iterator<Item = u128>) -> Option<Self> {
         let mut wrapped = Self { lo: 0, hi: 0 };
 
         for balance in balances {
@@ -402,6 +402,75 @@ impl WrappedBalanceSum {
     }
 }
 
+impl std::fmt::Display for WrappedBalanceSum {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if self.hi == 0 {
+            write!(f, "{}", self.lo)
+        } else {
+            write!(f, "{} * 2^128 + {}", self.hi, self.lo)
+        }
+    }
+}
+
+impl From<u128> for WrappedBalanceSum {
+    fn from(value: u128) -> Self {
+        Self { lo: value, hi: 0 }
+    }
+}
+
+#[derive(thiserror::Error, Debug)]
+pub enum ExecutionValidationError {
+    #[error("Pre-state account IDs are not unique")]
+    PreStateAccountIdsNotUnique,
+
+    #[error(
+        "Pre-state and post-state lengths do not match: pre-state length {pre_state_length}, post-state length {post_state_length}"
+    )]
+    MismatchedPreStatePostStateLength {
+        pre_state_length: usize,
+        post_state_length: usize,
+    },
+
+    #[error("Unallowed modification of nonce for account {account_id}")]
+    ModifiedNonce { account_id: AccountId },
+
+    #[error("Unallowed modification of program owner for account {account_id}")]
+    ModifiedProgramOwner { account_id: AccountId },
+
+    #[error(
+        "Trying to decrease balance of account {account_id} owned by {owner_program_id:?} in a program {executing_program_id:?} which is not the owner"
+    )]
+    DecreasingBalanceNotOwnedByProgram {
+        account_id: AccountId,
+        owner_program_id: ProgramId,
+        executing_program_id: ProgramId,
+    },
+
+    #[error(
+        "Unauthorized modification of data for account {account_id} which is not default and not owned by executing program {executing_program_id:?}"
+    )]
+    UnauthorizedDataModification {
+        account_id: AccountId,
+        executing_program_id: ProgramId,
+    },
+
+    #[error(
+        "Post-state for account {account_id} has default program owner but pre-state was not default"
+    )]
+    AccountWithDefaultOwnerMustBeDefault { account_id: AccountId },
+
+    #[error("Total balance across accounts overflowed 2^256 - 1")]
+    BalanceSumOverflow,
+
+    #[error(
+        "Total balance across accounts is not preserved: total balance in pre-states {total_balance_pre_states}, total balance in post-states {total_balance_post_states}"
+    )]
+    MismatchedTotalBalance {
+        total_balance_pre_states: WrappedBalanceSum,
+        total_balance_post_states: WrappedBalanceSum,
+    },
+}
+
 #[must_use]
 pub fn compute_authorized_pdas(
     caller_program_id: Option<ProgramId>,
@@ -440,31 +509,39 @@ pub fn read_nssa_inputs<T: DeserializeOwned>() -> (ProgramInput<T>, InstructionD
 /// - `pre_states`: The list of input accounts, each annotated with authorization metadata.
 /// - `post_states`: The list of resulting accounts after executing the program logic.
 /// - `executing_program_id`: The identifier of the program that was executed.
-#[must_use]
 pub fn validate_execution(
     pre_states: &[AccountWithMetadata],
     post_states: &[AccountPostState],
     executing_program_id: ProgramId,
-) -> bool {
+) -> Result<(), ExecutionValidationError> {
     // 1. Check account ids are all different
     if !validate_uniqueness_of_account_ids(pre_states) {
-        return false;
+        return Err(ExecutionValidationError::PreStateAccountIdsNotUnique);
     }
 
     // 2. Lengths must match
     if pre_states.len() != post_states.len() {
-        return false;
+        return Err(
+            ExecutionValidationError::MismatchedPreStatePostStateLength {
+                pre_state_length: pre_states.len(),
+                post_state_length: post_states.len(),
+            },
+        );
     }
 
     for (pre, post) in pre_states.iter().zip(post_states) {
         // 3. Nonce must remain unchanged
         if pre.account.nonce != post.account.nonce {
-            return false;
+            return Err(ExecutionValidationError::ModifiedNonce {
+                account_id: pre.account_id,
+            });
         }
 
         // 4. Program ownership changes are not allowed
         if pre.account.program_owner != post.account.program_owner {
-            return false;
+            return Err(ExecutionValidationError::ModifiedProgramOwner {
+                account_id: pre.account_id,
+            });
         }
 
         let account_program_owner = pre.account.program_owner;
@@ -473,7 +550,13 @@ pub fn validate_execution(
         if post.account.balance < pre.account.balance
             && account_program_owner != executing_program_id
         {
-            return false;
+            return Err(
+                ExecutionValidationError::DecreasingBalanceNotOwnedByProgram {
+                    account_id: pre.account_id,
+                    owner_program_id: account_program_owner,
+                    executing_program_id,
+                },
+            );
         }
 
         // 6. Data changes only allowed if owned by executing program or if account pre state has
@@ -482,13 +565,20 @@ pub fn validate_execution(
             && pre.account != Account::default()
             && account_program_owner != executing_program_id
         {
-            return false;
+            return Err(ExecutionValidationError::UnauthorizedDataModification {
+                account_id: pre.account_id,
+                executing_program_id,
+            });
         }
 
         // 7. If a post state has default program owner, the pre state must have been a default
         //    account
         if post.account.program_owner == DEFAULT_PROGRAM_ID && pre.account != Account::default() {
-            return false;
+            return Err(
+                ExecutionValidationError::AccountWithDefaultOwnerMustBeDefault {
+                    account_id: pre.account_id,
+                },
+            );
         }
     }
 
@@ -497,20 +587,23 @@ pub fn validate_execution(
     let Some(total_balance_pre_states) =
         WrappedBalanceSum::from_balances(pre_states.iter().map(|pre| pre.account.balance))
     else {
-        return false;
+        return Err(ExecutionValidationError::BalanceSumOverflow);
     };
 
     let Some(total_balance_post_states) =
         WrappedBalanceSum::from_balances(post_states.iter().map(|post| post.account.balance))
     else {
-        return false;
+        return Err(ExecutionValidationError::BalanceSumOverflow);
     };
 
     if total_balance_pre_states != total_balance_post_states {
-        return false;
+        return Err(ExecutionValidationError::MismatchedTotalBalance {
+            total_balance_pre_states,
+            total_balance_post_states,
+        });
     }
 
-    true
+    Ok(())
 }
 
 fn validate_uniqueness_of_account_ids(pre_states: &[AccountWithMetadata]) -> bool {

nssa/src/error.rs

@@ -1,12 +1,16 @@
 use std::io;
 
+use nssa_core::{
+    account::{Account, AccountId},
+    program::ProgramId,
+};
 use thiserror::Error;
 
 #[macro_export]
 macro_rules! ensure {
     ($cond:expr, $err:expr) => {
         if !$cond {
-            return Err($err);
+            return Err($err.into());
         }
     };
 }
@@ -17,7 +21,7 @@ pub enum NssaError {
     InvalidInput(String),
 
     #[error("Program violated execution rules")]
-    InvalidProgramBehavior,
+    InvalidProgramBehavior(#[from] InvalidProgramBehaviorError),
 
     #[error("Serialization error: {0}")]
     InstructionSerializationError(String),
@@ -32,15 +36,15 @@ pub enum NssaError {
     InvalidPublicKey(#[source] k256::schnorr::Error),
 
     #[error("Invalid hex for public key")]
-    InvalidHexPublicKey(hex::FromHexError),
+    InvalidHexPublicKey(#[source] hex::FromHexError),
 
-    #[error("Risc0 error: {0}")]
+    #[error("Failed to write program input: {0}")]
     ProgramWriteInputFailed(String),
 
-    #[error("Risc0 error: {0}")]
+    #[error("Failed to execute program: {0}")]
     ProgramExecutionFailed(String),
 
-    #[error("Risc0 error: {0}")]
+    #[error("Failed to prove program: {0}")]
     ProgramProveFailed(String),
 
     #[error("Invalid transaction: {0}")]
@@ -77,6 +81,55 @@ pub enum NssaError {
     OutOfValidityWindow,
 }
 
+#[derive(Error, Debug)]
+pub enum InvalidProgramBehaviorError {
+    #[error(
+        "Inconsistent pre-state for account {account_id} : expected {expected:?}, actual {actual:?}"
+    )]
+    InconsistentAccountPreState {
+        account_id: AccountId,
+        // Boxed to reduce the size of the error type
+        expected: Box<Account>,
+        actual: Box<Account>,
+    },
+
+    #[error(
+        "Inconsistent authorization for account {account_id} : expected {expected_authorization}, actual {actual_authorization}"
+    )]
+    InconsistentAccountAuthorization {
+        account_id: AccountId,
+        expected_authorization: bool,
+        actual_authorization: bool,
+    },
+
+    #[error("Program ID mismatch: expected {expected:?}, actual {actual:?}")]
+    MismatchedProgramId {
+        expected: ProgramId,
+        actual: ProgramId,
+    },
+
+    #[error(transparent)]
+    ExecutionValidationFailed(#[from] nssa_core::program::ExecutionValidationError),
+
+    #[error("Trying to claim account {account_id} which is not default")]
+    ClaimedNonDefaultAccount { account_id: AccountId },
+
+    #[error("Trying to claim account {account_id} which is not authorized")]
+    ClaimedUnauthorizedAccount { account_id: AccountId },
+
+    #[error("PDA claim mismatch: expected {expected:?}, actual {actual:?}")]
+    MismatchedPdaClaim {
+        expected: AccountId,
+        actual: AccountId,
+    },
+
+    #[error("Default account {account_id} was modified without being claimed")]
+    DefaultAccountModifiedWithoutClaim { account_id: AccountId },
+
+    #[error("Called program {program_id:?} which is not listed in dependencies")]
+    CalledProgramWhichIsNotListedInDependencies { program_id: ProgramId },
+}
+
 #[cfg(test)]
 mod tests {
 

nssa/src/privacy_preserving_transaction/circuit.rs

@@ -10,7 +10,7 @@ use nssa_core::{
 use risc0_zkvm::{ExecutorEnv, InnerReceipt, ProverOpts, Receipt, default_prover};
 
 use crate::{
-    error::NssaError,
+    error::{InvalidProgramBehaviorError, NssaError},
     program::Program,
     program_methods::{PRIVACY_PRESERVING_CIRCUIT_ELF, PRIVACY_PRESERVING_CIRCUIT_ID},
     state::MAX_NUMBER_CHAINED_CALLS,
@@ -112,9 +112,11 @@ pub fn execute_and_prove(
         env_builder.add_assumption(inner_receipt);
 
         for new_call in program_output.chained_calls.into_iter().rev() {
-            let next_program = dependencies
-                .get(&new_call.program_id)
-                .ok_or(NssaError::InvalidProgramBehavior)?;
+            let next_program = dependencies.get(&new_call.program_id).ok_or(
+                InvalidProgramBehaviorError::CalledProgramWhichIsNotListedInDependencies {
+                    program_id: new_call.program_id,
+                },
+            )?;
             chained_calls.push_front((new_call, next_program));
         }
 

nssa/src/public_transaction/transaction.rs

@@ -11,7 +11,7 @@ use sha2::{Digest as _, digest::FixedOutput as _};
 
 use crate::{
     V03State, ensure,
-    error::NssaError,
+    error::{InvalidProgramBehaviorError, NssaError},
     public_transaction::{Message, WitnessSet},
     state::MAX_NUMBER_CHAINED_CALLS,
 };
@@ -173,33 +173,43 @@ impl PublicTransaction {
                     .unwrap_or_else(|| state.get_account_by_id(account_id));
                 ensure!(
                     pre.account == expected_pre,
-                    NssaError::InvalidProgramBehavior
+                    InvalidProgramBehaviorError::InconsistentAccountPreState {
+                        account_id,
+                        expected: Box::new(expected_pre),
+                        actual: Box::new(pre.account.clone())
+                    }
                 );
 
                 // Check that authorization flags are consistent with the provided ones or
                 // authorized by program through the PDA mechanism
+                let expected_is_authorized = is_authorized(&account_id);
                 ensure!(
-                    pre.is_authorized == is_authorized(&account_id),
-                    NssaError::InvalidProgramBehavior
+                    pre.is_authorized == expected_is_authorized,
+                    InvalidProgramBehaviorError::InconsistentAccountAuthorization {
+                        account_id,
+                        expected_authorization: expected_is_authorized,
+                        actual_authorization: pre.is_authorized
+                    }
                 );
             }
 
             // Verify that the program output's self_program_id matches the expected program ID.
             ensure!(
                 program_output.self_program_id == chained_call.program_id,
-                NssaError::InvalidProgramBehavior
+                InvalidProgramBehaviorError::MismatchedProgramId {
+                    expected: chained_call.program_id,
+                    actual: program_output.self_program_id
+                }
             );
 
             // Verify execution corresponds to a well-behaved program.
             // See the # Programs section for the definition of the `validate_execution` method.
-            ensure!(
-                validate_execution(
-                    &program_output.pre_states,
-                    &program_output.post_states,
-                    chained_call.program_id,
-                ),
-                NssaError::InvalidProgramBehavior
-            );
+            validate_execution(
+                &program_output.pre_states,
+                &program_output.post_states,
+                chained_call.program_id,
+            )
+            .map_err(InvalidProgramBehaviorError::ExecutionValidationFailed)?;
 
             // Verify validity window
             ensure!(
@@ -214,27 +224,33 @@ impl PublicTransaction {
                 let Some(claim) = post.required_claim() else {
                     continue;
                 };
+                let account_id = program_output.pre_states[i].account_id;
+
                 // The invoked program can only claim accounts with default program id.
                 ensure!(
                     post.account().program_owner == DEFAULT_PROGRAM_ID,
-                    NssaError::InvalidProgramBehavior
+                    InvalidProgramBehaviorError::ClaimedNonDefaultAccount { account_id }
                 );
 
-                let account_id = program_output.pre_states[i].account_id;
-
                 match claim {
                     Claim::Authorized => {
                         // The program can only claim accounts that were authorized by the signer.
                         ensure!(
                             is_authorized(&account_id),
-                            NssaError::InvalidProgramBehavior
+                            InvalidProgramBehaviorError::ClaimedUnauthorizedAccount { account_id }
                         );
                     }
                     Claim::Pda(seed) => {
                         // The program can only claim accounts that correspond to the PDAs it is
                         // authorized to claim.
                         let pda = AccountId::from((&chained_call.program_id, &seed));
-                        ensure!(account_id == pda, NssaError::InvalidProgramBehavior);
+                        ensure!(
+                            account_id == pda,
+                            InvalidProgramBehaviorError::MismatchedPdaClaim {
+                                expected: pda,
+                                actual: account_id
+                            }
+                        );
                     }
                 }
 
@@ -260,7 +276,7 @@ impl PublicTransaction {
         }
 
         // Check that all modified uninitialized accounts where claimed
-        for post in state_diff.iter().filter_map(|(account_id, post)| {
+        for (account_id, post) in state_diff.iter().filter_map(|(account_id, post)| {
             let pre = state.get_account_by_id(*account_id);
             if pre.program_owner != DEFAULT_PROGRAM_ID {
                 return None;
@@ -268,11 +284,11 @@ impl PublicTransaction {
             if pre == *post {
                 return None;
             }
-            Some(post)
+            Some((*account_id, post))
         }) {
             ensure!(
                 post.program_owner != DEFAULT_PROGRAM_ID,
-                NssaError::InvalidProgramBehavior
+                InvalidProgramBehaviorError::DefaultAccountModifiedWithoutClaim { account_id }
             );
         }
 

nssa/src/state.rs

@@ -346,12 +346,15 @@ pub mod tests {
         Timestamp,
         account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data},
         encryption::{EphemeralPublicKey, Scalar, ViewingPublicKey},
-        program::{BlockValidityWindow, PdaSeed, ProgramId, TimestampValidityWindow},
+        program::{
+            BlockValidityWindow, ExecutionValidationError, PdaSeed, ProgramId,
+            TimestampValidityWindow, WrappedBalanceSum,
+        },
     };
 
     use crate::{
         PublicKey, PublicTransaction, V03State,
-        error::NssaError,
+        error::{InvalidProgramBehaviorError, NssaError},
         execute_and_prove,
         privacy_preserving_transaction::{
             PrivacyPreservingTransaction,
@@ -680,10 +683,11 @@ pub mod tests {
 
     #[test]
     fn program_should_fail_if_modifies_nonces() {
-        let initial_data = [(AccountId::new([1; 32]), 100)];
+        let account_id = AccountId::new([1; 32]);
+        let initial_data = [(account_id, 100)];
         let mut state =
             V03State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
-        let account_ids = vec![AccountId::new([1; 32])];
+        let account_ids = vec![account_id];
         let program_id = Program::nonce_changer_program().id();
         let message =
             public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap();
@@ -692,7 +696,14 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(
+                InvalidProgramBehaviorError::ExecutionValidationFailed(
+                    ExecutionValidationError::ModifiedNonce { account_id: err_account_id }
+                )
+            )) if err_account_id == account_id
+        ));
     }
 
     #[test]
@@ -709,7 +720,17 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(
+                InvalidProgramBehaviorError::ExecutionValidationFailed(
+                    ExecutionValidationError::MismatchedPreStatePostStateLength {
+                        pre_state_length,
+                        post_state_length
+                    }
+                )
+            )) if pre_state_length == 1 && post_state_length == 2
+        ));
     }
 
     #[test]
@@ -726,7 +747,17 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(
+                InvalidProgramBehaviorError::ExecutionValidationFailed(
+                    ExecutionValidationError::MismatchedPreStatePostStateLength {
+                        pre_state_length,
+                        post_state_length
+                    }
+                )
+            )) if pre_state_length == 2 && post_state_length == 1
+        ));
     }
 
     #[test]
@@ -750,7 +781,12 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id }
+            ))) if err_account_id == account_id
+        ));
     }
 
     #[test]
@@ -774,7 +810,12 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id }
+            ))) if err_account_id == account_id
+        ));
     }
 
     #[test]
@@ -798,7 +839,12 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id }
+            ))) if err_account_id == account_id
+        ));
     }
 
     #[test]
@@ -822,16 +868,21 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id }
+            ))) if err_account_id == account_id
+        ));
     }
 
     #[test]
     fn program_should_fail_if_transfers_balance_from_non_owned_account() {
-        let initial_data = [(AccountId::new([1; 32]), 100)];
-        let mut state =
-            V03State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
         let sender_account_id = AccountId::new([1; 32]);
         let receiver_account_id = AccountId::new([2; 32]);
+        let initial_data = [(sender_account_id, 100)];
+        let mut state =
+            V03State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
         let balance_to_move: u128 = 1;
         let program_id = Program::simple_balance_transfer().id();
         assert_ne!(
@@ -850,7 +901,12 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::DecreasingBalanceNotOwnedByProgram { account_id: err_account_id, owner_program_id, executing_program_id }
+            ))) if err_account_id == sender_account_id && owner_program_id != program_id && executing_program_id == program_id
+        ));
     }
 
     #[test]
@@ -875,7 +931,12 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::UnauthorizedDataModification { account_id: err_account_id, executing_program_id }
+            ))) if err_account_id == account_id && executing_program_id == program_id
+        ));
     }
 
     #[test]
@@ -893,7 +954,12 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states }
+            ))) if total_balance_pre_states == 0.into() && total_balance_post_states == 1.into()
+        ));
     }
 
     #[test]
@@ -922,7 +988,12 @@ pub mod tests {
         let tx = PublicTransaction::new(message, witness_set);
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states }
+            ))) if total_balance_pre_states == 100.into() && total_balance_post_states == 99.into()
+        ));
     }
 
     fn test_public_account_keys_1() -> TestPublicKeys {
@@ -2824,7 +2895,12 @@ pub mod tests {
 
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(
+                InvalidProgramBehaviorError::ClaimedNonDefaultAccount { account_id: err_account_id }
+            )) if err_account_id == account_id
+        ));
     }
 
     /// This test ensures that even if a malicious program tries to perform overflow of balances
@@ -2869,7 +2945,22 @@ pub mod tests {
         let witness_set = public_transaction::WitnessSet::for_message(&message, &[&sender_key]);
         let tx = PublicTransaction::new(message, witness_set);
         let res = state.transition_from_public_transaction(&tx, 1, 0);
-        assert!(matches!(res, Err(NssaError::InvalidProgramBehavior)));
+        let expected_total_balance_pre_states = WrappedBalanceSum::from_balances(
+            [sender_init_balance, recipient_init_balance].into_iter(),
+        )
+        .unwrap();
+        let expected_total_balance_post_states = WrappedBalanceSum::from_balances(
+            [sender_init_balance, recipient_init_balance, u128::MAX, 1].into_iter(),
+        )
+        .unwrap();
+        assert!(matches!(
+            res,
+            Err(NssaError::InvalidProgramBehavior(
+                InvalidProgramBehaviorError::ExecutionValidationFailed(
+                    ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states }
+                )
+            )) if total_balance_pre_states == expected_total_balance_pre_states && total_balance_post_states == expected_total_balance_post_states
+        ));
 
         let sender_post = state.get_account_by_id(sender_id);
         let recipient_post = state.get_account_by_id(recipient_id);
@@ -3114,7 +3205,14 @@ pub mod tests {
         let result = state.transition_from_public_transaction(&tx, 1, 0);
 
         // Should fail - cannot modify data without claiming the account
-        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+        assert!(matches!(
+            result,
+            Err(NssaError::InvalidProgramBehavior(
+                InvalidProgramBehaviorError::DefaultAccountModifiedWithoutClaim {
+                    account_id: err_account_id
+                }
+            )) if err_account_id == account_id
+        ));
     }
 
     #[test]

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -116,12 +116,17 @@ impl ExecutionState {
 
             // Check that the program is well behaved.
             // See the # Programs section for the definition of the `validate_execution` method.
-            let execution_valid = validate_execution(
+            let validated_execution = validate_execution(
                 &program_output.pre_states,
                 &program_output.post_states,
                 chained_call.program_id,
             );
-            assert!(execution_valid, "Bad behaved program");
+            if let Err(err) = validated_execution {
+                panic!(
+                    "Invalid program behavior in program {:?}: {err}",
+                    chained_call.program_id
+                );
+            }
 
             for next_call in program_output.chained_calls.iter().rev() {
                 chained_calls.push_front((next_call.clone(), Some(chained_call.program_id)));