lib/db: Fix race in NamespacedKV (#5496)

1 files changed, 22 insertions, 24 deletions

diff --git a/lib/db/namespaced.go b/lib/db/namespaced.go
index 7a26bdff52..fd6432726e 100644
--- a/lib/db/namespaced.go
+++ b/lib/db/namespaced.go
@@ -24,9 +24,14 @@ type NamespacedKV struct {
 // NewNamespacedKV returns a new NamespacedKV that lives in the namespace
 // specified by the prefix.
 func NewNamespacedKV(db *Lowlevel, prefix string) *NamespacedKV {
+	prefixBs := []byte(prefix)
+	// After the conversion from string the cap will be larger than the len (in Go 1.11.5,
+	// 32 bytes cap for small strings). We need to cut it down to ensure append() calls
+	// on the prefix make a new allocation.
+	prefixBs = prefixBs[:len(prefixBs):len(prefixBs)]
 	return &NamespacedKV{
 		db:     db,
-		prefix: []byte(prefix),
+		prefix: prefixBs,
 	}
 }
 
@@ -54,17 +59,15 @@ func (n *NamespacedKV) Reset() {
 // PutInt64 stores a new int64. Any existing value (even if of another type)
 // is overwritten.
 func (n *NamespacedKV) PutInt64(key string, val int64) {
-	keyBs := append(n.prefix, []byte(key)...)
 	var valBs [8]byte
 	binary.BigEndian.PutUint64(valBs[:], uint64(val))
-	n.db.Put(keyBs, valBs[:], nil)
+	n.db.Put(n.prefixedKey(key), valBs[:], nil)
 }
 
 // Int64 returns the stored value interpreted as an int64 and a boolean that
 // is false if no value was stored at the key.
 func (n *NamespacedKV) Int64(key string) (int64, bool) {
-	keyBs := append(n.prefix, []byte(key)...)
-	valBs, err := n.db.Get(keyBs, nil)
+	valBs, err := n.db.Get(n.prefixedKey(key), nil)
 	if err != nil {
 		return 0, false
 	}
@@ -75,17 +78,15 @@ func (n *NamespacedKV) Int64(key string) (int64, bool) {
 // PutTime stores a new time.Time. Any existing value (even if of another
 // type) is overwritten.
 func (n *NamespacedKV) PutTime(key string, val time.Time) {
-	keyBs := append(n.prefix, []byte(key)...)
 	valBs, _ := val.MarshalBinary() // never returns an error
-	n.db.Put(keyBs, valBs, nil)
+	n.db.Put(n.prefixedKey(key), valBs, nil)
 }
 
 // Time returns the stored value interpreted as a time.Time and a boolean
 // that is false if no value was stored at the key.
 func (n NamespacedKV) Time(key string) (time.Time, bool) {
 	var t time.Time
-	keyBs := append(n.prefix, []byte(key)...)
-	valBs, err := n.db.Get(keyBs, nil)
+	valBs, err := n.db.Get(n.prefixedKey(key), nil)
 	if err != nil {
 		return t, false
 	}
@@ -96,15 +97,13 @@ func (n NamespacedKV) Time(key string) (time.Time, bool) {
 // PutString stores a new string. Any existing value (even if of another type)
 // is overwritten.
 func (n *NamespacedKV) PutString(key, val string) {
-	keyBs := append(n.prefix, []byte(key)...)
-	n.db.Put(keyBs, []byte(val), nil)
+	n.db.Put(n.prefixedKey(key), []byte(val), nil)
 }
 
 // String returns the stored value interpreted as a string and a boolean that
 // is false if no value was stored at the key.
 func (n NamespacedKV) String(key string) (string, bool) {
-	keyBs := append(n.prefix, []byte(key)...)
-	valBs, err := n.db.Get(keyBs, nil)
+	valBs, err := n.db.Get(n.prefixedKey(key), nil)
 	if err != nil {
 		return "", false
 	}
@@ -114,15 +113,13 @@ func (n NamespacedKV) String(key string) (string, bool) {
 // PutBytes stores a new byte slice. Any existing value (even if of another type)
 // is overwritten.
 func (n *NamespacedKV) PutBytes(key string, val []byte) {
-	keyBs := append(n.prefix, []byte(key)...)
-	n.db.Put(keyBs, val, nil)
+	n.db.Put(n.prefixedKey(key), val, nil)
 }
 
 // Bytes returns the stored value as a raw byte slice and a boolean that
 // is false if no value was stored at the key.
 func (n NamespacedKV) Bytes(key string) ([]byte, bool) {
-	keyBs := append(n.prefix, []byte(key)...)
-	valBs, err := n.db.Get(keyBs, nil)
+	valBs, err := n.db.Get(n.prefixedKey(key), nil)
 	if err != nil {
 		return nil, false
 	}
@@ -132,19 +129,17 @@ func (n NamespacedKV) Bytes(key string) ([]byte, bool) {
 // PutBool stores a new boolean. Any existing value (even if of another type)
 // is overwritten.
 func (n *NamespacedKV) PutBool(key string, val bool) {
-	keyBs := append(n.prefix, []byte(key)...)
 	if val {
-		n.db.Put(keyBs, []byte{0x0}, nil)
+		n.db.Put(n.prefixedKey(key), []byte{0x0}, nil)
 	} else {
-		n.db.Put(keyBs, []byte{0x1}, nil)
+		n.db.Put(n.prefixedKey(key), []byte{0x1}, nil)
 	}
 }
 
 // Bool returns the stored value as a boolean and a boolean that
 // is false if no value was stored at the key.
 func (n NamespacedKV) Bool(key string) (bool, bool) {
-	keyBs := append(n.prefix, []byte(key)...)
-	valBs, err := n.db.Get(keyBs, nil)
+	valBs, err := n.db.Get(n.prefixedKey(key), nil)
 	if err != nil {
 		return false, false
 	}
@@ -154,8 +149,11 @@ func (n NamespacedKV) Bool(key string) (bool, bool) {
 // Delete deletes the specified key. It is allowed to delete a nonexistent
 // key.
 func (n NamespacedKV) Delete(key string) {
-	keyBs := append(n.prefix, []byte(key)...)
-	n.db.Delete(keyBs, nil)
+	n.db.Delete(n.prefixedKey(key), nil)
+}
+
+func (n NamespacedKV) prefixedKey(key string) []byte {
+	return append(n.prefix, []byte(key)...)
 }
 
 // Well known namespaces that can be instantiated without knowing the key