From 6d5527518d0b4e6369f0c6c12895231ae8b0be61 Mon Sep 17 00:00:00 2001
From: y0sher <lyosher@gmail.com>
Date: Wed, 28 Aug 2024 20:51:59 +0300
Subject: [PATCH] fix: return lock instead of locking inside `lock` function to
 prevent deadlock.

---
 signer/sign_attestation.go    |  5 +++--
 signer/sign_block.go          |  5 +++--
 signer/sign_sync_committee.go | 15 +++++++++------
 signer/validator_signer.go    | 16 +++-------------
 4 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/signer/sign_attestation.go b/signer/sign_attestation.go
index 7b6ce165..d03603e2 100644
--- a/signer/sign_attestation.go
+++ b/signer/sign_attestation.go
@@ -19,9 +19,10 @@ func (signer *SimpleSigner) SignBeaconAttestation(attestation *phase0.Attestatio
 	}
 
 	// 2. lock for current account
-	signer.lock(account.ID(), "attestation")
+	val := signer.lock(account.ID(), "attestation")
+	val.Lock()
 	defer func() {
-		signer.unlock(account.ID(), "attestation")
+		val.Unlock()
 	}()
 
 	// 3. far future check
diff --git a/signer/sign_block.go b/signer/sign_block.go
index 1817d188..616a9369 100644
--- a/signer/sign_block.go
+++ b/signer/sign_block.go
@@ -23,8 +23,9 @@ func (signer *SimpleSigner) SignBlock(block ssz.HashRoot, slot phase0.Slot, doma
 	}
 
 	// 2. lock for current account
-	signer.lock(account.ID(), "proposal")
-	defer signer.unlock(account.ID(), "proposal")
+	val := signer.lock(account.ID(), "proposal")
+	val.Lock()
+	defer val.Unlock()
 
 	// 3. far future check
 	if !IsValidFarFutureSlot(signer.network, slot) {
diff --git a/signer/sign_sync_committee.go b/signer/sign_sync_committee.go
index 7f2b0492..1d670045 100644
--- a/signer/sign_sync_committee.go
+++ b/signer/sign_sync_committee.go
@@ -21,8 +21,9 @@ func (signer *SimpleSigner) SignSyncCommittee(msgBlockRoot []byte, domain phase0
 	}
 
 	// 2. lock for current account
-	signer.lock(account.ID(), "sync_committee")
-	defer signer.unlock(account.ID(), "sync_committee")
+	val := signer.lock(account.ID(), "sync_committee")
+	val.Lock()
+	defer val.Unlock()
 
 	// 3. sign
 	sszRoot := SSZBytes(msgBlockRoot)
@@ -51,8 +52,9 @@ func (signer *SimpleSigner) SignSyncCommitteeSelectionData(data *altair.SyncAggr
 	}
 
 	// 2. lock for current account
-	signer.lock(account.ID(), "sync_committee_selection_data")
-	defer signer.unlock(account.ID(), "sync_committee_selection_data")
+	val := signer.lock(account.ID(), "sync_committee_selection_data")
+	val.Lock()
+	defer val.Unlock()
 
 	// 3. sign
 	if data == nil {
@@ -83,8 +85,9 @@ func (signer *SimpleSigner) SignSyncCommitteeContributionAndProof(contribAndProo
 	}
 
 	// 2. lock for current account
-	signer.lock(account.ID(), "sync_committee_selection_and_proof")
-	defer signer.unlock(account.ID(), "sync_committee_selection_and_proof")
+	val := signer.lock(account.ID(), "sync_committee_selection_and_proof")
+	val.Lock()
+	defer val.Unlock()
 
 	// 3. sign
 	if contribAndProof == nil {
diff --git a/signer/validator_signer.go b/signer/validator_signer.go
index a4e12b9c..b5e0d7c5 100644
--- a/signer/validator_signer.go
+++ b/signer/validator_signer.go
@@ -51,26 +51,16 @@ func NewSimpleSigner(wallet core.Wallet, slashingProtector core.SlashingProtecto
 }
 
 // lock locks signer
-func (signer *SimpleSigner) lock(accountID uuid.UUID, operation string) {
+func (signer *SimpleSigner) lock(accountID uuid.UUID, operation string) *sync.RWMutex {
 	signer.mapLock.Lock()
 	defer signer.mapLock.Unlock()
 
 	k := accountID.String() + "_" + operation
 	if val, ok := signer.signLocks[k]; ok {
-		val.Lock()
+		return val
 	} else {
 		signer.signLocks[k] = &sync.RWMutex{}
-		signer.signLocks[k].Lock()
-	}
-}
-
-func (signer *SimpleSigner) unlock(accountID uuid.UUID, operation string) {
-	signer.mapLock.RLock()
-	defer signer.mapLock.RUnlock()
-
-	k := accountID.String() + "_" + operation
-	if val, ok := signer.signLocks[k]; ok {
-		val.Unlock()
+		return signer.signLocks[k]
 	}
 }