From 70be670ddda70df7a496232edcc2b6a84b211aad Mon Sep 17 00:00:00 2001
From: Tao Zhu <tao@solana.com>
Date: Mon, 13 Mar 2023 21:38:22 +0000
Subject: [PATCH] checked result of consume_checked() if featureis activated

---
 programs/bpf_loader/src/lib.rs | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/programs/bpf_loader/src/lib.rs b/programs/bpf_loader/src/lib.rs
index 604b33a82ed0f1..58ce282989bf3b 100644
--- a/programs/bpf_loader/src/lib.rs
+++ b/programs/bpf_loader/src/lib.rs
@@ -321,13 +321,17 @@ pub fn create_ebpf_vm<'a, 'b>(
     orig_account_lengths: Vec<usize>,
     invoke_context: &'a mut InvokeContext<'b>,
 ) -> Result<EbpfVm<'a, RequisiteVerifier, InvokeContext<'b>>, EbpfError> {
-    let _ = invoke_context.consume_checked(calculate_heap_cost(
+    let round_up_heap_size = invoke_context
+        .feature_set
+        .is_active(&round_up_heap_size::id());
+    let heap_cost_result = invoke_context.consume_checked(calculate_heap_cost(
         heap.len() as u64,
         invoke_context.get_compute_budget().heap_cost,
-        invoke_context
-            .feature_set
-            .is_active(&round_up_heap_size::id()),
+        round_up_heap_size,
     ));
+    if round_up_heap_size {
+        heap_cost_result?;
+    }
     let check_aligned = bpf_loader_deprecated::id()
         != invoke_context
             .transaction_context