From 5d1adf1270759970731301f064ebf9748e1c59c9 Mon Sep 17 00:00:00 2001
From: Jack May <jack@solana.com>
Date: Tue, 19 Apr 2022 02:59:06 -0700
Subject: [PATCH] Fix signature count (#24471)

* Fix signature count

* protect the signature count futher
---
 runtime/src/bank.rs | 84 ++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 79 insertions(+), 5 deletions(-)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 9853e30e11..80d13a9f25 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -4131,8 +4131,6 @@ impl Bank {
                         feature_set_clone_time.as_us()
                     );
 
-                    signature_count += u64::from(tx.message().header().num_required_signatures);
-
                     let tx_wide_compute_cap = feature_set.is_active(&tx_wide_compute_cap::id());
                     let compute_budget_max_units = if tx_wide_compute_cap {
                         compute_budget::MAX_UNITS
@@ -4161,13 +4159,11 @@ impl Bank {
                         }
                     }
 
-                    let durable_nonce_fee = nonce.as_ref().map(DurableNonceFee::from);
-
                     self.execute_loaded_transaction(
                         tx,
                         loaded_transaction,
                         compute_budget,
-                        durable_nonce_fee,
+                        nonce.as_ref().map(DurableNonceFee::from),
                         enable_cpi_recording,
                         enable_log_recording,
                         enable_return_data_recording,
@@ -4268,6 +4264,10 @@ impl Bank {
             }
 
             if execution_result.was_executed() {
+                // Signature count must be accumulated only if the transaction
+                // is executed, otherwise a mismatched count between banking and
+                // replay could occur
+                signature_count += u64::from(tx.message().header().num_required_signatures);
                 executed_transactions_count += 1;
             }
 
@@ -16031,6 +16031,80 @@ pub(crate) mod tests {
         bank.process_transaction(&tx).unwrap();
     }
 
+    #[test]
+    fn test_failed_compute_request_instruction() {
+        solana_logger::setup();
+        let GenesisConfigInfo {
+            genesis_config,
+            mint_keypair,
+            ..
+        } = create_genesis_config_with_leader(
+            1_000_000_000_000_000,
+            &Pubkey::new_unique(),
+            bootstrap_validator_stake_lamports(),
+        );
+        let mut bank = Bank::new_for_tests(&genesis_config);
+
+        let payer0_keypair = Keypair::new();
+        let payer1_keypair = Keypair::new();
+        bank.transfer(10, &mint_keypair, &payer0_keypair.pubkey())
+            .unwrap();
+        bank.transfer(10, &mint_keypair, &payer1_keypair.pubkey())
+            .unwrap();
+
+        fn mock_ix_processor(
+            _first_instruction_account: usize,
+            invoke_context: &mut InvokeContext,
+        ) -> std::result::Result<(), InstructionError> {
+            let compute_budget = invoke_context.get_compute_budget();
+            assert_eq!(
+                *compute_budget,
+                ComputeBudget {
+                    max_units: 1,
+                    heap_size: Some(48 * 1024),
+                    ..ComputeBudget::default()
+                }
+            );
+            Ok(())
+        }
+        let program_id = solana_sdk::pubkey::new_rand();
+        bank.add_builtin("mock_program", &program_id, mock_ix_processor);
+
+        // This message will not be executed because the compute budget request is invalid
+        let message0 = Message::new(
+            &[
+                ComputeBudgetInstruction::request_heap_frame(1),
+                Instruction::new_with_bincode(program_id, &0, vec![]),
+            ],
+            Some(&payer0_keypair.pubkey()),
+        );
+        // This message will be processed successfully
+        let message1 = Message::new(
+            &[
+                ComputeBudgetInstruction::request_units(1, 0),
+                ComputeBudgetInstruction::request_heap_frame(48 * 1024),
+                Instruction::new_with_bincode(program_id, &0, vec![]),
+            ],
+            Some(&payer1_keypair.pubkey()),
+        );
+        let txs = vec![
+            Transaction::new(&[&payer0_keypair], message0, bank.last_blockhash()),
+            Transaction::new(&[&payer1_keypair], message1, bank.last_blockhash()),
+        ];
+        let results = bank.process_transactions(txs.iter());
+
+        assert_eq!(
+            results[0],
+            Err(TransactionError::InstructionError(
+                0,
+                InstructionError::InvalidInstructionData
+            ))
+        );
+        assert_eq!(results[1], Ok(()));
+        // two transfers and the mock program
+        assert_eq!(bank.signature_count(), 3);
+    }
+
     #[test]
     fn test_verify_and_hash_transaction_sig_len() {
         let GenesisConfigInfo {