pr comment, wrap pk

btcstaking-tracker/btcslasher/slasher.go

@@ -3,6 +3,8 @@ package btcslasher
 import (
 	"context"
 	"fmt"
+	"github.com/babylonlabs-io/vigilante/types"
+	"github.com/decred/dcrd/dcrec/secp256k1/v4"
 	"sync"
 	"time"
 
@@ -233,8 +235,8 @@ func (bs *BTCSlasher) equivocationTracker() {
 // SlashFinalityProvider slashes all BTC delegations under a given finality provider
 // the checkBTC option indicates whether to check the slashing tx's input is still spendable
 // on Bitcoin (including mempool txs).
-func (bs *BTCSlasher) SlashFinalityProvider(extractedfpBTCSK *btcec.PrivateKey) error {
-	fpBTCPK := bbn.NewBIP340PubKeyFromBTCPK(extractedfpBTCSK.PubKey())
+func (bs *BTCSlasher) SlashFinalityProvider(extractedFpBTCSK *btcec.PrivateKey) error {
+	fpBTCPK := bbn.NewBIP340PubKeyFromBTCPK(extractedFpBTCSK.PubKey())
 	bs.logger.Infof("start slashing finality provider %s", fpBTCPK.MarshalHex())
 
 	// get all active and unbonded BTC delegations at the current BTC height
@@ -246,8 +248,8 @@ func (bs *BTCSlasher) SlashFinalityProvider(extractedfpBTCSK *btcec.PrivateKey)
 		return fmt.Errorf("failed to get BTC delegations under finality provider %s: %w", fpBTCPK.MarshalHex(), err)
 	}
 
-	// Initialize a mutex to protect access to extractedfpBTCSK *btcec.PrivateKey
-	var mu sync.Mutex
+	// Initialize a mutex protected *btcec.PrivateKey
+	safeExtractedFpBTCSK := types.NewPrivateKeyWithMutex(extractedFpBTCSK)
 
 	// try to slash both staking and unbonding txs for each BTC delegation
 	// sign and submit slashing tx for each active delegation
@@ -256,9 +258,9 @@ func (bs *BTCSlasher) SlashFinalityProvider(extractedfpBTCSK *btcec.PrivateKey)
 		bs.wg.Add(1)
 		go func(d *bstypes.BTCDelegationResponse) {
 			defer bs.wg.Done()
-			mu.Lock()
-			defer mu.Unlock()
-			bs.slashBTCDelegation(fpBTCPK, extractedfpBTCSK, d)
+			safeExtractedFpBTCSK.UseKey(func(key *secp256k1.PrivateKey) {
+				bs.slashBTCDelegation(fpBTCPK, key, d)
+			})
 		}(del)
 	}
 	// sign and submit slashing tx for each unbonded delegation
@@ -267,9 +269,9 @@ func (bs *BTCSlasher) SlashFinalityProvider(extractedfpBTCSK *btcec.PrivateKey)
 		bs.wg.Add(1)
 		go func(d *bstypes.BTCDelegationResponse) {
 			defer bs.wg.Done()
-			mu.Lock()
-			defer mu.Unlock()
-			bs.slashBTCDelegation(fpBTCPK, extractedfpBTCSK, d)
+			safeExtractedFpBTCSK.UseKey(func(key *secp256k1.PrivateKey) {
+				bs.slashBTCDelegation(fpBTCPK, key, d)
+			})
 		}(del)
 	}
 

types/safeprivatekey.go

@@ -0,0 +1,33 @@
+package types
+
+import (
+	"github.com/decred/dcrd/dcrec/secp256k1/v4"
+	"sync"
+)
+
+// PrivateKeyWithMutex wraps a btcec.PrivateKey with a mutex to ensure thread-safe access.
+type PrivateKeyWithMutex struct {
+	mu  sync.Mutex
+	key *secp256k1.PrivateKey
+}
+
+// NewPrivateKeyWithMutex creates a new PrivateKeyWithMutex.
+func NewPrivateKeyWithMutex(key *secp256k1.PrivateKey) *PrivateKeyWithMutex {
+	return &PrivateKeyWithMutex{
+		key: key,
+	}
+}
+
+// GetKey safely retrieves the private key.
+func (p *PrivateKeyWithMutex) GetKey() *secp256k1.PrivateKey {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.key
+}
+
+// UseKey performs an operation with the private key in a thread-safe manner.
+func (p *PrivateKeyWithMutex) UseKey(operation func(key *secp256k1.PrivateKey)) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	operation(p.key)
+}