From e3f0993b133d152bd32cbc98cd612643a198d859 Mon Sep 17 00:00:00 2001
From: 0xfourzerofour <joshpauline@gmail.com>
Date: Fri, 8 Sep 2023 16:18:18 -0400
Subject: [PATCH] feat(tracer): fix one more of the tests

---
 src/common/simulation.rs       | 21 ++++++++++++
 src/common/tracer.rs           | 11 ++++++-
 tracer/src/validationTracer.ts | 60 ++++++++++++++++++++++++++++++----
 3 files changed, 85 insertions(+), 7 deletions(-)

diff --git a/src/common/simulation.rs b/src/common/simulation.rs
index 4a0c4ca8a..1eae161ad 100644
--- a/src/common/simulation.rs
+++ b/src/common/simulation.rs
@@ -238,6 +238,27 @@ where
                     ViolationOpCode(opcode),
                 ));
             }
+
+            for (addr, opcode) in &phase.ext_code_access_info {
+                if *addr == self.entry_point.address() {
+                    violations.push(SimulationViolation::UsedForbiddenOpcode(
+                        entity,
+                        *addr,
+                        ViolationOpCode(*opcode),
+                    ));
+                }
+            }
+
+            for (addr, size) in &phase.contract_size {
+                if *addr != context.entity_infos.sender_address() && size.contract_size <= 2 {
+                    violations.push(SimulationViolation::UsedForbiddenOpcode(
+                        entity,
+                        *addr,
+                        ViolationOpCode(size.opcode),
+                    ));
+                }
+            }
+
             for precompile in &phase.forbidden_precompiles_used {
                 let (contract, precompile) = parse_combined_tracer_str(precompile)?;
                 violations.push(SimulationViolation::UsedForbiddenPrecompile(
diff --git a/src/common/tracer.rs b/src/common/tracer.rs
index ab6cf0b2c..87420094b 100644
--- a/src/common/tracer.rs
+++ b/src/common/tracer.rs
@@ -5,7 +5,7 @@ use std::{
 };
 
 use anyhow::Context;
-use ethers::types::{transaction::eip2718::TypedTransaction, Address, BlockId, U256};
+use ethers::types::{transaction::eip2718::TypedTransaction, Address, BlockId, Opcode, U256};
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use super::{
@@ -36,6 +36,8 @@ pub struct Phase {
     pub called_non_entry_point_with_value: bool,
     pub ran_out_of_gas: bool,
     pub undeployed_contract_accesses: Vec<Address>,
+    pub contract_size: HashMap<Address, ContractSize>,
+    pub ext_code_access_info: HashMap<Address, Opcode>,
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
@@ -45,6 +47,13 @@ pub struct StorageAccess {
     pub slots: Vec<U256>,
 }
 
+#[derive(Clone, Debug, Deserialize, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ContractSize {
+    pub contract_size: u64,
+    pub opcode: Opcode,
+}
+
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct AssociatedSlotsByAddress(HashMap<Address, BTreeSet<U256>>);
 
diff --git a/tracer/src/validationTracer.ts b/tracer/src/validationTracer.ts
index 117c9d216..77a4b6fff 100644
--- a/tracer/src/validationTracer.ts
+++ b/tracer/src/validationTracer.ts
@@ -23,6 +23,8 @@ interface Phase {
   calledNonEntryPointWithValue: boolean;
   ranOutOfGas: boolean;
   undeployedContractAccesses: string[];
+  contractSize: Record<string, ContractSize>;
+  extCodeAccessInfo: { [addr: string]: string }
 }
 
 interface StorageAccess {
@@ -30,6 +32,15 @@ interface StorageAccess {
   slots: string[];
 }
 
+interface ContractSize {
+  contractSize: number;
+  opcode: string;
+}
+
+interface RelevantStepData {
+  opcode: string
+  stackTop3: any[]
+}
 type InternalPhase = Omit<
   Phase,
   | "forbiddenOpcodesUsed"
@@ -37,12 +48,14 @@ type InternalPhase = Omit<
   | "storageAccesses"
   | "addressesCallingWithValue"
   | "undeployedContractAccesses"
+  | "contractSize"
 > & {
   forbiddenOpcodesUsed: StringSet;
   forbiddenPrecompilesUsed: StringSet;
   storageAccesses: Record<string, StringSet>;
   addressesCallingWithValue: StringSet;
   undeployedContractAccesses: StringSet;
+  contractSize: Record<string, ContractSize>;
 };
 
 type StringSet = Record<string, boolean | undefined>;
@@ -108,7 +121,7 @@ type StringSet = Record<string, boolean | undefined>;
   let entryPointAddress = "";
   let justCalledGas = false;
   let pendingKeccakAddress = "";
-  let lastThreeOpcodes: string[] = [];
+  let lastThreeOpcodes: RelevantStepData[] = [];
 
   function newInternalPhase(): InternalPhase {
     return {
@@ -120,6 +133,9 @@ type StringSet = Record<string, boolean | undefined>;
       calledNonEntryPointWithValue: false,
       ranOutOfGas: false,
       undeployedContractAccesses: {},
+      contractSize: {},
+      extCodeAccessInfo: {},
+
     };
   }
 
@@ -128,6 +144,8 @@ type StringSet = Record<string, boolean | undefined>;
       calledBannedEntryPointMethod,
       calledNonEntryPointWithValue,
       ranOutOfGas,
+      contractSize,
+      extCodeAccessInfo,
     } = currentPhase;
     const forbiddenOpcodesUsed = Object.keys(currentPhase.forbiddenOpcodesUsed);
     const forbiddenPrecompilesUsed = Object.keys(
@@ -154,10 +172,11 @@ type StringSet = Record<string, boolean | undefined>;
       calledNonEntryPointWithValue,
       ranOutOfGas,
       undeployedContractAccesses,
+      contractSize,
+      extCodeAccessInfo,
     };
     phases.push(phase);
     currentPhase = newInternalPhase();
-    lastThreeOpcodes = [];
   }
 
   function bigIntToNumber(n: BigInt): number {
@@ -237,22 +256,30 @@ type StringSet = Record<string, boolean | undefined>;
         )[keccakResult] = true;
         pendingKeccakAddress = "";
       }
-      const opcode = log.op.toString();
 
-      lastThreeOpcodes.push(opcode);
+      const opcode = log.op.toString()
+
+      const stackSize = log.stack.length()
+      const stackTop3 = []
+      for (let i = 0; i < 3 && i < stackSize; i++) {
+        stackTop3.push(log.stack.peek(i))
+      }
+
+      lastThreeOpcodes.push({ opcode, stackTop3 })
 
       if (lastThreeOpcodes.length > 3) {
-          lastThreeOpcodes.shift();
+        lastThreeOpcodes.shift()
       }
 
       const entryPointIsExecuting = log.getDepth() === 1;
       if (entryPointIsExecuting) {
         if (opcode === "NUMBER") {
           concludePhase();
-        } else if (opcode === "REVERT") {
+        } else if (opcode === "REVERT" || opcode == "RETURN") {
           const offset = bigIntToNumber(log.stack.peek(0));
           const length = bigIntToNumber(log.stack.peek(1));
           revertData = toHex(log.memory.slice(offset, offset + length));
+          lastThreeOpcodes = [];
         }
       } else {
         // The entry point is allowed to freely call `GAS`, but otherwise we
@@ -272,6 +299,19 @@ type StringSet = Record<string, boolean | undefined>;
     
         }
       }
+
+      const lastOpInfo = lastThreeOpcodes[lastThreeOpcodes.length - 2]
+      // store all addresses touched by EXTCODE* opcodes
+      if (lastOpInfo?.opcode?.match(/^(EXT.*)$/) != null) {
+        const addr = toAddress(lastOpInfo.stackTop3[0].toString(16))
+        const addrHex = toHex(addr)
+        const last3opcodesString = lastThreeOpcodes.map(x => x.opcode).join(' ')
+        // only store the last EXTCODE* opcode per address - could even be a boolean for our current use-case
+        if (last3opcodesString.match(/^(\w+) EXTCODESIZE ISZERO$/) == null) {
+          currentPhase.extCodeAccessInfo[addrHex] = opcode
+        } 
+      }
+
       if (opcode === "CREATE2") {
         if (phases.length === 0) {
           // In factory phase.
@@ -323,6 +363,14 @@ type StringSet = Record<string, boolean | undefined>;
         const index = EXT_OPCODES[opcode] ? 0 : 1;
         const address = toAddress(log.stack.peek(index).toString(16));
         const addressHex = toHex(address);
+
+        if (currentPhase.contractSize[addressHex] == null && !PRECOMPILE_WHILTELIST[addressHex]) {
+          currentPhase.contractSize[addressHex] = {
+            contractSize: db.getCode(address).length,
+            opcode
+          }
+        }
+
         if (!isPrecompiled(address)) {
           if (
             !accessedContractAddresses[addressHex] ||