Repository: celestiaorg/nmt (go)

Issue:
VerifyLeafHashes can panic in the completeness check if proof.nodes are not namespaced
# Problem
Link to the completeness check: https://github.com/celestiaorg/nmt/blob/master/proof.go#L199
The proof.nodes length are not checked before extracting their namespace IDs. This may cause panic.

# Acceptance Criteria
Call ValidateNode() on the proof.nodes at the beginning of the VerifyLeafHashes and make an early return if they are invalid. 


Pull request:
feat: checks that all the arguments of VerifyNamespace conform to the same namespace hash format
## Overview

Closes #141 

## Checklist
- [x] New and updated code has appropriate documentation
- [x] New and updated code has new and/or updated testing
- [x] Required CI checks are passing
- [x] Visual proof for any user facing features like CLI or documentation updates
- [x] Linked issues closed with keywords


Patch:
diff --git a/proof.go b/proof.go
index 9916b6a..32d8480 100644
--- a/proof.go
+++ b/proof.go
@@ -126,28 +126,43 @@ func NewAbsenceProof(proofStart, proofEnd int, proofNodes [][]byte, leafHash []b
 //	`end-1` of the tree.
 //
 // `root` is the root of the NMT against which the `proof` is verified.
-func (proof Proof) VerifyNamespace(h hash.Hash, nID namespace.ID, data [][]byte, root []byte) bool {
-	nth := NewNmtHasher(h, nID.Size(), proof.isMaxNamespaceIDIgnored)
-	min := namespace.ID(MinNamespace(root, nID.Size()))
-	max := namespace.ID(MaxNamespace(root, nID.Size()))
-	if nID.Size() != min.Size() || nID.Size() != max.Size() {
-		// conflicting namespace sizes
+func (proof Proof) VerifyNamespace(h hash.Hash, nID namespace.ID, leaves [][]byte, root []byte) bool {
+	nIDLen := nID.Size()
+	nth := NewNmtHasher(h, nIDLen, proof.isMaxNamespaceIDIgnored)
+
+	// perform some consistency checks:
+	// check that the root is valid w.r.t the NMT hasher
+	if err := nth.ValidateNodeFormat(root); err != nil {
 		return false
 	}
+	// check that all the proof.nodes are valid w.r.t the NMT hasher
+	for _, node := range proof.nodes {
+		if err := nth.ValidateNodeFormat(node); err != nil {
+			return false
+		}
+	}
+
+	// if the proof is an absence proof, the leafHash must be valid w.r.t the NMT hasher
+	if proof.IsOfAbsence() {
+		if err := nth.ValidateNodeFormat(proof.leafHash); err != nil {
+			return false
+		}
+	}
 
 	isEmptyRange := proof.start == proof.end
-	if len(data) == 0 && isEmptyRange && len(proof.nodes) == 0 {
+	if len(leaves) == 0 && isEmptyRange && len(proof.nodes) == 0 {
 		// empty proofs are always rejected unless nID is outside the range of
 		// namespaces covered by the root we special case the empty root, since
 		// it purports to cover the zero namespace but does not actually include
 		// any such nodes
+		min := namespace.ID(MinNamespace(root, nIDLen))
+		max := namespace.ID(MaxNamespace(root, nIDLen))
 		if nID.Less(min) || max.Less(nID) || bytes.Equal(root, nth.EmptyRoot()) {
 			return true
 		}
 		return false
 	}
-	gotLeafHashes := make([][]byte, 0, len(data))
-	nIDLen := nID.Size()
+	gotLeafHashes := make([][]byte, 0, len(leaves))
 	if proof.IsOfAbsence() {
 		gotLeafHashes = append(gotLeafHashes, proof.leafHash)
 		// conduct some sanity checks:
@@ -160,7 +175,7 @@ func (proof Proof) VerifyNamespace(h hash.Hash, nID namespace.ID, data [][]byte,
 	} else {
 		// collect leaf hashes from provided data and do some sanity checks:
 		hashLeafFunc := nth.HashLeaf
-		for _, gotLeaf := range data {
+		for _, gotLeaf := range leaves {
 			if nth.ValidateLeaf(gotLeaf) != nil {
 				return false
 			}
@@ -179,7 +194,8 @@ func (proof Proof) VerifyNamespace(h hash.Hash, nID namespace.ID, data [][]byte,
 	}
 	// check whether the number of leaves match the proof range i.e., end-start.
 	// If not, make an early return.
-	if !proof.IsOfAbsence() && len(gotLeafHashes) != (proof.End()-proof.Start()) {
+	expectedLeafCount := proof.End() - proof.Start()
+	if !proof.IsOfAbsence() && len(gotLeafHashes) != expectedLeafCount {
 		return false
 	}
 	// with verifyCompleteness set to true:
@@ -200,6 +216,27 @@ func (proof Proof) VerifyNamespace(h hash.Hash, nID namespace.ID, data [][]byte,
 // tree represented by the root parameter that matches the namespace ID nID
 // but is not present in the leafHashes list.
 func (proof Proof) verifyLeafHashes(nth *Hasher, verifyCompleteness bool, nID namespace.ID, leafHashes [][]byte, root []byte) (bool, error) {
+	// 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())
+	}
+	// check that the root is valid w.r.t the NMT hasher
+	if err := nth.ValidateNodeFormat(root); err != nil {
+		return false, fmt.Errorf("root does not match the NMT hasher's hash format: %w", err)
+	}
+	// check that all the proof.nodes are valid w.r.t the NMT hasher
+	for _, node := range proof.nodes {
+		if err := nth.ValidateNodeFormat(node); err != nil {
+			return false, fmt.Errorf("proof nodes do not match the NMT hasher's hash format: %w", err)
+		}
+	}
+	// check that all the proof.nodes are valid w.r.t the NMT hasher
+	for _, leaf := range leafHashes {
+		if err := nth.ValidateNodeFormat(leaf); err != nil {
+			return false, fmt.Errorf("leaf hash does not match the NMT hasher's hash format: %w", err)
+		}
+	}
+
 	var leafIndex uint64
 	// leftSubtrees is to be populated by the subtree roots upto [0, r.Start)
 	leftSubtrees := make([][]byte, 0, len(proof.nodes))


Test patch:
diff --git a/proof_test.go b/proof_test.go
index d4afec6..6cd3f0e 100644
--- a/proof_test.go
+++ b/proof_test.go
@@ -243,9 +243,9 @@ func TestVerifyLeafHashes_Err(t *testing.T) {
 		root               []byte
 		wantErr            bool
 	}{
-		{" wrong leafHash: not namespaced", proof5, hasher, true, nID5, [][]byte{leafHash5}, root, true},
-		{" wrong leafHash: incorrect namespace", proof5, hasher, true, nID5, [][]byte{{10, 10, 10, 10}}, root, true},
-		{" wrong proof.nodes: the last node has an incorrect format", proof4, hasher, false, nID4, [][]byte{leafHash4}, root, true},
+		{"wrong leafHash: not namespaced", proof5, hasher, true, nID5, [][]byte{leafHash5}, root, true},
+		{"wrong leafHash: incorrect namespace", proof5, hasher, true, nID5, [][]byte{{10, 10, 10, 10}}, root, true},
+		{"wrong proof.nodes: the last node has an incorrect format", proof4, hasher, false, nID4, [][]byte{leafHash4}, root, true},
 		//  the verifyCompleteness parameter in the verifyProof function should be set to false in order to bypass nodes correctness check during the completeness verification (otherwise it panics).
 	}
 	for _, tt := range tests {
@@ -306,3 +306,122 @@ func TestVerifyInclusion_False(t *testing.T) {
 		})
 	}
 }
+
+func TestVerifyNamespace_False(t *testing.T) {
+	nIDs := []byte{1, 2, 3, 4, 5, 6, 7, 8, 11}
+
+	// create a sample tree with namespace ID size of 1
+	nmt1 := exampleNMT(1, nIDs...)
+	root1, err := nmt1.Root()
+	require.NoError(t, err)
+	nid4_1 := namespace.ID{4}
+	proof4_1, err := nmt1.ProveNamespace(nid4_1) // leaf at index 3 has namespace ID 4
+	require.NoError(t, err)
+
+	// create a sample tree with namespace ID size of 2
+	nmt2 := exampleNMT(2, nIDs...)
+	root2, err := nmt2.Root()
+	require.NoError(t, err)
+	nid4_2 := namespace.ID{4, 4}
+	proof4_2, err := nmt2.ProveNamespace(nid4_2) // leaf at index 3 has namespace ID 4
+	require.NoError(t, err)
+
+	leaf := nmt1.leaves[3]
+
+	// create an absence proof with namespace ID size of 1
+	nid9_1 := namespace.ID{9}
+	absenceProof9_1, err := nmt1.ProveNamespace(nid9_1)
+	require.NoError(t, err)
+	require.True(t, absenceProof9_1.IsOfAbsence())
+
+	// create an absence proof with namespace ID size of 2
+	nid9_2 := namespace.ID{9, 9}
+	absenceProof9_2, err := nmt2.ProveNamespace(nid9_2)
+	require.NoError(t, err)
+	require.True(t, absenceProof9_2.IsOfAbsence())
+
+	// swap leafHashes of the absence proofs
+	buffer := absenceProof9_2.leafHash
+	absenceProof9_2.leafHash = absenceProof9_1.leafHash
+	absenceProof9_1.leafHash = buffer
+
+	hasher := sha256.New()
+
+	type args struct {
+		hasher hash.Hash
+		nID    namespace.ID
+		leaves [][]byte
+		root   []byte
+	}
+	tests := []struct {
+		name   string
+		proof  Proof
+		args   args
+		result bool
+	}{
+		{"nID size of proof < nID size of VerifyNamespace's nmt hasher", proof4_1, args{hasher, nid4_2, [][]byte{leaf}, root2}, false},
+		{"nID size of proof > nID size of VerifyNamespace's nmt hasher", proof4_2, args{hasher, nid4_1, [][]byte{leaf}, root1}, false},
+		{"nID size of root < nID size of VerifyNamespace's nmt hasher", proof4_2, args{hasher, nid4_2, [][]byte{leaf}, root1}, false},
+		{"nID size of root > nID size of VerifyNamespace's nmt hasher", proof4_1, args{hasher, nid4_1, [][]byte{leaf}, root2}, false},
+		{"nID size of proof.leafHash < nID size of VerifyNamespace's nmt hasher", absenceProof9_2, args{hasher, nid9_2, [][]byte{}, root2}, false},
+		{"nID size of proof.leafHash > nID size of VerifyNamespace's nmt hasher", absenceProof9_1, args{hasher, nid9_1, [][]byte{}, root1}, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := tt.proof.VerifyNamespace(tt.args.hasher, tt.args.nID, tt.args.leaves, tt.args.root)
+			assert.Equal(t, tt.result, got)
+		})
+	}
+}
+
+func TestVerifyLeafHashes_False(t *testing.T) {
+	nIDs := []byte{1, 2, 3, 4, 5, 6, 7, 8}
+
+	// create a sample tree with namespace ID size of 1
+	nmt1 := exampleNMT(1, nIDs...)
+	root1, err := nmt1.Root()
+	require.NoError(t, err)
+	nid4_1 := namespace.ID{4}
+	proof4_1, err := nmt1.ProveNamespace(nid4_1) // leaf at index 3 has namespace ID 4
+	require.NoError(t, err)
+
+	// create a sample tree with namespace ID size of 2
+	nmt2 := exampleNMT(2, nIDs...)
+	root2, err := nmt2.Root()
+	require.NoError(t, err)
+	nid4_2 := namespace.ID{4, 4}
+	proof4_2, err := nmt2.ProveNamespace(nid4_2) // leaf at index 3 has namespace ID 4
+	require.NoError(t, err)
+
+	leafHash1 := nmt1.leafHashes[3]
+	leafHash2 := nmt2.leafHashes[3]
+
+	type args struct {
+		nIDSize namespace.IDSize
+		nID     namespace.ID
+		leaves  [][]byte
+		root    []byte
+	}
+	tests := []struct {
+		name   string
+		proof  Proof
+		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_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},
+		{"size of queried nID > nID size of verifyLeafHashes' nmt hasher", proof4_1, args{1, nid4_2, [][]byte{leafHash1}, root1}, false},
+		{"size of queried nID < nID size of verifyLeafHashes' nmt hasher", proof4_2, args{2, nid4_1, [][]byte{leafHash2}, root2}, false},
+		{"nID size of leafHash < nID size of verifyLeafHashes' nmt hasher", proof4_2, args{2, nid4_2, [][]byte{leafHash1}, root2}, false},
+		{"nID size of leafHash > nID size of verifyLeafHashes' nmt hasher", proof4_1, args{1, nid4_1, [][]byte{leafHash2}, root1}, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			hasher := NewNmtHasher(sha256.New(), tt.args.nIDSize, true)
+			got, _ := tt.proof.verifyLeafHashes(hasher, true, tt.args.nID, tt.args.leaves, tt.args.root)
+			assert.Equal(t, tt.result, got)
+		})
+	}
+}