From aefcd2d64b9eb03144cb07824d7d50d1251ff95d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ecl=C3=A9sio=20Junior?= <eclesiomelo.1@gmail.com>
Date: Fri, 4 Feb 2022 12:43:38 -0400
Subject: [PATCH] fix(lib/trie): remove map deletion at `loadProof` (#2259)

* fix: remove map deletion to every proof hash

* chore: use `node` as variable name
---
 lib/trie/database.go   |  1 -
 lib/trie/proof_test.go | 82 +++++++++++++++++++++++++++++++++++++-----
 2 files changed, 73 insertions(+), 10 deletions(-)

diff --git a/lib/trie/database.go b/lib/trie/database.go
index 87a8425f71..312da1ca3d 100644
--- a/lib/trie/database.go
+++ b/lib/trie/database.go
@@ -137,7 +137,6 @@ func (t *Trie) loadProof(proofHashToNode map[string]Node, n Node) {
 		if !ok {
 			continue
 		}
-		delete(proofHashToNode, proofHash)
 
 		branch.Children[i] = node
 		t.loadProof(proofHashToNode, node)
diff --git a/lib/trie/proof_test.go b/lib/trie/proof_test.go
index 7c190d1c3c..7ab5c60361 100644
--- a/lib/trie/proof_test.go
+++ b/lib/trie/proof_test.go
@@ -78,11 +78,10 @@ func testGenerateProof(t *testing.T, entries []Pair, keys [][]byte) ([]byte, [][
 		value := trie.Get(key)
 		require.NotNil(t, value)
 
-		itemFromDB := Pair{
+		items[idx] = Pair{
 			Key:   key,
 			Value: value,
 		}
-		items[idx] = itemFromDB
 	}
 
 	return root, proof, items
@@ -95,23 +94,23 @@ func TestVerifyProof_ShouldReturnTrue(t *testing.T) {
 		{Key: []byte("alpha"), Value: make([]byte, 32)},
 		{Key: []byte("bravo"), Value: []byte("bravo")},
 		{Key: []byte("do"), Value: []byte("verb")},
-		{Key: []byte("dog"), Value: []byte("puppy")},
-		{Key: []byte("doge"), Value: make([]byte, 32)},
+		{Key: []byte("dogea"), Value: []byte("puppy")},
+		{Key: []byte("dogeb"), Value: []byte("puppy")},
 		{Key: []byte("horse"), Value: []byte("stallion")},
 		{Key: []byte("house"), Value: []byte("building")},
 	}
 
 	keys := [][]byte{
 		[]byte("do"),
-		[]byte("dog"),
-		[]byte("doge"),
+		[]byte("dogea"),
+		[]byte("dogeb"),
 	}
 
-	root, proof, pl := testGenerateProof(t, entries, keys)
+	root, proof, pairs := testGenerateProof(t, entries, keys)
+	v, err := VerifyProof(proof, root, pairs)
 
-	v, err := VerifyProof(proof, root, pl)
-	require.True(t, v)
 	require.NoError(t, err)
+	require.True(t, v)
 }
 
 func TestVerifyProof_ShouldReturnDuplicateKeysError(t *testing.T) {
@@ -158,3 +157,68 @@ func TestVerifyProof_ShouldReturnTrueWithouCompareValues(t *testing.T) {
 	require.True(t, v)
 	require.NoError(t, err)
 }
+
+func TestBranchNodes_SameHash_DiferentPaths_GenerateAndVerifyProof(t *testing.T) {
+	value := []byte("somevalue")
+	entries := []Pair{
+		{Key: []byte("d"), Value: value},
+		{Key: []byte("b"), Value: value},
+		{Key: []byte("dxyz"), Value: value},
+		{Key: []byte("bxyz"), Value: value},
+		{Key: []byte("dxyzi"), Value: value},
+		{Key: []byte("bxyzi"), Value: value},
+	}
+
+	keys := [][]byte{
+		[]byte("d"),
+		[]byte("b"),
+		[]byte("dxyz"),
+		[]byte("bxyz"),
+		[]byte("dxyzi"),
+		[]byte("bxyzi"),
+	}
+
+	root, proof, pairs := testGenerateProof(t, entries, keys)
+
+	ok, err := VerifyProof(proof, root, pairs)
+	require.NoError(t, err)
+	require.True(t, ok)
+}
+
+func TestLeafNodes_SameHash_DifferentPaths_GenerateAndVerifyProof(t *testing.T) {
+	tmp := t.TempDir()
+
+	memdb, err := chaindb.NewBadgerDB(&chaindb.Config{
+		InMemory: true,
+		DataDir:  tmp,
+	})
+	require.NoError(t, err)
+
+	var (
+		value = []byte("somevalue")
+		key1  = []byte("worlda")
+		key2  = []byte("worldb")
+	)
+
+	tt := NewEmptyTrie()
+	tt.Put(key1, value)
+	tt.Put(key2, value)
+
+	err = tt.Store(memdb)
+	require.NoError(t, err)
+
+	hash, err := tt.Hash()
+	require.NoError(t, err)
+
+	proof, err := GenerateProof(hash.ToBytes(), [][]byte{key1, key2}, memdb)
+	require.NoError(t, err)
+
+	pairs := []Pair{
+		{Key: key1, Value: value},
+		{Key: key2, Value: value},
+	}
+
+	ok, err := VerifyProof(proof, hash.ToBytes(), pairs)
+	require.NoError(t, err)
+	require.True(t, ok)
+}