From 4f3ba6742fbf1cc585048fcc45072523c6de173d 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 a198eb204..c34ac6c78 100644
--- a/trie/stacktrie.go
+++ b/trie/stacktrie.go
@@ -346,8 +346,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
@@ -357,17 +356,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 29706f2e9..50a4eda32 100644
--- a/trie/stacktrie_test.go
+++ b/trie/stacktrie_test.go
@@ -1,9 +1,12 @@
 package trie
 
 import (
+	"bytes"
+	"math/big"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/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)
+		}
+	}
+}