Explicitly sanitize program id indexes before usage

1. check transaction has valid program_id before using it to avoid possible panic;
2. change calculate_cost function signature to return Result;
3. add CostModelError enum, update return type from Result<_, str> to Result<_, CostModelError>

core/src/banking_stage.rs

@@ -1217,7 +1217,14 @@ impl BankingStage {
                 cost_tracker
                     .write()
                     .unwrap()
-                    .add_transaction_cost(tx.transaction());
+                    .add_transaction_cost(tx.transaction())
+                    .unwrap_or_else(|err| {
+                        warn!(
+                            "failed to track transaction cost, err {:?}, tx {:?}",
+                            err,
+                            tx.transaction()
+                        )
+                    });
             }
         });
         cost_tracking_time.stop();

core/src/cost_model.rs

@@ -28,6 +28,19 @@ pub const BLOCK_MAX_COST: u64 = 2_500_000_000;
 
 const MAX_WRITABLE_ACCOUNTS: usize = 256;
 
+#[derive(Debug, Clone)]
+pub enum CostModelError {
+    /// transaction that would fail sanitize, cost model is not able to process
+    /// such transaction.
+    InvalidTransaction,
+
+    /// would exceed block max limit
+    WouldExceedBlockMaxLimit,
+
+    /// would exceed account max limit
+    WouldExceedAccountMaxLimit,
+}
+
 // cost of transaction is made of account_access_cost and instruction execution_cost
 // where
 // account_access_cost is the sum of read/write/sign all accounts included in the transaction
@@ -113,9 +126,16 @@ impl CostModel {
         );
     }
 
-    pub fn calculate_cost(&mut self, transaction: &Transaction) -> &TransactionCost {
+    pub fn calculate_cost(
+        &mut self,
+        transaction: &Transaction,
+    ) -> Result<&TransactionCost, CostModelError> {
         self.transaction_cost.reset();
 
+        // calculate transaction exeution cost
+        self.transaction_cost.execution_cost = self.find_transaction_cost(transaction)?;
+
+        // calculate account access cost
         let message = transaction.message();
         message.account_keys.iter().enumerate().for_each(|(i, k)| {
             let is_signer = message.is_signer(i);
@@ -135,12 +155,11 @@ impl CostModel {
                     NON_SIGNED_READONLY_ACCOUNT_ACCESS_COST;
             }
         });
-        self.transaction_cost.execution_cost = self.find_transaction_cost(transaction);
         debug!(
             "transaction {:?} has cost {:?}",
             transaction, self.transaction_cost
         );
-        &self.transaction_cost
+        Ok(&self.transaction_cost)
     }
 
     // To update or insert instruction cost to table.
@@ -175,10 +194,15 @@ impl CostModel {
         }
     }
 
-    fn find_transaction_cost(&self, transaction: &Transaction) -> u64 {
+    fn find_transaction_cost(&self, transaction: &Transaction) -> Result<u64, CostModelError> {
         let mut cost: u64 = 0;
 
         for instruction in &transaction.message().instructions {
+            // The Transaction may not be sanitized at this point
+            if instruction.program_id_index as usize >= transaction.message().account_keys.len() {
+                return Err(CostModelError::InvalidTransaction);
+            }
+
             let program_id =
                 transaction.message().account_keys[instruction.program_id_index as usize];
             let instruction_cost = self.find_instruction_cost(&program_id);
@@ -189,7 +213,7 @@ impl CostModel {
             );
             cost += instruction_cost;
         }
-        cost
+        Ok(cost)
     }
 }
 
@@ -271,7 +295,7 @@ mod tests {
             .unwrap();
         assert_eq!(
             expected_cost,
-            testee.find_transaction_cost(&simple_transaction)
+            testee.find_transaction_cost(&simple_transaction).unwrap()
         );
     }
 
@@ -295,7 +319,7 @@ mod tests {
         testee
             .upsert_instruction_cost(&system_program::id(), program_cost)
             .unwrap();
-        assert_eq!(expected_cost, testee.find_transaction_cost(&tx));
+        assert_eq!(expected_cost, testee.find_transaction_cost(&tx).unwrap());
     }
 
     #[test]
@@ -321,7 +345,7 @@ mod tests {
         debug!("many random transaction {:?}", tx);
 
         let testee = CostModel::default();
-        let result = testee.find_transaction_cost(&tx);
+        let result = testee.find_transaction_cost(&tx).unwrap();
 
         // expected cost for two random/unknown program is
         let expected_cost = testee.instruction_execution_cost_table.get_mode() * 2;
@@ -350,7 +374,7 @@ mod tests {
         );
 
         let mut cost_model = CostModel::default();
-        let tx_cost = cost_model.calculate_cost(&tx);
+        let tx_cost = cost_model.calculate_cost(&tx).unwrap();
         assert_eq!(2 + 2, tx_cost.writable_accounts.len());
         assert_eq!(signer1.pubkey(), tx_cost.writable_accounts[0]);
         assert_eq!(signer2.pubkey(), tx_cost.writable_accounts[1]);
@@ -392,7 +416,7 @@ mod tests {
         cost_model
             .upsert_instruction_cost(&system_program::id(), expected_execution_cost)
             .unwrap();
-        let tx_cost = cost_model.calculate_cost(&tx);
+        let tx_cost = cost_model.calculate_cost(&tx).unwrap();
         assert_eq!(expected_account_cost, tx_cost.account_access_cost);
         assert_eq!(expected_execution_cost, tx_cost.execution_cost);
         assert_eq!(2, tx_cost.writable_accounts.len());
@@ -460,7 +484,7 @@ mod tests {
                 } else {
                     thread::spawn(move || {
                         let mut cost_model = cost_model.write().unwrap();
-                        let tx_cost = cost_model.calculate_cost(&tx);
+                        let tx_cost = cost_model.calculate_cost(&tx).unwrap();
                         assert_eq!(3, tx_cost.writable_accounts.len());
                         assert_eq!(expected_account_cost, tx_cost.account_access_cost);
                     })

core/src/cost_tracker.rs

@@ -4,7 +4,7 @@
 //! - would_transaction_fit(&tx), immutable function to test if `tx` would fit into current block
 //! - add_transaction_cost(&tx), mutable function to accumulate `tx` cost to tracker.
 //!
-use crate::cost_model::{CostModel, TransactionCost};
+use crate::cost_model::{CostModel, CostModelError, TransactionCost};
 use solana_sdk::{clock::Slot, pubkey::Pubkey, transaction::Transaction};
 use std::{
     collections::HashMap,
@@ -43,18 +43,21 @@ impl CostTracker {
         }
     }
 
-    pub fn would_transaction_fit(&self, transaction: &Transaction) -> Result<(), &'static str> {
+    pub fn would_transaction_fit(&self, transaction: &Transaction) -> Result<(), CostModelError> {
         let mut cost_model = self.cost_model.write().unwrap();
-        let tx_cost = cost_model.calculate_cost(transaction);
+        let tx_cost = cost_model.calculate_cost(transaction)?;
         self.would_fit(
             &tx_cost.writable_accounts,
             &(tx_cost.account_access_cost + tx_cost.execution_cost),
         )
     }
 
-    pub fn add_transaction_cost(&mut self, transaction: &Transaction) {
+    pub fn add_transaction_cost(
+        &mut self,
+        transaction: &Transaction,
+    ) -> Result<(), CostModelError> {
         let mut cost_model = self.cost_model.write().unwrap();
-        let tx_cost = cost_model.calculate_cost(transaction);
+        let tx_cost = cost_model.calculate_cost(transaction)?;
         let cost = tx_cost.account_access_cost + tx_cost.execution_cost;
         for account_key in tx_cost.writable_accounts.iter() {
             *self
@@ -63,6 +66,7 @@ impl CostTracker {
                 .or_insert(0) += cost;
         }
         self.block_cost += cost;
+        Ok(())
     }
 
     pub fn reset_if_new_bank(&mut self, slot: Slot) {
@@ -73,7 +77,7 @@ impl CostTracker {
         }
     }
 
-    pub fn try_add(&mut self, transaction_cost: &TransactionCost) -> Result<u64, &'static str> {
+    pub fn try_add(&mut self, transaction_cost: &TransactionCost) -> Result<u64, CostModelError> {
         let cost = transaction_cost.account_access_cost + transaction_cost.execution_cost;
         self.would_fit(&transaction_cost.writable_accounts, &cost)?;
 
@@ -81,15 +85,15 @@ impl CostTracker {
         Ok(self.block_cost)
     }
 
-    fn would_fit(&self, keys: &[Pubkey], cost: &u64) -> Result<(), &'static str> {
+    fn would_fit(&self, keys: &[Pubkey], cost: &u64) -> Result<(), CostModelError> {
         // check against the total package cost
         if self.block_cost + cost > self.block_cost_limit {
-            return Err("would exceed block cost limit");
+            return Err(CostModelError::WouldExceedBlockMaxLimit);
         }
 
         // check if the transaction itself is more costly than the account_cost_limit
         if *cost > self.account_cost_limit {
-            return Err("Transaction is too expansive, exceeds account cost limit");
+            return Err(CostModelError::WouldExceedAccountMaxLimit);
         }
 
         // check each account against account_cost_limit,
@@ -97,7 +101,7 @@ impl CostTracker {
             match self.cost_by_writable_accounts.get(account_key) {
                 Some(chained_cost) => {
                     if chained_cost + cost > self.account_cost_limit {
-                        return Err("would exceed account cost limit");
+                        return Err(CostModelError::WouldExceedAccountMaxLimit);
                     } else {
                         continue;
                     }

ledger-tool/src/main.rs

@@ -751,7 +751,16 @@ fn compute_slot_cost(blockstore: &Blockstore, slot: Slot) -> Result<(), String>
         let mut cost_model = cost_model.write().unwrap();
         for transaction in &entry.transactions {
             programs += transaction.message().instructions.len();
-            let tx_cost = cost_model.calculate_cost(transaction);
+            let tx_cost = match cost_model.calculate_cost(transaction) {
+                Err(err) => {
+                    warn!(
+                        "failed to calculate transaction cost, err {:?}, tx {:?}",
+                        err, transaction
+                    );
+                    continue;
+                }
+                Ok(cost) => cost,
+            };
             if cost_tracker.try_add(tx_cost).is_err() {
                 println!(
                     "Slot: {}, CostModel rejected transaction {:?}, stats {:?}!",