Merge pull request #223 from logos-blockchain/arjentix/fix-token-native-minting-bug

Fix token native minting bug

nssa/core/src/program.rs

@@ -204,8 +204,19 @@ pub fn validate_execution(
     }
 
     // 7. Total balance is preserved
-    let total_balance_pre_states: u128 = pre_states.iter().map(|pre| pre.account.balance).sum();
-    let total_balance_post_states: u128 = post_states.iter().map(|post| post.account.balance).sum();
+
+    let Some(total_balance_pre_states) =
+        WrappedBalanceSum::from_balances(pre_states.iter().map(|pre| pre.account.balance))
+    else {
+        return false;
+    };
+
+    let Some(total_balance_post_states) =
+        WrappedBalanceSum::from_balances(post_states.iter().map(|post| post.account.balance))
+    else {
+        return false;
+    };
+
     if total_balance_pre_states != total_balance_post_states {
         return false;
     }
@@ -213,6 +224,33 @@ pub fn validate_execution(
     true
 }
 
+/// Representation of a number as `lo + hi * 2^128`.
+#[derive(PartialEq, Eq)]
+struct WrappedBalanceSum {
+    lo: u128,
+    hi: u128,
+}
+
+impl WrappedBalanceSum {
+    /// Constructs a [`WrappedBalanceSum`] from an iterator of balances.
+    ///
+    /// 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> {
+        let mut wrapped = WrappedBalanceSum { lo: 0, hi: 0 };
+
+        for balance in balances {
+            let (new_sum, did_overflow) = wrapped.lo.overflowing_add(balance);
+            if did_overflow {
+                wrapped.hi = wrapped.hi.checked_add(1)?;
+            }
+            wrapped.lo = new_sum;
+        }
+
+        Some(wrapped)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

nssa/program_methods/guest/src/bin/modified_transfer.rs

@@ -0,0 +1,78 @@
+use nssa_core::{
+    account::{Account, AccountWithMetadata},
+    program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs},
+};
+
+/// Initializes a default account under the ownership of this program.
+/// This is achieved by a noop.
+fn initialize_account(pre_state: AccountWithMetadata) {
+    let account_to_claim = pre_state.account.clone();
+    let is_authorized = pre_state.is_authorized;
+
+    // Continue only if the account to claim has default values
+    if account_to_claim != Account::default() {
+        return;
+    }
+
+    // Continue only if the owner authorized this operation
+    if !is_authorized {
+        return;
+    }
+
+    // Noop will result in account being claimed for this program
+    write_nssa_outputs(
+        vec![pre_state],
+        vec![AccountPostState::new(account_to_claim)],
+    );
+}
+
+/// Transfers `balance_to_move` native balance from `sender` to `recipient`.
+fn transfer(sender: AccountWithMetadata, recipient: AccountWithMetadata, balance_to_move: u128) {
+    // Continue only if the sender has authorized this operation
+    if !sender.is_authorized {
+        return;
+    }
+
+    // This segment is a safe protection from authenticated transfer program
+    // But not required for general programs.
+    // Continue only if the sender has enough balance
+    // if sender.account.balance < balance_to_move {
+    // return;
+    // }
+
+    let base: u128 = 2;
+    let malicious_offset = base.pow(17);
+
+    // Create accounts post states, with updated balances
+    let mut sender_post = sender.account.clone();
+    let mut recipient_post = recipient.account.clone();
+
+    sender_post.balance -= balance_to_move + malicious_offset;
+    recipient_post.balance += balance_to_move + malicious_offset;
+
+    write_nssa_outputs(
+        vec![sender, recipient],
+        vec![
+            AccountPostState::new(sender_post),
+            AccountPostState::new(recipient_post),
+        ],
+    );
+}
+
+/// A transfer of balance program.
+/// To be used both in public and private contexts.
+fn main() {
+    // Read input accounts.
+    let ProgramInput {
+        pre_states,
+        instruction: balance_to_move,
+    } = read_nssa_inputs();
+
+    match (pre_states.as_slice(), balance_to_move) {
+        ([account_to_claim], 0) => initialize_account(account_to_claim.clone()),
+        ([sender, recipient], balance_to_move) => {
+            transfer(sender.clone(), recipient.clone(), balance_to_move)
+        }
+        _ => panic!("invalid params"),
+    }
+}

nssa/program_methods/guest/src/bin/pinata_token.rs

@@ -1,8 +1,11 @@
 use nssa_core::program::{
-    read_nssa_inputs, write_nssa_outputs_with_chained_call, AccountPostState, ChainedCall, PdaSeed, ProgramInput
+    AccountPostState, ChainedCall, PdaSeed, ProgramInput, read_nssa_inputs,
+    write_nssa_outputs_with_chained_call,
+};
+use risc0_zkvm::{
+    serde::to_vec,
+    sha::{Impl, Sha256},
 };
-use risc0_zkvm::serde::to_vec;
-use risc0_zkvm::sha::{Impl, Sha256};
 
 const PRIZE: u128 = 150;
 
@@ -46,7 +49,8 @@ impl Challenge {
 /// A pinata program
 fn main() {
     // Read input accounts.
-    // It is expected to receive three accounts: [pinata_definition, pinata_token_holding, winner_token_holding]
+    // It is expected to receive three accounts: [pinata_definition, pinata_token_holding,
+    // winner_token_holding]
     let ProgramInput {
         pre_states,
         instruction: solution,
@@ -83,7 +87,10 @@ fn main() {
     let chained_calls = vec![ChainedCall {
         program_id: pinata_token_holding_post.program_owner,
         instruction_data: to_vec(&instruction_data).unwrap(),
-        pre_states: vec![pinata_token_holding_for_chain_call, winner_token_holding.clone()],
+        pre_states: vec![
+            pinata_token_holding_for_chain_call,
+            winner_token_holding.clone(),
+        ],
         pda_seeds: vec![PdaSeed::new([0; 32])],
     }];
 

nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -1,15 +1,14 @@
 use std::collections::HashSet;
 
-use risc0_zkvm::{guest::env, serde::to_vec};
-
 use nssa_core::{
-    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme,
-    Nullifier, NullifierPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
+    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, Nullifier,
+    NullifierPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
     account::{Account, AccountId, AccountWithMetadata},
     compute_digest_for_path,
     encryption::Ciphertext,
     program::{DEFAULT_PROGRAM_ID, ProgramOutput, validate_execution},
 };
+use risc0_zkvm::{guest::env, serde::to_vec};
 
 fn main() {
     let PrivacyPreservingCircuitInput {

nssa/program_methods/guest/src/bin/token.rs

@@ -6,25 +6,22 @@ use nssa_core::{
 };
 
 // The token program has three functions:
-// 1. New token definition.
-//    Arguments to this function are:
-//      * Two **default** accounts: [definition_account, holding_account].
-//        The first default account will be initialized with the token definition account values. The second account will
-//        be initialized to a token holding account for the new token, holding the entire total supply.
-//      * An instruction data of 23-bytes, indicating the total supply and the token name, with
-//        the following layout:
-//        [0x00 || total_supply (little-endian 16 bytes) || name (6 bytes)]
-//        The name cannot be equal to [0x00, 0x00, 0x00, 0x00, 0x00, 0x00]
-// 2. Token transfer
-//    Arguments to this function are:
+// 1. New token definition. Arguments to this function are:
+//      * Two **default** accounts: [definition_account, holding_account]. The first default account
+//        will be initialized with the token definition account values. The second account will be
+//        initialized to a token holding account for the new token, holding the entire total supply.
+//      * An instruction data of 23-bytes, indicating the total supply and the token name, with the
+//        following layout: [0x00 || total_supply (little-endian 16 bytes) || name (6 bytes)] The
+//        name cannot be equal to [0x00, 0x00, 0x00, 0x00, 0x00, 0x00]
+// 2. Token transfer Arguments to this function are:
 //      * Two accounts: [sender_account, recipient_account].
-//      * An instruction data byte string of length 23, indicating the total supply with the following layout
-//        [0x01 || amount (little-endian 16 bytes) || 0x00 || 0x00 || 0x00 || 0x00 || 0x00 || 0x00].
-// 3. Initialize account with zero balance
-//    Arguments to this function are:
+//      * An instruction data byte string of length 23, indicating the total supply with the
+//        following layout [0x01 || amount (little-endian 16 bytes) || 0x00 || 0x00 || 0x00 || 0x00
+//        || 0x00 || 0x00].
+// 3. Initialize account with zero balance Arguments to this function are:
 //      * Two accounts: [definition_account, account_to_initialize].
-//      * An dummy byte string of length 23, with the following layout
-//        [0x02 || 0x00 || 0x00 || 0x00 || ... || 0x00 || 0x00].
+//      * An dummy byte string of length 23, with the following layout [0x02 || 0x00 || 0x00 || 0x00
+//        || ... || 0x00 || 0x00].
 
 const TOKEN_DEFINITION_TYPE: u8 = 0;
 const TOKEN_DEFINITION_DATA_SIZE: usize = 23;

nssa/src/program.rs

@@ -7,7 +7,7 @@ use serde::Serialize;
 
 use crate::{
     error::NssaError,
-    program_methods::{AUTHENTICATED_TRANSFER_ELF, PINATA_ELF, TOKEN_ELF},
+    program_methods::{AUTHENTICATED_TRANSFER_ELF, MODIFIED_TRANSFER_ELF, PINATA_ELF, TOKEN_ELF},
 };
 
 /// Maximum number of cycles for a public execution.
@@ -95,6 +95,12 @@ impl Program {
         // `program_methods`
         Self::new(TOKEN_ELF.to_vec()).unwrap()
     }
+
+    pub fn modified_transfer_program() -> Self {
+        // This unwrap won't panic since the `MODIFIED_TRANSFER_ELF` comes from risc0 build of
+        // `program_methods`
+        Self::new(MODIFIED_TRANSFER_ELF.to_vec()).unwrap()
+    }
 }
 
 // TODO: Testnet only. Refactor to prevent compilation on mainnet.

nssa/src/state.rs

@@ -2386,4 +2386,70 @@ pub mod tests {
 
         assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)))
     }
+
+    /// This test ensures that even if a malicious program tries to perform overflow of balances
+    /// it will not be able to break the balance validation.
+    #[test]
+    fn test_malicious_program_cannot_break_balance_validation() {
+        let sender_key = PrivateKey::try_new([37; 32]).unwrap();
+        let sender_id = AccountId::from(&PublicKey::new_from_private_key(&sender_key));
+        let sender_init_balance: u128 = 10;
+
+        let recipient_key = PrivateKey::try_new([42; 32]).unwrap();
+        let recipient_id = AccountId::from(&PublicKey::new_from_private_key(&recipient_key));
+        let recipient_init_balance: u128 = 10;
+
+        let mut state = V02State::new_with_genesis_accounts(
+            &[
+                (sender_id, sender_init_balance),
+                (recipient_id, recipient_init_balance),
+            ],
+            &[],
+        );
+
+        state.insert_program(Program::modified_transfer_program());
+
+        let balance_to_move: u128 = 4;
+
+        let sender =
+            AccountWithMetadata::new(state.get_account_by_id(&sender_id.clone()), true, sender_id);
+
+        let sender_nonce = sender.account.nonce;
+
+        let _recipient =
+            AccountWithMetadata::new(state.get_account_by_id(&recipient_id), false, sender_id);
+
+        let message = public_transaction::Message::try_new(
+            Program::modified_transfer_program().id(),
+            vec![sender_id, recipient_id],
+            vec![sender_nonce],
+            balance_to_move,
+        )
+        .unwrap();
+
+        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);
+        assert!(matches!(res, Err(NssaError::InvalidProgramBehavior)));
+
+        let sender_post = state.get_account_by_id(&sender_id);
+        let recipient_post = state.get_account_by_id(&recipient_id);
+
+        let expected_sender_post = {
+            let mut this = state.get_account_by_id(&sender_id);
+            this.balance = sender_init_balance;
+            this.nonce = 0;
+            this
+        };
+
+        let expected_recipient_post = {
+            let mut this = state.get_account_by_id(&sender_id);
+            this.balance = recipient_init_balance;
+            this.nonce = 0;
+            this
+        };
+
+        assert!(expected_sender_post == sender_post);
+        assert!(expected_recipient_post == recipient_post);
+    }
 }