From dc4c117b8daccf7aeaffa93df88c626698865c7e Mon Sep 17 00:00:00 2001
From: Tao Zhu <tao@solana.com>
Date: Thu, 18 Apr 2024 16:15:15 +0000
Subject: [PATCH] Fail transaction if requested cu-limit exceeds max limit

---
 .../src/compute_budget_processor.rs           | 19 +++++++++++++++++++
 sdk/src/transaction/error.rs                  |  5 +++++
 storage-proto/proto/transaction_by_addr.proto |  1 +
 storage-proto/src/convert.rs                  | 12 ++++++++++++
 4 files changed, 37 insertions(+)

diff --git a/program-runtime/src/compute_budget_processor.rs b/program-runtime/src/compute_budget_processor.rs
index f87bbcd6c17fd7..2005b031aae6d0 100644
--- a/program-runtime/src/compute_budget_processor.rs
+++ b/program-runtime/src/compute_budget_processor.rs
@@ -120,6 +120,13 @@ pub fn process_compute_budget_instructions<'a>(
         }
     }
 
+    // fail transaction if requesting cu-limit more than max
+    if let Some(user_requested_compute_unit_limit) = updated_compute_unit_limit {
+        if user_requested_compute_unit_limit > MAX_COMPUTE_UNIT_LIMIT {
+            return Err(TransactionError::InvalidComputeUnitLimitRequested);
+        }
+    }
+
     // sanitize limits
     let updated_heap_bytes = requested_heap_size
         .unwrap_or(u32::try_from(MIN_HEAP_FRAME_BYTES).unwrap()) // loader's default heap_size
@@ -497,4 +504,16 @@ mod tests {
             })
         );
     }
+
+    #[test]
+    fn test_process_instructions_invalid_compute_unit_limit() {
+        // larger than MAX_COMPUTE_UNIT_LIMIT
+        test!(
+            &[
+                ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT + 1),
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
+            ],
+            Err(TransactionError::InvalidComputeUnitLimitRequested)
+        );
+    }
 }
diff --git a/sdk/src/transaction/error.rs b/sdk/src/transaction/error.rs
index ffd4e70795e7a8..17ada67a0bb39b 100644
--- a/sdk/src/transaction/error.rs
+++ b/sdk/src/transaction/error.rs
@@ -173,6 +173,11 @@ pub enum TransactionError {
     /// Program cache hit max limit.
     #[error("Program cache hit max limit")]
     ProgramCacheHitMaxLimit,
+
+    /// The requested compute-unit limit is either lesser than minimally required by the
+    /// transaction, or greater than maximal transaction compute-unit limit
+    #[error("Requested compute-unit limit is invalid")]
+    InvalidComputeUnitLimitRequested,
 }
 
 impl From<SanitizeError> for TransactionError {
diff --git a/storage-proto/proto/transaction_by_addr.proto b/storage-proto/proto/transaction_by_addr.proto
index d0fa74a2104707..3c97cc8f631d0c 100644
--- a/storage-proto/proto/transaction_by_addr.proto
+++ b/storage-proto/proto/transaction_by_addr.proto
@@ -63,6 +63,7 @@ enum TransactionErrorType {
     PROGRAM_EXECUTION_TEMPORARILY_RESTRICTED = 35;
     UNBALANCED_TRANSACTION = 36;
     PROGRAM_CACHE_HIT_MAX_LIMIT = 37;
+    INVALID_COMPUTE_UNIT_LIMIT_REQUESTED = 38;
 }
 
 message InstructionError {
diff --git a/storage-proto/src/convert.rs b/storage-proto/src/convert.rs
index 754afc8a18475b..731e86391414fe 100644
--- a/storage-proto/src/convert.rs
+++ b/storage-proto/src/convert.rs
@@ -819,6 +819,7 @@ impl TryFrom<tx_by_addr::TransactionError> for TransactionError {
             34 => TransactionError::ResanitizationNeeded,
             36 => TransactionError::UnbalancedTransaction,
             37 => TransactionError::ProgramCacheHitMaxLimit,
+            38 => TransactionError::InvalidComputeUnitLimitRequested,
             _ => return Err("Invalid TransactionError"),
         })
     }
@@ -940,6 +941,9 @@ impl From<TransactionError> for tx_by_addr::TransactionError {
                 TransactionError::ProgramCacheHitMaxLimit => {
                     tx_by_addr::TransactionErrorType::ProgramCacheHitMaxLimit
                 }
+                TransactionError::InvalidComputeUnitLimitRequested => {
+                    tx_by_addr::TransactionErrorType::InvalidComputeUnitLimitRequested
+                }
             } as i32,
             instruction_error: match transaction_error {
                 TransactionError::InstructionError(index, ref instruction_error) => {
@@ -1855,6 +1859,14 @@ mod test {
             transaction_error,
             tx_by_addr_transaction_error.try_into().unwrap()
         );
+
+        let transaction_error = TransactionError::InvalidComputeUnitLimitRequested;
+        let tx_by_addr_transaction_error: tx_by_addr::TransactionError =
+            transaction_error.clone().into();
+        assert_eq!(
+            transaction_error,
+            tx_by_addr_transaction_error.try_into().unwrap()
+        );
     }
 
     #[test]