fix(trie): differentiate []byte(nil) and []byte{} for storage values (#2964)

Author: qdm12

dot/state/storage_test.go

@@ -113,7 +113,8 @@ func TestStorage_LoadFromDB(t *testing.T) {
 	trieKV := []struct {
 		key   []byte
 		value []byte
-	}{{},
+	}{
+		{value: []byte{}},
 		{[]byte("key1"), []byte("value1")},
 		{[]byte("key2"), []byte("value2")},
 		{[]byte("xyzKey1"), []byte("xyzValue1")},
@@ -147,7 +148,7 @@ func TestStorage_LoadFromDB(t *testing.T) {
 
 	entries, err := storage.Entries(&root)
 	require.NoError(t, err)
-	require.Equal(t, 3, len(entries))
+	require.Equal(t, 4, len(entries))
 }
 
 func TestStorage_StoreTrie_NotSyncing(t *testing.T) {

internal/trie/node/decode.go

@@ -132,15 +132,10 @@ func decodeLeaf(reader io.Reader, partialKeyLength uint16) (node *Node, err erro
 	}
 
 	sd := scale.NewDecoder(reader)
-	var storageValue []byte
-	err = sd.Decode(&storageValue)
-	if err != nil && !errors.Is(err, io.EOF) {
+	err = sd.Decode(&node.StorageValue)
+	if err != nil {
 		return nil, fmt.Errorf("%w: %s", ErrDecodeStorageValue, err)
 	}
 
-	if len(storageValue) > 0 {
-		node.StorageValue = storageValue
-	}
-
 	return node, nil
 }

internal/trie/node/decode_test.go

@@ -334,19 +334,19 @@ func Test_decodeLeaf(t *testing.T) {
 			}),
 			variant:          leafVariant.bits,
 			partialKeyLength: 1,
-			leaf: &Node{
-				PartialKey: []byte{9},
-			},
+			errWrapped:       ErrDecodeStorageValue,
+			errMessage:       "cannot decode storage value: reading byte: EOF",
 		},
 		"empty storage value data": {
 			reader: bytes.NewBuffer(concatByteSlices([][]byte{
-				{9}, // key data
-				scaleEncodeByteSlice(t, nil),
+				{9},                               // key data
+				scaleEncodeByteSlice(t, []byte{}), // results to []byte{0}
 			})),
 			variant:          leafVariant.bits,
 			partialKeyLength: 1,
 			leaf: &Node{
-				PartialKey: []byte{9},
+				PartialKey:   []byte{9},
+				StorageValue: []byte{},
 			},
 		},
 		"success": {

internal/trie/node/encode.go

@@ -37,9 +37,10 @@ func (n *Node) Encode(buffer Buffer) (err error) {
 		}
 	}
 
-	// Only encode node storage value if the node is a leaf or
-	// the node is a branch with a non empty storage value.
-	if !nodeIsBranch || (nodeIsBranch && n.StorageValue != nil) {
+	// Only encode node storage value if the node has a storage value,
+	// even if it is empty. Do not encode if the branch is without value.
+	// Note leaves and branches with value cannot have a `nil` storage value.
+	if n.StorageValue != nil {
 		encoder := scale.NewEncoder(buffer)
 		err = encoder.Encode(n.StorageValue)
 		if err != nil {

internal/trie/node/encode_test.go

@@ -97,13 +97,14 @@ func Test_Node_Encode(t *testing.T) {
 		},
 		"leaf with empty storage value success": {
 			node: &Node{
-				PartialKey: []byte{1, 2, 3},
+				PartialKey:   []byte{1, 2, 3},
+				StorageValue: []byte{},
 			},
 			writes: []writeCall{
 				{written: []byte{leafVariant.bits | 3}}, // partial key length 3
 				{written: []byte{0x01, 0x23}},           // partial key
 				{written: []byte{0}},                    // node storage value encoded length
-				{written: nil},                          // node storage value
+				{written: []byte{}},                     // node storage value
 			},
 			expectedEncoding: []byte{1, 2, 3},
 		},

lib/trie/trie.go

@@ -333,10 +333,10 @@ func (t *Trie) Put(keyLE, value []byte) {
 
 func (t *Trie) insertKeyLE(keyLE, value []byte, deletedMerkleValues map[string]struct{}) {
 	nibblesKey := codec.KeyLEToNibbles(keyLE)
-	if len(value) == 0 {
-		// Force value to be inserted to nil since we don't
-		// differentiate between nil and empty values.
-		value = nil
+	if value == nil {
+		// Force nil value to be inserted to []byte{} since `nil` means there
+		// is no value.
+		value = []byte{}
 	}
 	t.root, _, _ = t.insert(t.root, nibblesKey, value, deletedMerkleValues)
 }

lib/trie/trie_endtoend_test.go

@@ -106,11 +106,7 @@ func Fuzz_Trie_PutAndGet_Single(f *testing.F) {
 		trie := NewEmptyTrie()
 		trie.Put(key, value)
 		retrievedValue := trie.Get(key)
-		if retrievedValue == nil {
-			assert.Empty(t, value)
-		} else {
-			assert.Equal(t, value, retrievedValue)
-		}
+		assert.Equal(t, value, retrievedValue)
 	})
 }
 

pkg/scale/decode.go

@@ -592,9 +592,12 @@ func (ds *decodeState) decodeBytes(dstv reflect.Value) (err error) {
 	}
 
 	b := make([]byte, length)
-	_, err = ds.Read(b)
-	if err != nil {
-		return
+
+	if length > 0 {
+		_, err = ds.Read(b)
+		if err != nil {
+			return
+		}
 	}
 
 	in := dstv.Interface()