fix(trie): disallow empty byte slice node values (#2927)

- Always encode empty values for leaf nodes only
- Force empty byte slice values to `nil` when inserting
- Remove now unneeded `SubValueEqual` method
- Add extra test cases

Author: qdm12

internal/trie/node/copy.go

@@ -87,7 +87,8 @@ func (n *Node) Copy(settings CopySettings) *Node {
 		copy(cpy.Key, n.Key)
 	}
 
-	// nil and []byte{} are encoded differently, watch out!
+	// nil and []byte{} values for branches result in a different node encoding,
+	// so we ensure to keep the `nil` value.
 	if settings.CopyValue && n.SubValue != nil {
 		cpy.SubValue = make([]byte, len(n.SubValue))
 		copy(cpy.SubValue, n.SubValue)

internal/trie/node/decode_test.go

@@ -327,7 +327,7 @@ func Test_decodeLeaf(t *testing.T) {
 			errWrapped:       ErrDecodeValue,
 			errMessage:       "cannot decode value: unknown prefix for compact uint: 255",
 		},
-		"zero value": {
+		"missing value data": {
 			reader: bytes.NewBuffer([]byte{
 				9, // key data
 				// missing value data
@@ -338,6 +338,17 @@ func Test_decodeLeaf(t *testing.T) {
 				Key: []byte{9},
 			},
 		},
+		"empty value data": {
+			reader: bytes.NewBuffer(concatByteSlices([][]byte{
+				{9}, // key data
+				scaleEncodeByteSlice(t, nil),
+			})),
+			variant:          leafVariant.bits,
+			partialKeyLength: 1,
+			leaf: &Node{
+				Key: []byte{9},
+			},
+		},
 		"success": {
 			reader: bytes.NewBuffer(
 				concatByteSlices([][]byte{

internal/trie/node/encode.go

@@ -35,17 +35,19 @@ func (n *Node) Encode(buffer Buffer) (err error) {
 		return fmt.Errorf("cannot write LE key to buffer: %w", err)
 	}
 
-	if n.Kind() == Branch {
+	kind := n.Kind()
+	nodeIsBranch := kind == Branch
+	if nodeIsBranch {
 		childrenBitmap := common.Uint16ToBytes(n.ChildrenBitmap())
 		_, err = buffer.Write(childrenBitmap)
 		if err != nil {
 			return fmt.Errorf("cannot write children bitmap to buffer: %w", err)
 		}
 	}
 
-	// check value is not nil for branch nodes, even though
-	// leaf nodes always have a non-nil value.
-	if n.SubValue != nil {
+	// Only encode node value if the node is a leaf or
+	// the node is a branch with a non empty value.
+	if !nodeIsBranch || (nodeIsBranch && n.SubValue != nil) {
 		encodedValue, err := scale.Marshal(n.SubValue) // TODO scale encoder to write to buffer
 		if err != nil {
 			return fmt.Errorf("cannot scale encode value: %w", err)
@@ -57,14 +59,14 @@ func (n *Node) Encode(buffer Buffer) (err error) {
 		}
 	}
 
-	if n.Kind() == Branch {
+	if nodeIsBranch {
 		err = encodeChildrenOpportunisticParallel(n.Children, buffer)
 		if err != nil {
 			return fmt.Errorf("cannot encode children of branch: %w", err)
 		}
 	}
 
-	if n.Kind() == Leaf {
+	if kind == Leaf {
 		// TODO cache this for branches too and update test cases.
 		// TODO remove this copying since it defeats the purpose of `buffer`
 		// and the sync.Pool.

internal/trie/node/encode_test.go

@@ -127,6 +127,25 @@ func Test_Node_Encode(t *testing.T) {
 			bufferBytesCall:  true,
 			expectedEncoding: []byte{1, 2, 3},
 		},
+		"leaf with empty value success": {
+			node: &Node{
+				Key: []byte{1, 2, 3},
+			},
+			writes: []writeCall{
+				{
+					written: []byte{leafVariant.bits | 3}, // partial key length 3
+				},
+				{
+					written: []byte{0x01, 0x23},
+				},
+				{
+					written: []byte{0},
+				},
+			},
+			bufferLenCall:    true,
+			bufferBytesCall:  true,
+			expectedEncoding: []byte{1, 2, 3},
+		},
 		"clean branch with encoding": {
 			node: &Node{
 				Children: make([]*Node, ChildrenCapacity),
@@ -297,6 +316,32 @@ func Test_Node_Encode(t *testing.T) {
 				},
 			},
 		},
+		"branch without value and with children success": {
+			node: &Node{
+				Key: []byte{1, 2, 3},
+				Children: []*Node{
+					nil, nil, nil, {Key: []byte{9}, SubValue: []byte{1}},
+					nil, nil, nil, {Key: []byte{11}, SubValue: []byte{1}},
+				},
+			},
+			writes: []writeCall{
+				{ // header
+					written: []byte{branchVariant.bits | 3}, // partial key length 3
+				},
+				{ // key LE
+					written: []byte{0x01, 0x23},
+				},
+				{ // children bitmap
+					written: []byte{136, 0},
+				},
+				{ // first children
+					written: []byte{16, 65, 9, 4, 1},
+				},
+				{ // second children
+					written: []byte{16, 65, 11, 4, 1},
+				},
+			},
+		},
 	}
 
 	for name, testCase := range testCases {

internal/trie/node/subvalue.go

@@ -342,6 +342,11 @@ 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
+	}
 	t.root, _, _ = t.insert(t.root, nibblesKey, value, deletedMerkleValues)
 }
 
@@ -374,7 +379,7 @@ func (t *Trie) insertInLeaf(parentLeaf *Node, key, value []byte,
 	newParent *Node, mutated bool, nodesCreated uint32) {
 	if bytes.Equal(parentLeaf.Key, key) {
 		nodesCreated = 0
-		if parentLeaf.SubValueEqual(value) {
+		if bytes.Equal(parentLeaf.SubValue, value) {
 			mutated = false
 			return parentLeaf, mutated, nodesCreated
 		}
@@ -455,7 +460,7 @@ func (t *Trie) insertInBranch(parentBranch *Node, key, value []byte,
 	copySettings := node.DefaultCopySettings
 
 	if bytes.Equal(key, parentBranch.Key) {
-		if parentBranch.SubValueEqual(value) {
+		if bytes.Equal(parentBranch.SubValue, value) {
 			mutated = false
 			return parentBranch, mutated, 0
 		}

internal/trie/node/subvalue_test.go

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