From 3d2079788de1d88981c3ebd39a5bfbf0a29364a0 Mon Sep 17 00:00:00 2001
From: Mason Liang <mason@scroll.io>
Date: Mon, 29 Apr 2024 21:59:53 -0400
Subject: [PATCH] Add failing write protection test for TStore

---
 .../execution/error_write_protection.rs       | 69 ++++++++++++-------
 1 file changed, 44 insertions(+), 25 deletions(-)

diff --git a/zkevm-circuits/src/evm_circuit/execution/error_write_protection.rs b/zkevm-circuits/src/evm_circuit/execution/error_write_protection.rs
index 02663549890..06ebc9b3ae5 100644
--- a/zkevm-circuits/src/evm_circuit/execution/error_write_protection.rs
+++ b/zkevm-circuits/src/evm_circuit/execution/error_write_protection.rs
@@ -155,16 +155,25 @@ mod test {
         Account::mock_code_balance(code)
     }
 
+    enum FailureReason {
+        Sstore,
+        TStore,
+        CallWithValue,
+    }
+
     #[test]
     fn test_write_protection() {
-        // test sstore with write protection error
-        test_internal_write_protection(false);
-        // test call with write protection error
-        test_internal_write_protection(true);
+        for reason in [
+            FailureReason::Sstore,
+            FailureReason::CallWithValue,
+            FailureReason::TStore,
+        ] {
+            test_internal_write_protection(reason)
+        }
     }
 
     // ErrorWriteProtection error happen in internal call
-    fn test_internal_write_protection(is_call: bool) {
+    fn test_internal_write_protection(reason: FailureReason) {
         let mut caller_bytecode = bytecode! {
             PUSH1(0)
             PUSH1(0)
@@ -185,26 +194,36 @@ mod test {
             PUSH1(0x02)
         };
 
-        if is_call {
-            callee_bytecode.append(&bytecode! {
-                PUSH1(0)
-                PUSH1(0)
-                PUSH1(10)
-                PUSH1(200)  // non zero value
-                PUSH20(Address::repeat_byte(0xff).to_word())
-                PUSH2(10000)  // gas
-                //this call got error: ErrorWriteProtection
-                CALL
-                RETURN
-                STOP
-            });
-        } else {
-            callee_bytecode.append(&bytecode! {
-                // this SSTORE got error: ErrorWriteProtection
-                SSTORE
-                STOP
-            });
-        }
+        match reason {
+            CallWithValue => {
+                callee_bytecode.append(&bytecode! {
+                    PUSH1(0)
+                    PUSH1(0)
+                    PUSH1(10)
+                    PUSH1(200)  // non zero value
+                    PUSH20(Address::repeat_byte(0xff).to_word())
+                    PUSH2(10000)  // gas
+                    //this call got error: ErrorWriteProtection
+                    CALL
+                    RETURN
+                    STOP
+                });
+            }
+            Sstore => {
+                callee_bytecode.append(&bytecode! {
+                    // this SSTORE got error: ErrorWriteProtection
+                    SSTORE
+                    STOP
+                });
+            }
+            TStore => {
+                callee_bytecode.append(&bytecode! {
+                    // this TSTORE got error: ErrorWriteProtection
+                    TSTORE
+                    STOP
+                });
+            }
+        };
 
         test_ok(
             Account::mock_100_ether(caller_bytecode),