From 159228b34ffe0875d9acb5a1d2592d44fb262c77 Mon Sep 17 00:00:00 2001
From: David Theodore <29786815+infosecual@users.noreply.github.com>
Date: Thu, 15 Dec 2022 00:50:12 -0600
Subject: [PATCH] SparseMerkleTrie depth fix & regression tests (#11778)

---
 container/trie/sparse_merkle.go      | 10 +++++++--
 container/trie/sparse_merkle_test.go | 31 ++++++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/container/trie/sparse_merkle.go b/container/trie/sparse_merkle.go
index ce362c7cc9..bbba0643ac 100644
--- a/container/trie/sparse_merkle.go
+++ b/container/trie/sparse_merkle.go
@@ -57,8 +57,9 @@ func (m *SparseMerkleTrie) validate() error {
 	if m.depth >= uint(len(m.branches)) {
 		return errors.New("depth is greater than or equal to number of branches")
 	}
-	if m.depth >= 64 {
-		return errors.New("depth exceeds 64") // PowerOf2 would overflow.
+	if m.depth >= 63 {
+		return errors.New("supported merkle trie depth exceeded (max uint64 depth is 63, " +
+			"theoretical max sparse merkle trie depth is 64)") // PowerOf2 would overflow.
 	}
 
 	return nil
@@ -69,6 +70,11 @@ func GenerateTrieFromItems(items [][]byte, depth uint64) (*SparseMerkleTrie, err
 	if len(items) == 0 {
 		return nil, errors.New("no items provided to generate Merkle trie")
 	}
+	if depth >= 63 {
+		return nil, errors.New("supported merkle trie depth exceeded (max uint64 depth is 63, " +
+			"theoretical max sparse merkle trie depth is 64)") // PowerOf2 would overflow
+	}
+
 	leaves := items
 	layers := make([][][]byte, depth+1)
 	transformedLeaves := make([][]byte, len(leaves))
diff --git a/container/trie/sparse_merkle_test.go b/container/trie/sparse_merkle_test.go
index 63ed108639..f32df7f37e 100644
--- a/container/trie/sparse_merkle_test.go
+++ b/container/trie/sparse_merkle_test.go
@@ -65,9 +65,10 @@ func TestCreateTrieFromProto_Validation(t *testing.T) {
 		{
 			trie: &ethpb.SparseMerkleTrie{
 				Layers: genValidLayers(66),
-				Depth:  65,
+				Depth:  63,
 			},
-			errString: "depth exceeds 64",
+			errString: "supported merkle trie depth exceeded (max uint64 depth is 63, " +
+				"theoretical max sparse merkle trie depth is 64)",
 		},
 	}
 	for _, tt := range tests {
@@ -153,6 +154,32 @@ func TestGenerateTrieFromItems_NoItemsProvided(t *testing.T) {
 	}
 }
 
+func TestGenerateTrieFromItems_DepthSupport(t *testing.T) {
+	items := [][]byte{
+		[]byte("A"),
+		[]byte("BB"),
+		[]byte("CCC"),
+		[]byte("DDDD"),
+		[]byte("EEEEE"),
+		[]byte("FFFFFF"),
+		[]byte("GGGGGGG"),
+	}
+	// max supported depth is 62 (uint64 will overflow above this)
+	// max theoretical depth is 64
+	var max_supported_trie_depth uint64 = 62
+	// Supported depth
+	m1, err := trie.GenerateTrieFromItems(items, max_supported_trie_depth)
+	require.NoError(t, err)
+	proof, err := m1.MerkleProof(2)
+	require.NoError(t, err)
+	require.Equal(t, len(proof), int(max_supported_trie_depth)+1)
+	// Unsupported depth
+	_, err = trie.GenerateTrieFromItems(items, max_supported_trie_depth+1)
+	errString := "supported merkle trie depth exceeded (max uint64 depth is 63, " +
+		"theoretical max sparse merkle trie depth is 64)"
+	require.ErrorContains(t, errString, err)
+}
+
 func TestMerkleTrie_VerifyMerkleProofWithDepth(t *testing.T) {
 	items := [][]byte{
 		[]byte("A"),