Fix assertion tx env gas fields

crates/cheatcodes/src/credible.rs

@@ -103,6 +103,7 @@ pub struct TxAttributes {
     pub data: Bytes,
     pub caller: Address,
     pub kind: TxKind,
+    pub gas_limit: u64,
 }
 
 /// Maximum gas allowed for assertion execution (300k gas).
@@ -123,6 +124,24 @@ fn check_assertion_gas_limit(gas_used: u64) -> Option<String> {
     }
 }
 
+fn build_tx_env(
+    tx_attributes: TxAttributes,
+    base_tx_env: &TxEnv,
+    chain_id: u64,
+    nonce: u64,
+) -> TxEnv {
+    let tx_gas_limit = tx_attributes.gas_limit.min(TX_GAS_LIMIT_CAP);
+    let mut tx_env = base_tx_env.clone();
+    tx_env.caller = tx_attributes.caller;
+    tx_env.gas_limit = tx_gas_limit;
+    tx_env.chain_id = Some(chain_id);
+    tx_env.value = tx_attributes.value;
+    tx_env.data = tx_attributes.data;
+    tx_env.kind = tx_attributes.kind;
+    tx_env.nonce = nonce;
+    tx_env
+}
+
 /// Used to handle assertion execution in inspector in calls after the cheatcode was called.
 pub fn execute_assertion(
     assertion: &Assertion,
@@ -135,10 +154,12 @@ pub fn execute_assertion(
     let block = ecx.block.clone();
     let state = ecx.journaled_state.state.clone();
     let chain_id = ecx.cfg.chain_id;
+    let base_tx_env = ecx.tx.clone();
 
-    let (db, journal, _) = ecx.as_db_env_and_journal();
-    let nonce =
-        journal.load_account(db, tx_attributes.caller).map(|acc| acc.info.nonce).unwrap_or(0);
+    let nonce = {
+        let (db, journal, _) = ecx.as_db_env_and_journal();
+        journal.load_account(db, tx_attributes.caller).map(|acc| acc.info.nonce).unwrap_or(0)
+    };
     // Setup assertion database
     let db = ThreadSafeDb::new(*ecx.db_mut());
 
@@ -166,19 +187,8 @@ pub fn execute_assertion(
     });
 
     store.insert(assertion.adopter, assertion_state).expect("Failed to store assertions");
-    // transaction gas limit should respect new hardfork rules of max 16m gas
-    let tx_gas_limit = block.gas_limit.min(TX_GAS_LIMIT_CAP);
-    let tx_env = TxEnv {
-        caller: tx_attributes.caller,
-        gas_limit: tx_gas_limit,
-        gas_price: block.basefee.into(),
-        chain_id: Some(chain_id),
-        value: tx_attributes.value,
-        data: tx_attributes.data,
-        kind: tx_attributes.kind,
-        nonce,
-        ..Default::default()
-    };
+
+    let tx_env = build_tx_env(tx_attributes, &base_tx_env, chain_id, nonce);
 
     let mut assertion_executor = config.build(store);
 
@@ -358,6 +368,55 @@ mod tests {
         assert_eq!(decoded, "Halt reason: CallTooDeep");
     }
 
+    #[test]
+    fn test_build_tx_env_uses_base_gas_fields() {
+        let base_tx_env = TxEnv {
+            gas_limit: 50_000,
+            gas_price: 123,
+            gas_priority_fee: Some(7),
+            tx_type: 2,
+            ..Default::default()
+        };
+        let tx_attributes = TxAttributes {
+            value: U256::from(1),
+            data: Bytes::from(vec![0x01, 0x02]),
+            caller: Address::from([0x11; 20]),
+            kind: TxKind::Call(Address::from([0x22; 20])),
+            gas_limit: 40_000,
+        };
+        let tx_env = build_tx_env(tx_attributes, &base_tx_env, 1, 9);
+
+        assert_eq!(tx_env.gas_limit, 40_000);
+        assert_eq!(tx_env.gas_price, 123);
+        assert_eq!(tx_env.gas_priority_fee, Some(7));
+        assert_eq!(tx_env.tx_type, 2);
+        assert_eq!(tx_env.chain_id, Some(1));
+        assert_eq!(tx_env.nonce, 9);
+        assert_eq!(tx_env.caller, Address::from([0x11; 20]));
+        assert_eq!(tx_env.value, U256::from(1));
+        assert_eq!(tx_env.data, Bytes::from(vec![0x01, 0x02]));
+        assert_eq!(tx_env.kind, TxKind::Call(Address::from([0x22; 20])));
+    }
+
+    #[test]
+    fn test_build_tx_env_caps_gas_limit() {
+        let base_tx_env = TxEnv {
+            gas_limit: TX_GAS_LIMIT_CAP.saturating_add(1),
+            gas_price: 1,
+            ..Default::default()
+        };
+        let tx_attributes = TxAttributes {
+            value: U256::ZERO,
+            data: Bytes::new(),
+            caller: Address::from([0x33; 20]),
+            kind: TxKind::Call(Address::from([0x44; 20])),
+            gas_limit: TX_GAS_LIMIT_CAP.saturating_add(1),
+        };
+        let tx_env = build_tx_env(tx_attributes, &base_tx_env, 1, 0);
+
+        assert_eq!(tx_env.gas_limit, TX_GAS_LIMIT_CAP);
+    }
+
     #[test]
     fn test_assertion_gas_limit_constant() {
         // Ensure the gas limit is set to the expected value (300k)

crates/cheatcodes/src/inspector.rs

@@ -1051,6 +1051,7 @@ impl Cheatcodes {
                 data: call.input.bytes(ecx),
                 caller: call.caller,
                 kind: TxKind::Call(call.target_address),
+                gas_limit: call.gas_limit,
             };
 
             return match crate::credible::execute_assertion(
@@ -1790,6 +1791,7 @@ impl Inspector<EthEvmContext<&mut dyn DatabaseExt>> for Cheatcodes {
                 data: input.init_code(),
                 caller: input.caller(),
                 kind: TxKind::Create,
+                gas_limit: input.gas_limit(),
             };
 
             return match crate::credible::execute_assertion(