core, core: crypto: dkg: fix concurrent access bug (#551)

* core: crypto: dkg: fix concurrent access bug

* core: fix concurrent bug

2 files changed, 30 insertions, 9 deletions

diff --git a/core/crypto/dkg/dkg.go b/core/crypto/dkg/dkg.go
index b89ad10..ab43f51 100644
--- a/core/crypto/dkg/dkg.go
+++ b/core/crypto/dkg/dkg.go
@@ -450,11 +450,10 @@ func (prvs *PrivateKeyShares) Share(ID ID) (*PrivateKey, bool) {
 
 // NewEmptyPublicKeyShares creates an empty public key shares.
 func NewEmptyPublicKeyShares() *PublicKeyShares {
-	cache := &publicKeySharesCache{
-		index: make(map[ID]int),
-	}
 	pubShares := &PublicKeyShares{}
-	pubShares.cache.Store(cache)
+	pubShares.cache.Store(&publicKeySharesCache{
+		index: make(map[ID]int),
+	})
 	return pubShares
 }
 
@@ -481,9 +480,9 @@ func (pubs *PublicKeyShares) Share(ID ID) (*PublicKey, error) {
 }
 
 // AddShare adds a share.
-func (pubs *PublicKeyShares) AddShare(ID ID, share *PublicKey) error {
+func (pubs *PublicKeyShares) AddShare(shareID ID, share *PublicKey) error {
 	cache := pubs.cache.Load().(*publicKeySharesCache)
-	if idx, exist := cache.index[ID]; exist {
+	if idx, exist := cache.index[shareID]; exist {
 		if !share.publicKey.IsEqual(&cache.share[idx].publicKey) {
 			return ErrDuplicatedShare
 		}
@@ -492,9 +491,17 @@ func (pubs *PublicKeyShares) AddShare(ID ID, share *PublicKey) error {
 	pubs.lock.Lock()
 	defer pubs.lock.Unlock()
 	cache = pubs.cache.Load().(*publicKeySharesCache)
-	cache.index[ID] = len(cache.share)
-	cache.share = append(cache.share, *share)
-	pubs.cache.Store(cache)
+	newCache := &publicKeySharesCache{
+		index: make(map[ID]int, len(cache.index)+1),
+		share: make([]PublicKey, len(cache.share), len(cache.share)+1),
+	}
+	for k, v := range cache.index {
+		newCache.index[k] = v
+	}
+	copy(newCache.share, cache.share)
+	newCache.index[shareID] = len(newCache.share)
+	newCache.share = append(newCache.share, *share)
+	pubs.cache.Store(newCache)
 	return nil
 }
 
diff --git a/core/crypto/dkg/dkg_test.go b/core/crypto/dkg/dkg_test.go
index a679cf4..84dc2d4 100644
--- a/core/crypto/dkg/dkg_test.go
+++ b/core/crypto/dkg/dkg_test.go
@@ -394,6 +394,20 @@ func (s *DKGTestSuite) TestPublicKeySharesMove() {
 	req.True(pubShares2.Equal(pubShares3))
 }
 
+func (s *DKGTestSuite) TestPublicKeySharesConcurrent() {
+	t := 5
+	n := 10
+	IDs := make(IDs, n)
+	for i := range IDs {
+		id := common.NewRandomHash()
+		IDs[i] = NewID(id[:])
+	}
+	_, pubShare := NewPrivateKeyShares(t)
+	for _, id := range IDs {
+		go pubShare.Share(id)
+	}
+}
+
 func (s *DKGTestSuite) TestPrivateKeySharesEquality() {
 	var req = s.Require()
 	IDs := s.genID(2)