fixed concurrent access

primev · Apr 5, 2024 · 49546a0 · 49546a0
1 parent 620124e
commit 49546a0
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 16 deletions.
diff --git a/pkg/keyexchange/keyexchange.go b/pkg/keyexchange/keyexchange.go
@@ -27,11 +27,11 @@ func New(
 	signer signer.Signer,
 ) *KeyExchange {
 	return &KeyExchange{
-		topo:      topo,
-		streamer:  streamer,
-		keyKeeper: keyKeeper,
-		logger:    logger,
-		signer:    signer,
+		topo:              topo,
+		streamer:          streamer,
+		keyKeeper:         keyKeeper,
+		logger:            logger,
+		signer:            signer,
 	}
 }
 
@@ -195,7 +195,7 @@ func (ke *KeyExchange) handleTimestampMessage(ctx context.Context, peer p2p.Peer
 		return fmt.Errorf("validate and process timestamp failed: %w", err)
 	}
 
-	ke.keyKeeper.(*keykeeper.ProviderKeyKeeper).BiddersAESKeys[peer.EthAddress] = aesKey
+	ke.keyKeeper.(*keykeeper.ProviderKeyKeeper).SetAESKey(peer.EthAddress, aesKey)
 
 	return nil
 }
@@ -245,7 +245,7 @@ func (ke *KeyExchange) decryptMessage(ekmWithSignature *keyexchangepb.EKMWithSig
 		message   keyexchangepb.EncryptedKeysMessage
 	)
 
-    err = proto.Unmarshal(ekmWithSignature.Message, &message)
+	err = proto.Unmarshal(ekmWithSignature.Message, &message)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to unmarshal message: %w", err)
 	}

diff --git a/pkg/keyexchange/keyexchange_test.go b/pkg/keyexchange/keyexchange_test.go
@@ -94,8 +94,8 @@ func TestKeyExchange_SendAndHandleTimestampMessage(t *testing.T) {
 		if time.Since(start) > 5*time.Second {
 			t.Fatal("timed out")
 		}
-		if _, exists := providerKK.BiddersAESKeys[bidderPeer.EthAddress]; exists {
-			aesKey := providerKK.BiddersAESKeys[bidderPeer.EthAddress]
+		aesKey, exists := providerKK.GetAESKey(bidderPeer.EthAddress)
+		if exists {
 			if !bytes.Equal(bidderKK.AESKey, aesKey) {
 				t.Fatal("AES keys are not equal")
 			}

diff --git a/pkg/keykeeper/keykeeper.go b/pkg/keykeeper/keykeeper.go
@@ -6,6 +6,7 @@ import (
 	"crypto/elliptic"
 	"crypto/rand"
 	"encoding/hex"
+	"sync"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto/ecies"
@@ -72,8 +73,9 @@ func NewProviderKeyKeeper(keysigner keysigner.KeySigner) (*ProviderKeyKeeper, er
 	}
 
 	return &ProviderKeyKeeper{
-		BaseKeyKeeper:  NewBaseKeyKeeper(keysigner),
-		BiddersAESKeys: biddersAESKeys,
+		BaseKeyKeeper:      NewBaseKeyKeeper(keysigner),
+		BiddersAESKeys:     biddersAESKeys,
+		bidderAESKeysMutex: &sync.RWMutex{},
 		keys: ProviderKeys{
 			EncryptionPrivateKey: encryptionPrivateKey,
 			EncryptionPublicKey:  &encryptionPrivateKey.PublicKey,
@@ -98,3 +100,16 @@ func (pkk *ProviderKeyKeeper) DecryptWithECIES(message []byte) ([]byte, error) {
 func (pkk *ProviderKeyKeeper) GetNIKEPrivateKey() *ecdh.PrivateKey {
 	return pkk.keys.NIKEPrivateKey
 }
+
+func (pkk *ProviderKeyKeeper) SetAESKey(bidder common.Address, key []byte) {
+	pkk.bidderAESKeysMutex.Lock()
+	defer pkk.bidderAESKeysMutex.Unlock()
+	pkk.BiddersAESKeys[bidder] = key
+}
+
+func (pkk *ProviderKeyKeeper) GetAESKey(bidder common.Address) ([]byte, bool) {
+	pkk.bidderAESKeysMutex.RLock()
+	defer pkk.bidderAESKeysMutex.RUnlock()
+	key, exists := pkk.BiddersAESKeys[bidder]
+	return key, exists
+}
diff --git a/pkg/keykeeper/models.go b/pkg/keykeeper/models.go
@@ -3,6 +3,7 @@ package keykeeper
 import (
 	"crypto/ecdh"
 	"crypto/ecdsa"
+	"sync"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto/ecies"
@@ -29,8 +30,9 @@ type ProviderKeys struct {
 
 type ProviderKeyKeeper struct {
 	*BaseKeyKeeper
-	keys           ProviderKeys
-	BiddersAESKeys map[common.Address][]byte
+	keys               ProviderKeys
+	bidderAESKeysMutex *sync.RWMutex
+	BiddersAESKeys     map[common.Address][]byte
 }
 
 type BidderKeyKeeper struct {

diff --git a/pkg/signer/preconfencryptor/encryptor_test.go b/pkg/signer/preconfencryptor/encryptor_test.go
@@ -41,7 +41,7 @@ func TestBids(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		providerKeyKeeper.BiddersAESKeys[address] = keyKeeper.AESKey
+		providerKeyKeeper.SetAESKey(address, keyKeeper.AESKey)
 		encryptorProvider := preconfencryptor.NewEncryptor(providerKeyKeeper)
 		bid, err := encryptorProvider.DecryptBidData(address, encryptedBid)
 		if err != nil {
@@ -85,8 +85,7 @@ func TestBids(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		providerKeyKeeper.BiddersAESKeys[bidderAddress] = bidderKeyKeeper.AESKey
-
+		providerKeyKeeper.SetAESKey(bidderAddress, bidderKeyKeeper.AESKey)
 		providerEncryptor := preconfencryptor.NewEncryptor(providerKeyKeeper)
 		start := time.Now().UnixMilli()
 		end := start + 100000