From 334dee684474b51dfd50ebe159bcc3f592400040 Mon Sep 17 00:00:00 2001
From: hazim-j <hazim.jumali@gmail.com>
Date: Fri, 4 Aug 2023 00:18:05 +1000
Subject: [PATCH] Add mutex locks to prevent db transaction conflicts

---
 pkg/modules/paymaster/reputation.go | 11 ++++++++++-
 pkg/modules/relay/relayer.go        |  7 +++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/pkg/modules/paymaster/reputation.go b/pkg/modules/paymaster/reputation.go
index 2e734431..ec0dcb34 100644
--- a/pkg/modules/paymaster/reputation.go
+++ b/pkg/modules/paymaster/reputation.go
@@ -4,6 +4,7 @@ package paymaster
 
 import (
 	"errors"
+	"sync"
 
 	mapset "github.com/deckarep/golang-set/v2"
 	"github.com/dgraph-io/badger/v3"
@@ -15,12 +16,14 @@ import (
 // UserOperation.
 type Reputation struct {
 	db *badger.DB
+	mu *sync.Mutex
 }
 
 // New returns an instance of a Reputation object to track and appropriately process userOps by paymaster
 // status.
 func New(db *badger.DB) *Reputation {
-	return &Reputation{db}
+	mu := &sync.Mutex{}
+	return &Reputation{db, mu}
 }
 
 // CheckStatus returns a UserOpHandler that is used by the Client to determine if the userOp is allowed based
@@ -37,7 +40,9 @@ func (r *Reputation) CheckStatus() modules.UserOpHandlerFunc {
 				return nil
 			}
 
+			r.mu.Lock()
 			status, err := getStatus(txn, paymaster)
+			r.mu.Unlock()
 			if err != nil {
 				return err
 			}
@@ -62,6 +67,8 @@ func (r *Reputation) IncOpsSeen() modules.UserOpHandlerFunc {
 				return nil
 			}
 
+			r.mu.Lock()
+			defer r.mu.Unlock()
 			return incrementOpsSeenByPaymaster(txn, paymaster)
 		})
 	}
@@ -87,6 +94,8 @@ func (r *Reputation) IncOpsIncluded() modules.BatchHandlerFunc {
 				}
 			}
 
+			r.mu.Lock()
+			defer r.mu.Unlock()
 			return incrementOpsIncludedByPaymasters(txn, c, ps.ToSlice()...)
 		})
 	}
diff --git a/pkg/modules/relay/relayer.go b/pkg/modules/relay/relayer.go
index eb40607f..278680e2 100644
--- a/pkg/modules/relay/relayer.go
+++ b/pkg/modules/relay/relayer.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"math/big"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/dgraph-io/badger/v3"
@@ -45,6 +46,7 @@ type Relayer struct {
 	bannedThreshold  int
 	bannedTimeWindow time.Duration
 	waitTimeout      time.Duration
+	mu               *sync.Mutex
 }
 
 // New initializes a new EOA relayer for sending batches to the EntryPoint with IP throttling protection.
@@ -66,6 +68,7 @@ func New(
 		bannedThreshold:  DefaultBanThreshold,
 		bannedTimeWindow: DefaultBanTimeWindow,
 		waitTimeout:      DefaultWaitTimeout,
+		mu:               &sync.Mutex{},
 	}
 }
 
@@ -185,6 +188,8 @@ func (r *Relayer) MapUserOpHashToClientID() gin.HandlerFunc {
 				return err
 			}
 
+			r.mu.Lock()
+			defer r.mu.Unlock()
 			return incrementOpsSeenByClientID(txn, cid, r.bannedTimeWindow)
 		})
 		if err != nil {
@@ -260,6 +265,8 @@ func (r *Relayer) SendUserOperation() modules.BatchHandlerFunc {
 
 			hashes := getUserOpHashesFromOps(ctx.EntryPoint, ctx.ChainID, ctx.Batch...)
 			del = append([]string{}, hashes...)
+			r.mu.Lock()
+			defer r.mu.Unlock()
 			return incrementOpsIncludedByUserOpHashes(txn, r.bannedTimeWindow, hashes...)
 		})
 		if err != nil {