From a8bd860e2f879a3b1cd561102be25abb9aac688f Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Fri, 16 Apr 2021 14:21:01 +0200
Subject: [PATCH] trie: make stacktrie not mutate input values (#22673)

The stacktrie is a bit un-untuitive, API-wise: since it mutates input values.
Such behaviour is dangerous, and easy to get wrong if the calling code 'forgets' this quirk. The behaviour is fixed by this PR, so that the input values are not modified by the stacktrie.

Note: just as with the Trie, the stacktrie still references the live input objects, so it's still _not_ safe to mutate the values form the callsite.
---
 trie/stacktrie.go      | 16 +++++--------
 trie/stacktrie_test.go | 51 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 11 deletions(-)

diff --git a/trie/stacktrie.go b/trie/stacktrie.go
index ac3fcb6dc1..d224a983f4 100644
--- a/trie/stacktrie.go
+++ b/trie/stacktrie.go
@@ -347,8 +347,7 @@ func (st *StackTrie) hash() {
 			panic(err)
 		}
 	case emptyNode:
-		st.val = st.val[:0]
-		st.val = append(st.val, emptyRoot[:]...)
+		st.val = emptyRoot.Bytes()
 		st.key = st.key[:0]
 		st.nodeType = hashedNode
 		return
@@ -358,17 +357,12 @@ func (st *StackTrie) hash() {
 	st.key = st.key[:0]
 	st.nodeType = hashedNode
 	if len(h.tmp) < 32 {
-		st.val = st.val[:0]
-		st.val = append(st.val, h.tmp...)
+		st.val = common.CopyBytes(h.tmp)
 		return
 	}
-	// Going to write the hash to the 'val'. Need to ensure it's properly sized first
-	// Typically, 'branchNode's will have no 'val', and require this allocation
-	if required := 32 - len(st.val); required > 0 {
-		buf := make([]byte, required)
-		st.val = append(st.val, buf...)
-	}
-	st.val = st.val[:32]
+	// Write the hash to the 'val'. We allocate a new val here to not mutate
+	// input values
+	st.val = make([]byte, 32)
 	h.sha.Reset()
 	h.sha.Write(h.tmp)
 	h.sha.Read(st.val)
diff --git a/trie/stacktrie_test.go b/trie/stacktrie_test.go
index 562b1c76f4..eaacf2434c 100644
--- a/trie/stacktrie_test.go
+++ b/trie/stacktrie_test.go
@@ -1,9 +1,12 @@
 package trie
 
 import (
+	"bytes"
+	"math/big"
 	"testing"
 
 	"github.com/XinFinOrg/XDPoSChain/common"
+	"github.com/XinFinOrg/XDPoSChain/crypto"
 	"github.com/XinFinOrg/XDPoSChain/ethdb/memorydb"
 )
 
@@ -119,3 +122,51 @@ func TestUpdateVariableKeys(t *testing.T) {
 		t.Fatalf("error %x != %x", st.Hash(), nt.Hash())
 	}
 }
+
+// TestStacktrieNotModifyValues checks that inserting blobs of data into the
+// stacktrie does not mutate the blobs
+func TestStacktrieNotModifyValues(t *testing.T) {
+	st := NewStackTrie(nil)
+	{ // Test a very small trie
+		// Give it the value as a slice with large backing alloc,
+		// so if the stacktrie tries to append, it won't have to realloc
+		value := make([]byte, 1, 100)
+		value[0] = 0x2
+		want := common.CopyBytes(value)
+		st.TryUpdate([]byte{0x01}, value)
+		st.Hash()
+		if have := value; !bytes.Equal(have, want) {
+			t.Fatalf("tiny trie: have %#x want %#x", have, want)
+		}
+		st = NewStackTrie(nil)
+	}
+	// Test with a larger trie
+	keyB := big.NewInt(1)
+	keyDelta := big.NewInt(1)
+	var vals [][]byte
+	getValue := func(i int) []byte {
+		if i%2 == 0 { // large
+			return crypto.Keccak256(big.NewInt(int64(i)).Bytes())
+		} else { //small
+			return big.NewInt(int64(i)).Bytes()
+		}
+	}
+
+	for i := 0; i < 1000; i++ {
+		key := common.BigToHash(keyB)
+		value := getValue(i)
+		st.TryUpdate(key.Bytes(), value)
+		vals = append(vals, value)
+		keyB = keyB.Add(keyB, keyDelta)
+		keyDelta.Add(keyDelta, common.Big1)
+	}
+	st.Hash()
+	for i := 0; i < 1000; i++ {
+		want := getValue(i)
+
+		have := vals[i]
+		if !bytes.Equal(have, want) {
+			t.Fatalf("item %d, have %#x want %#x", i, have, want)
+		}
+	}
+}