From 083936f231f64f6360bbe96d52dcd5a8753ef910 Mon Sep 17 00:00:00 2001
From: Alex Forshtat <forshtat1@gmail.com>
Date: Mon, 2 Oct 2023 16:25:22 +0300
Subject: [PATCH] STO-040 rule check implementation

---
 packages/bundler/src/modules/BundleManager.ts |  4 +-
 .../bundler/src/modules/MempoolManager.ts     | 52 ++++++++++++++++++-
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/packages/bundler/src/modules/BundleManager.ts b/packages/bundler/src/modules/BundleManager.ts
index c03ae5f0..276f4311 100644
--- a/packages/bundler/src/modules/BundleManager.ts
+++ b/packages/bundler/src/modules/BundleManager.ts
@@ -156,9 +156,7 @@ export class BundleManager {
     const senders = new Set<string>()
 
     // all entities that are known to be valid senders in the mempool
-    const knownSenders = entries.map(it => {
-      return it.userOp.sender.toLowerCase()
-    })
+    const knownSenders = this.mempoolManager.getKnownSenders()
 
     const storageMap: StorageMap = {}
     let totalGas = BigNumber.from(0)
diff --git a/packages/bundler/src/modules/MempoolManager.ts b/packages/bundler/src/modules/MempoolManager.ts
index 44e480e1..1b5a0be1 100644
--- a/packages/bundler/src/modules/MempoolManager.ts
+++ b/packages/bundler/src/modules/MempoolManager.ts
@@ -96,8 +96,8 @@ export class MempoolManager {
       if (factory != null) {
         this.incrementEntryCount(factory)
       }
-      // this.checkSenderCountInMempool(userOp, senderInfo)
       this.checkReputation(senderInfo, paymasterInfo, factoryInfo, aggregatorInfo)
+      this.checkMultipleRolesViolation(userOp)
       this.mempool.push(entry)
     }
     this.updateSeenStatus(aggregatorInfo?.addr, userOp, senderInfo)
@@ -137,6 +137,27 @@ export class MempoolManager {
     }
   }
 
+  private checkMultipleRolesViolation (userOp: UserOperation): void {
+    const knownEntities = this.getKnownEntities()
+    requireCond(
+      !knownEntities.includes(userOp.sender.toLowerCase()),
+      `The sender address "${userOp.sender}" is used as a different entity in another UserOperation currently in mempool`,
+      ValidationErrors.OpcodeValidation
+    )
+
+    const knownSenders = this.getKnownSenders()
+    const paymaster = getAddr(userOp.paymasterAndData)?.toLowerCase()
+    const factory = getAddr(userOp.initCode)?.toLowerCase()
+
+    const isPaymasterSenderViolation = knownSenders.includes(paymaster?.toLowerCase() ?? '')
+    const isFactorySenderViolation = knownSenders.includes(factory?.toLowerCase() ?? '')
+    requireCond(
+      !isPaymasterSenderViolation && !isFactorySenderViolation,
+      `An entity in this UserOperation is used as a sender entity in another UserOperation currently in mempool. (${paymaster}:${isPaymasterSenderViolation};${factory}:${isFactorySenderViolation})`,
+      ValidationErrors.OpcodeValidation
+    )
+  }
+
   private checkReputationStatus (
     title: 'account' | 'paymaster' | 'aggregator' | 'deployer',
     stakeInfo: StakeInfo,
@@ -234,4 +255,33 @@ export class MempoolManager {
     this.mempool = []
     this._entryCount = {}
   }
+
+  /**
+   * Returns all addresses that are currently known to be "senders" according to the current mempool.
+   */
+  getKnownSenders (): string[] {
+    return this.getSortedForInclusion().map(it => {
+      return it.userOp.sender.toLowerCase()
+    })
+  }
+
+  /**
+   * Returns all addresses that are currently known to be any kind of entity according to the current mempool.
+   * Note that "sender" addresses are not returned by this function. Use {@link getKnownSenders} instead.
+   */
+  getKnownEntities (): string[] {
+    const res = []
+    const userOps = this.getSortedForInclusion()
+    res.push(
+      ...userOps.map(it => {
+        return getAddr(it.userOp.paymasterAndData)
+      })
+    )
+    res.push(
+      ...userOps.map(it => {
+        return getAddr(it.userOp.initCode)
+      })
+    )
+    return res.filter(it => it != null).map(it => (it as string).toLowerCase())
+  }
 }