From 568f0e2d62987a29c65db9b8232b4242fd9d2b13 Mon Sep 17 00:00:00 2001
From: CHAMI Rachid <chamirachid1@gmail.com>
Date: Fri, 2 Aug 2024 17:48:19 +0200
Subject: [PATCH] fix: check the correctness of the leaf ranges (#268)

<!--
Please read and fill out this form before submitting your PR.

Please make sure you have reviewed our contributors guide before
submitting your
first PR.
-->

## Overview

Closes: https://github.com/celestiaorg/nmt/issues/267

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit


- **New Features**
- Strengthened error handling for subtree root verification and leaf
range calculations to improve robustness against invalid inputs.

- **Bug Fixes**
- Enhanced validation processes to prevent incorrect operations and
ensure reliable range processing.

- **Tests**
- Refined test cases to focus on relevant scenarios and added new tests
for improved coverage of edge cases related to error handling.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 proof.go      | 12 +++++++-
 proof_test.go | 77 ++++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 69 insertions(+), 20 deletions(-)

diff --git a/proof.go b/proof.go
index bde6723..b260e17 100644
--- a/proof.go
+++ b/proof.go
@@ -528,6 +528,10 @@ func (proof Proof) VerifySubtreeRootInclusion(nth *NmtHasher, subtreeRoots [][]b
 			return popIfNonEmpty(&subtreeRoots), nil
 		}
 
+		if end-start == 1 {
+			return nil, fmt.Errorf("the provided range [%d, %d) does not reference a valid inner node", proof.start, proof.end)
+		}
+
 		// Recursively get left and right subtree
 		k := getSplitPoint(end - start)
 		left, err := computeRoot(start, start+k)
@@ -634,7 +638,13 @@ func nextLeafRange(currentStart, currentEnd, subtreeWidth int) (LeafRange, error
 	if err != nil {
 		return LeafRange{}, err
 	}
-	return LeafRange{Start: currentStart, End: currentStart + currentRange}, nil
+	rangeEnd := currentStart + currentRange
+	idealTreeSize := nextSubtreeSize(uint64(currentStart), uint64(rangeEnd))
+	if currentStart+idealTreeSize != rangeEnd {
+		// this will happen if the calculated range does not correctly reference an inner node in the tree.
+		return LeafRange{}, fmt.Errorf("provided subtree width %d doesn't allow creating a valid leaf range [%d, %d)", subtreeWidth, currentStart, rangeEnd)
+	}
+	return LeafRange{Start: currentStart, End: rangeEnd}, nil
 }
 
 // largestPowerOfTwo calculates the largest power of two
diff --git a/proof_test.go b/proof_test.go
index b2ea98f..cb208cd 100644
--- a/proof_test.go
+++ b/proof_test.go
@@ -1332,18 +1332,6 @@ func TestNextLeafRange(t *testing.T) {
 			subtreeRootMaximumLeafRange: 8,
 			expectedRange:               LeafRange{Start: 4, End: 8},
 		},
-		{
-			currentStart:                4,
-			currentEnd:                  20,
-			subtreeRootMaximumLeafRange: 16,
-			expectedRange:               LeafRange{Start: 4, End: 20},
-		},
-		{
-			currentStart:                4,
-			currentEnd:                  20,
-			subtreeRootMaximumLeafRange: 1,
-			expectedRange:               LeafRange{Start: 4, End: 5},
-		},
 		{
 			currentStart:                4,
 			currentEnd:                  20,
@@ -1356,12 +1344,6 @@ func TestNextLeafRange(t *testing.T) {
 			subtreeRootMaximumLeafRange: 4,
 			expectedRange:               LeafRange{Start: 4, End: 8},
 		},
-		{
-			currentStart:                4,
-			currentEnd:                  20,
-			subtreeRootMaximumLeafRange: 8,
-			expectedRange:               LeafRange{Start: 4, End: 12},
-		},
 		{
 			currentStart:                0,
 			currentEnd:                  1,
@@ -1392,6 +1374,36 @@ func TestNextLeafRange(t *testing.T) {
 			subtreeRootMaximumLeafRange: 0,
 			expectError:                 true,
 		},
+		{ // A range not referencing any inner node
+			currentStart:                1,
+			currentEnd:                  3,
+			subtreeRootMaximumLeafRange: 4,
+			expectError:                 true,
+		},
+		{ // A range not referencing any inner node
+			currentStart:                1,
+			currentEnd:                  5,
+			subtreeRootMaximumLeafRange: 4,
+			expectError:                 true,
+		},
+		{ // A range not referencing any inner node
+			currentStart:                1,
+			currentEnd:                  6,
+			subtreeRootMaximumLeafRange: 4,
+			expectError:                 true,
+		},
+		{ // A range not referencing any inner node
+			currentStart:                1,
+			currentEnd:                  7,
+			subtreeRootMaximumLeafRange: 4,
+			expectError:                 true,
+		},
+		{ // A range not referencing any inner node
+			currentStart:                2,
+			currentEnd:                  8,
+			subtreeRootMaximumLeafRange: 4,
+			expectError:                 true,
+		},
 	}
 
 	for _, tt := range tests {
@@ -1798,7 +1810,6 @@ func TestVerifySubtreeRootInclusion(t *testing.T) {
 			root:         root,
 			expectError:  true,
 		},
-
 		{
 			proof: func() Proof {
 				p, err := tree.ProveRange(0, 8)
@@ -1828,3 +1839,31 @@ func TestVerifySubtreeRootInclusion(t *testing.T) {
 		})
 	}
 }
+
+// TestVerifySubtreeRootInclusion_infiniteRecursion is motivated by a failing test
+// case in celestia-node
+func TestVerifySubtreeRootInclusion_infiniteRecursion(t *testing.T) {
+	namespaceIDs := bytes.Repeat([]byte{1}, 64)
+	tree := exampleNMT(1, true, namespaceIDs...)
+	root, err := tree.Root()
+	require.NoError(t, err)
+
+	nmthasher := tree.treeHasher
+	hasher := nmthasher.(*NmtHasher)
+	subtreeRoot, err := tree.ComputeSubtreeRoot(0, 4)
+	require.NoError(t, err)
+	subtreeRoots := [][]byte{subtreeRoot, subtreeRoot, subtreeRoot, subtreeRoot, subtreeRoot, subtreeRoot, subtreeRoot}
+	subtreeWidth := 8
+
+	proof, err := tree.ProveRange(19, 64)
+	require.NoError(t, err)
+
+	require.NotPanics(t, func() {
+		// This previously hits:
+		// runtime: goroutine stack exceeds 1000000000-byte limit
+		// runtime: sp=0x14020160480 stack=[0x14020160000, 0x14040160000]
+		// fatal error: stack overflow
+		_, err = proof.VerifySubtreeRootInclusion(hasher, subtreeRoots, subtreeWidth, root)
+		require.Error(t, err)
+	})
+}