From e2641d1ef10e116c1702c065f2b7e8699a8a397f Mon Sep 17 00:00:00 2001
From: Yun Yeo <sw.yunsuk@gmail.com>
Date: Mon, 30 Aug 2021 16:26:12 +0900
Subject: [PATCH] add mutex to prevent concurrent map access

---
 mutable_tree.go | 29 ++++++++++++++++++++++++++++-
 nodedb.go       |  6 +++---
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/mutable_tree.go b/mutable_tree.go
index b7943f409..764467089 100644
--- a/mutable_tree.go
+++ b/mutable_tree.go
@@ -5,6 +5,7 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"sort"
+	"sync"
 
 	"github.com/pkg/errors"
 
@@ -29,6 +30,8 @@ type MutableTree struct {
 	versions       map[int64]bool   // The previous, saved versions of the tree.
 	allRootLoaded  bool             // Whether all roots are loaded or not(by LazyLoadVersion)
 	ndb            *nodeDB
+
+	mtx sync.RWMutex // versions Read/write lock.
 }
 
 // NewMutableTree returns a new tree with the specified cache size and datastore.
@@ -59,6 +62,9 @@ func (tree *MutableTree) IsEmpty() bool {
 
 // VersionExists returns whether or not a version exists.
 func (tree *MutableTree) VersionExists(version int64) bool {
+	tree.mtx.RLock()
+	defer tree.mtx.RUnlock()
+
 	if tree.allRootLoaded {
 		return tree.versions[version]
 	}
@@ -74,6 +80,9 @@ func (tree *MutableTree) VersionExists(version int64) bool {
 
 // AvailableVersions returns all available versions in ascending order
 func (tree *MutableTree) AvailableVersions() []int {
+	tree.mtx.RLock()
+	defer tree.mtx.RUnlock()
+
 	res := make([]int, 0, len(tree.versions))
 	for i, v := range tree.versions {
 		if v {
@@ -318,6 +327,8 @@ func (tree *MutableTree) LazyLoadVersion(targetVersion int64) (int64, error) {
 		return latestVersion, ErrVersionDoesNotExist
 	}
 
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
 	tree.versions[targetVersion] = true
 
 	iTree := &ImmutableTree{
@@ -354,6 +365,9 @@ func (tree *MutableTree) LoadVersion(targetVersion int64) (int64, error) {
 	firstVersion := int64(0)
 	latestVersion := int64(0)
 
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
+
 	var latestRoot []byte
 	for version, r := range roots {
 		tree.versions[version] = true
@@ -411,6 +425,9 @@ func (tree *MutableTree) LoadVersionForOverwriting(targetVersion int64) (int64,
 
 	tree.ndb.resetLatestVersion(latestVersion)
 
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
+
 	for v := range tree.versions {
 		if v > targetVersion {
 			delete(tree.versions, v)
@@ -429,7 +446,11 @@ func (tree *MutableTree) GetImmutable(version int64) (*ImmutableTree, error) {
 	}
 	if rootHash == nil {
 		return nil, ErrVersionDoesNotExist
-	} else if len(rootHash) == 0 {
+	}
+
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
+	if len(rootHash) == 0 {
 		tree.versions[version] = true
 		return &ImmutableTree{
 			ndb:     tree.ndb,
@@ -526,6 +547,8 @@ func (tree *MutableTree) SaveVersion() ([]byte, int64, error) {
 		return nil, version, err
 	}
 
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
 	tree.version = version
 	tree.versions[version] = true
 
@@ -605,6 +628,8 @@ func (tree *MutableTree) DeleteVersionsRange(fromVersion, toVersion int64) error
 		return err
 	}
 
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
 	for version := fromVersion; version < toVersion; version++ {
 		delete(tree.versions, version)
 	}
@@ -625,6 +650,8 @@ func (tree *MutableTree) DeleteVersion(version int64) error {
 		return err
 	}
 
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
 	delete(tree.versions, version)
 	return nil
 }
diff --git a/nodedb.go b/nodedb.go
index ea26fc821..6b29cf473 100644
--- a/nodedb.go
+++ b/nodedb.go
@@ -36,7 +36,7 @@ var (
 )
 
 type nodeDB struct {
-	mtx            sync.Mutex       // Read/write lock.
+	mtx            sync.RWMutex     // Read/write lock.
 	db             dbm.DB           // Persistent node storage.
 	batch          dbm.Batch        // Batched writing buffer.
 	opts           Options          // Options to customize for pruning/writing
@@ -68,8 +68,8 @@ func newNodeDB(db dbm.DB, cacheSize int, opts *Options) *nodeDB {
 // GetNode gets a node from memory or disk. If it is an inner node, it does not
 // load its children.
 func (ndb *nodeDB) GetNode(hash []byte) *Node {
-	ndb.mtx.Lock()
-	defer ndb.mtx.Unlock()
+	ndb.mtx.RLock()
+	defer ndb.mtx.RUnlock()
 
 	if len(hash) == 0 {
 		panic("nodeDB.GetNode() requires hash")