Repository: celestiaorg/nmt
Pull request #253 (2024-05-13 18:19:15)

feat!: adds input range check to optimize VerifyLeafHashes and VerifyInclusion methods
This PR optimizes the `VerifyLeafHashes` and `VerifyInclusion` methods by ensuring that the size of the provided leaves or leaf hashes matches the proof range and, if not, making an early return in order to prevent unnecessary hashing operations that would otherwise occur.
It introduces a breaking change by altering the behaviour of `VerifyLeafHashes` and `VerifyInclusion`. Previously, verification would succeed even if the provided leaves or leaf hashes exceeded the proof range, as long as the proof was valid and the extra leaves were appended at the end. In the new implementation, the verification will return false or an error in such cases.

staheri14: [To the reviewers] The codecov [error](https://github.com/celestiaorg/nmt/actions/runs/9068204992/job/24914967372?pr=254#step:8:32) is not specific to this PR, I am investigating it.  
staheri14: Waiting on #255 to merge
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/celestiaorg/nmt/pull/253?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=celestiaorg) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 66.50%. Comparing base [(`2c0563a`)](https://app.codecov.io/gh/celestiaorg/nmt/commit/2c0563aa85a8a087dde05fa57a32b4fe6938d46a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=celestiaorg) to head [(`864d8dc`)](https://app.codecov.io/gh/celestiaorg/nmt/pull/253?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=celestiaorg).


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #253      +/-   ##
==========================================
+ Coverage   66.24%   66.50%   +0.26%     
==========================================
  Files           6        6              
  Lines        1022     1030       +8     
==========================================
+ Hits          677      685       +8     
  Misses        328      328              
  Partials       17       17              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/celestiaorg/nmt/pull/253?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=celestiaorg).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=celestiaorg).


Patch:
diff --git a/docs/spec/nmt.md b/docs/spec/nmt.md
index 40bf3ff..784a7cf 100644
--- a/docs/spec/nmt.md
+++ b/docs/spec/nmt.md
@@ -198,21 +198,21 @@ More formally, the short namespace absence proof consists of the following compo
    1) Find the index of a leaf in the tree that meets two conditions:
       1) Its namespace is the smallest namespace greater than `NS`.
       1) The namespace of the leaf to its left is smaller than `NS`.
-   1) Traverse up the branch connecting that leaf to the root and locate one of the parents/grandparents of that leaf whose namespace range does not overlap with the queried namespace. 
+   1) Traverse up the branch connecting that leaf to the root and locate one of the parents/grandparents of that leaf whose namespace range does not overlap with the queried namespace.
    The `SubtreeHash` is the hash of that node.
-1) `start` and `end` range: These represent the indices of the `SubtreeHash` within its respective level. 
+1) `start` and `end` range: These represent the indices of the `SubtreeHash` within its respective level.
 Nodes at each level are indexed from left to right starting at index `0`.
 1) `nodes`: This set comprises the index-based  Merkle inclusion proof of the `SubtreeHash` to the tree root `T`.
 
 Below, we illustrate the short namespace absence proof for namespace `NS = 02` in an 8-leaf tree:
-The namespace `03` is the smallest namespace larger than `02`. 
-By traversing the branch from the leaf with namespace `03` to the root, we find a node with hash `03 04 52c7c03` whose namespace range doesn't overlap with `02`. 
+The namespace `03` is the smallest namespace larger than `02`.
+By traversing the branch from the leaf with namespace `03` to the root, we find a node with hash `03 04 52c7c03` whose namespace range doesn't overlap with `02`.
 This node is the highest such node along the branch.
 The `SubtreeHash` is the hash of that node, which is `03 04 52c7c03`.
-The `start` and `end` indices indicate its position in the respective level. 
-In this case, `start = 1` and `end = 2`. 
+The `start` and `end` indices indicate its position in the respective level.
+In this case, `start = 1` and `end = 2`.
 Note that node indices start at `0` from left to right at each level.
-The `nodes` form the index-based Merkle inclusion proof of the `SubtreeHash` to the tree root `T`. 
+The `nodes` form the index-based Merkle inclusion proof of the `SubtreeHash` to the tree root `T`.
 The `nodes` set includes `00 00 ead8d25`, the left sibling of `03 04 52c7c03`.
 
 In summary, the short namespace absence proof for `NS = 02` in this tree consists of `SubtreeHash = 03 04 52c7c03`, `start = 1`, `end = 2`, and the `nodes` set containing `00 00 ead8d25`.
diff --git a/proof.go b/proof.go
index 083370e..998b9a8 100644
--- a/proof.go
+++ b/proof.go
@@ -12,8 +12,11 @@ import (
 	pb "github.com/celestiaorg/nmt/pb"
 )
 
-// ErrFailedCompletenessCheck indicates that the verification of a namespace proof failed due to the lack of completeness property.
-var ErrFailedCompletenessCheck = errors.New("failed completeness check")
+var (
+	// ErrFailedCompletenessCheck indicates that the verification of a namespace proof failed due to the lack of completeness property.
+	ErrFailedCompletenessCheck = errors.New("failed completeness check")
+	ErrWrongLeafHashesSize     = errors.New("wrong leafHashes size")
+)
 
 // Proof represents a namespace proof of a namespace.ID in an NMT. In case this
 // proof proves the absence of a namespace.ID in a tree it also contains the
@@ -250,6 +253,7 @@ func (proof Proof) VerifyNamespace(h hash.Hash, nID namespace.ID, leaves [][]byt
 // If there is an issue during the proof verification e.g., a node does not conform to the namespace hash format, then a proper error is returned to indicate the root cause of the issue.
 // The leafHashes parameter is a list of leaf hashes, where each leaf hash is represented
 // by a byte slice.
+// The size of leafHashes should match the proof range i.e., end-start.
 // If the verifyCompleteness parameter is set to true, the function also checks
 // the completeness of the proof by verifying that there is no leaf in the
 // tree represented by the root parameter that matches the namespace ID nID
@@ -260,6 +264,15 @@ func (proof Proof) VerifyLeafHashes(nth *NmtHasher, verifyCompleteness bool, nID
 		return false, fmt.Errorf("proof range [proof.start=%d, proof.end=%d) is not valid: %w", proof.Start(), proof.End(), ErrInvalidRange)
 	}
 
+	// check whether the number of leaves match the proof range i.e., end-start.
+	// If not, make an early return.
+	expectedLeafHashesCount := proof.End() - proof.Start()
+	if len(leafHashes) != expectedLeafHashesCount {
+		return false, fmt.Errorf(
+			"supplied leafHashes size  %d, expected size %d: %w",
+			len(leafHashes), expectedLeafHashesCount, ErrWrongLeafHashesSize)
+	}
+
 	// perform some consistency checks:
 	if nID.Size() != nth.NamespaceSize() {
 		return false, fmt.Errorf("namespace ID size (%d) does not match the namespace size of the NMT hasher (%d)", nID.Size(), nth.NamespaceSize())
@@ -395,6 +408,7 @@ func (proof Proof) VerifyLeafHashes(nth *NmtHasher, verifyCompleteness bool, nID
 // and the provided proof to regenerate and compare the root. Note that the leavesWithoutNamespace data should not contain the prefixed namespace, unlike the tree.Push method,
 // which takes prefixed data. All leaves implicitly have the same namespace ID:
 // `nid`.
+// The size of the leavesWithoutNamespace should be equal to the proof range i.e., end-start.
 // VerifyInclusion does not verify the completeness of the proof, so it's possible for leavesWithoutNamespace to be a subset of the leaves in the tree that have the namespace ID nid.
 func (proof Proof) VerifyInclusion(h hash.Hash, nid namespace.ID, leavesWithoutNamespace [][]byte, root []byte) bool {
 	// check the range of the proof
@@ -411,6 +425,13 @@ func (proof Proof) VerifyInclusion(h hash.Hash, nid namespace.ID, leavesWithoutN
 		return false
 	}
 
+	// check whether the number of leavesWithoutNamespace match the proof range i.e., end-start.
+	// If not, make an early return.
+	expectedLeavesCount := proof.End() - proof.Start()
+	if len(leavesWithoutNamespace) != expectedLeavesCount {
+		return false
+	}
+
 	nth := NewNmtHasher(h, nid.Size(), proof.isMaxNamespaceIDIgnored)
 
 	// perform some consistency checks:


Test patch:
diff --git a/proof_test.go b/proof_test.go
index 2dba96f..235a403 100644
--- a/proof_test.go
+++ b/proof_test.go
@@ -546,6 +546,168 @@ func TestVerifyNamespace_False(t *testing.T) {
 	}
 }
 
+func TestVerifyInclusion_MismatchingRange(t *testing.T) {
+	nIDs := []byte{1, 2, 3, 4, 6, 6, 6, 9}
+	nmt := exampleNMT(1, true, nIDs...)
+	root, err := nmt.Root()
+	require.NoError(t, err)
+
+	nid6 := namespace.ID{6}
+	// node at index 5 has namespace ID 6
+	incProof6, err := nmt.ProveNamespace(nid6)
+	require.NoError(t, err)
+	// leaves with namespace ID 6
+	leaf4 := nmt.leaves[4][nmt.NamespaceSize():]
+	leaf5 := nmt.leaves[5][nmt.NamespaceSize():]
+	leaf6 := nmt.leaves[6][nmt.NamespaceSize():]
+
+	type args struct {
+		nIDSize                namespace.IDSize
+		nID                    namespace.ID
+		leavesWithoutNamespace [][]byte
+		root                   []byte
+	}
+	tests := []struct {
+		name   string
+		proof  Proof
+		args   args
+		result bool
+	}{
+		{
+			"inclusion proof: size of proof's range = size of leavesWithoutNamespace",
+			incProof6,
+			args{1, nid6, [][]byte{leaf4, leaf5, leaf6}, root},
+			true,
+		},
+		{
+			"inclusion proof: size of proof's range > size of" +
+				" a non-empty leavesWithoutNamespace",
+			incProof6,
+			args{1, nid6, [][]byte{leaf4, leaf5}, root},
+			false,
+		},
+		{
+			"inclusion proof: size of proof's range > size of" +
+				" an empty leavesWithoutNamespace",
+			incProof6,
+			args{1, nid6, [][]byte{}, root},
+			false,
+		},
+		{
+			"inclusion proof: size of proof's range < size of" +
+				" leavesWithoutNamespace",
+			incProof6,
+			args{1, nid6, [][]byte{leaf4, leaf5, leaf6, leaf6}, root},
+			false,
+		},
+		{
+			// in this testcase the nameID does not really matter since the
+			// leaves are empty
+			"empty proof: size of proof's range = size of leavesWithoutNamespace",
+			Proof{start: 1, end: 1},
+			args{1, nid6, [][]byte{}, root},
+			true,
+		},
+		{
+			"empty proof: size of proof's range < size of" +
+				" leavesWithoutNamespace",
+			Proof{start: 1, end: 1},
+			args{1, nid6, [][]byte{leaf4, leaf5, leaf6}, root},
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			hasher := sha256.New()
+			got := tt.proof.VerifyInclusion(hasher, tt.args.nID,
+				tt.args.leavesWithoutNamespace, tt.args.root)
+			assert.Equal(t, tt.result, got)
+		})
+	}
+}
+
+func TestVerifyLeafHashes_MismatchingRange(t *testing.T) {
+	nIDs := []byte{1, 2, 3, 4, 6, 6, 6, 9}
+	nmt := exampleNMT(1, true, nIDs...)
+	root, err := nmt.Root()
+	require.NoError(t, err)
+
+	nid5 := namespace.ID{5}
+	// namespace 5 does not exist in the tree, hence the proof is an absence proof
+	absenceProof5, err := nmt.ProveNamespace(nid5)
+	require.NoError(t, err)
+	leafHash5 := nmt.leafHashes[4]
+
+	nid6 := namespace.ID{6}
+	// node at index 5 has namespace ID 6
+	incProof6, err := nmt.Prove(5)
+	require.NoError(t, err)
+	leafHash6 := nmt.leafHashes[5]
+
+	type args struct {
+		nIDSize    namespace.IDSize
+		nID        namespace.ID
+		leafHashes [][]byte
+		root       []byte
+	}
+	tests := []struct {
+		name   string
+		proof  Proof
+		args   args
+		result bool
+		err    error
+	}{
+		{
+			"absence proof: size of proof's range = size of leafHashes",
+			absenceProof5,
+			args{1, namespace.ID{5}, [][]byte{leafHash5}, root},
+			true, nil,
+		},
+		{
+			"absence proof: size of proof's range > size of leafHashes",
+			absenceProof5,
+			args{1, nid5, [][]byte{}, root},
+			false, ErrWrongLeafHashesSize,
+		},
+		{
+			"absence proof: size of proof's range < size of leafHashes",
+			absenceProof5,
+			args{1, nid5, [][]byte{leafHash5, leafHash5}, root},
+			false, ErrWrongLeafHashesSize,
+		},
+		{
+			"inclusion proof: size of proof's range = size of leafHashes",
+			incProof6,
+			args{1, nid6, [][]byte{leafHash6}, root},
+			true, nil,
+		},
+		{
+			"inclusion proof: size of proof's range > size of leafHashes",
+			incProof6,
+			args{1, nid6, [][]byte{}, root},
+			false, ErrWrongLeafHashesSize,
+		},
+		{
+			"inclusion proof: size of proof's range < size of leafHashes",
+			incProof6,
+			args{1, nid6, [][]byte{leafHash6, leafHash6}, root},
+			false,
+			ErrWrongLeafHashesSize,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			hasher := NewNmtHasher(sha256.New(), tt.args.nIDSize, true)
+			got, err := tt.proof.VerifyLeafHashes(hasher, false, tt.args.nID,
+				tt.args.leafHashes, tt.args.root)
+			assert.Equal(t, tt.result, got)
+			if tt.err != nil {
+				assert.ErrorAs(t, err, &tt.err)
+			}
+		})
+	}
+}
+
 func TestVerifyLeafHashes_False(t *testing.T) {
 	nIDs := []byte{1, 2, 3, 4, 6, 7, 8, 9}
 
@@ -594,7 +756,12 @@ func TestVerifyLeafHashes_False(t *testing.T) {
 		args   args
 		result bool
 	}{
-		{"nID size of proof < nID size of VerifyLeafHashes' nmt hasher", proof4_1, args{2, nid4_2, [][]byte{leafHash2}, root2}, false},
+		{
+			"nID size of proof < nID size of VerifyLeafHashes' nmt hasher",
+			proof4_1,
+			args{2, nid4_2, [][]byte{leafHash2}, root2},
+			false,
+		},
 		{"nID size of proof > nID size of VerifyLeafHashes' nmt hasher", proof4_2, args{1, nid4_1, [][]byte{leafHash1}, root1}, false},
 		{"nID size of root < nID size of VerifyLeafHashes' nmt hasher", proof4_2, args{2, nid4_2, [][]byte{leafHash2}, root1}, false},
 		{"nID size of root > nID size of VerifyLeafHashes' nmt hasher", proof4_1, args{1, nid4_1, [][]byte{leafHash1}, root2}, false},