From 276de79f4d319c5411d7e78b07a2e535011e91ce Mon Sep 17 00:00:00 2001
From: Alex Forshtat <forshtat1@gmail.com>
Date: Thu, 16 May 2024 16:40:04 +0200
Subject: [PATCH] Add support for the transient storage (TSTORE/TLOAD opcodes)
 (#192)

* Add support for the transient storage (TSTORE/TLOAD opcodes)
---
 .../src/BundlerCollectorTracer.ts             | 16 ++++++++++++---
 .../src/TracerResultParser.ts                 | 20 ++++++++++++++-----
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/packages/validation-manager/src/BundlerCollectorTracer.ts b/packages/validation-manager/src/BundlerCollectorTracer.ts
index cc4d0fe6..4da32c6f 100644
--- a/packages/validation-manager/src/BundlerCollectorTracer.ts
+++ b/packages/validation-manager/src/BundlerCollectorTracer.ts
@@ -75,6 +75,10 @@ export interface AccessInfo {
   reads: { [slot: string]: string }
   // count of writes.
   writes: { [slot: string]: number }
+  // count of transient reads
+  transientReads: { [slot: string]: number }
+  // count of transient writes
+  transientWrites: { [slot: string]: number }
 }
 
 export interface LogInfo {
@@ -291,7 +295,7 @@ export function bundlerCollectorTracer (): BundlerCollectorTracer {
       }
       this.lastOp = opcode
 
-      if (opcode === 'SLOAD' || opcode === 'SSTORE') {
+      if (opcode === 'SLOAD' || opcode === 'SSTORE' || opcode === 'TLOAD' || opcode === 'TSTORE') {
         const slot = toWord(log.stack.peek(0).toString(16))
         const slotHex = toHex(slot)
         const addr = log.contract.getAddress()
@@ -300,7 +304,9 @@ export function bundlerCollectorTracer (): BundlerCollectorTracer {
         if (access == null) {
           access = {
             reads: {},
-            writes: {}
+            writes: {},
+            transientReads: {},
+            transientWrites: {}
           }
           this.currentLevel.access[addrHex] = access
         }
@@ -310,8 +316,12 @@ export function bundlerCollectorTracer (): BundlerCollectorTracer {
           if (access.reads[slotHex] == null && access.writes[slotHex] == null) {
             access.reads[slotHex] = toHex(db.getState(addr, slot))
           }
-        } else {
+        } else if (opcode === 'SSTORE') {
           this.countSlot(access.writes, slotHex)
+        } else if (opcode === 'TLOAD') {
+          this.countSlot(access.transientReads, slotHex)
+        } else if (opcode === 'TSTORE') {
+          this.countSlot(access.transientWrites, slotHex)
         }
       }
 
diff --git a/packages/validation-manager/src/TracerResultParser.ts b/packages/validation-manager/src/TracerResultParser.ts
index b91b838c..1e684bee 100644
--- a/packages/validation-manager/src/TracerResultParser.ts
+++ b/packages/validation-manager/src/TracerResultParser.ts
@@ -255,7 +255,9 @@ export function tracerResultParser (
 
     Object.entries(access).forEach(([addr, {
       reads,
-      writes
+      writes,
+      transientReads,
+      transientWrites
     }]) => {
       // testing read/write access on contract "addr"
       if (addr === sender) {
@@ -305,7 +307,12 @@ export function tracerResultParser (
       // scan all slots. find a referenced slot
       // at the end of the scan, we will check if the entity has stake, and report that slot if not.
       let requireStakeSlot: string | undefined
-      [...Object.keys(writes), ...Object.keys(reads)].forEach(slot => {
+      [
+        ...Object.keys(writes),
+        ...Object.keys(reads),
+        ...Object.keys(transientWrites),
+        ...Object.keys(transientReads)
+      ].forEach(slot => {
         // slot associated with sender is allowed (e.g. token.balanceOf(sender)
         // but during initial UserOp (where there is an initCode), it is allowed only for staked entity
         if (associatedWith(slot, sender, entitySlots)) {
@@ -324,14 +331,17 @@ export function tracerResultParser (
           // [STO-031]
           // accessing storage member of entity itself requires stake.
           requireStakeSlot = slot
-        } else if (writes[slot] == null) {
+        } else if (writes[slot] == null && transientWrites[slot] == null) {
           // [STO-033]: staked entity have read-only access to any storage in non-entity contract.
           requireStakeSlot = slot
         } else {
           // accessing arbitrary storage of another contract is not allowed
-          const readWrite = Object.keys(writes).includes(addr) ? 'write to' : 'read from'
+          const isWrite = Object.keys(writes).includes(slot) || Object.keys(transientWrites).includes(slot)
+          const isTransient = Object.keys(transientReads).includes(slot) || Object.keys(transientWrites).includes(slot)
+          const readWrite = isWrite ? 'write to' : 'read from'
+          const transientStr = isTransient ? 'transient ' : ''
           requireCond(false,
-            `${entityTitle} has forbidden ${readWrite} ${nameAddr(addr, entityTitle)} slot ${slot}`,
+            `${entityTitle} has forbidden ${readWrite} ${transientStr}${nameAddr(addr, entityTitle)} slot ${slot}`,
             ValidationErrors.OpcodeValidation, { [entityTitle]: entStakes?.addr })
         }
       })