From 9ea65a34f8312d280f292e10dbfb82155683c26f Mon Sep 17 00:00:00 2001
From: h5law <h5law.curve514@passfwd.com>
Date: Wed, 20 Mar 2024 16:59:52 +0000
Subject: [PATCH] feat: enforce path length check in ClosestProof and its
 validation

---
 errors.go |  3 +++
 proofs.go | 12 ++++++++++++
 smt.go    |  5 +++++
 3 files changed, 20 insertions(+)

diff --git a/errors.go b/errors.go
index b8d6bbe..8b89097 100644
--- a/errors.go
+++ b/errors.go
@@ -9,4 +9,7 @@ var (
 	ErrBadProof = errors.New("bad proof")
 	// ErrKeyNotFound is returned when a key is not found in the tree.
 	ErrKeyNotFound = errors.New("key not found")
+	// ErrInvalidClosestPath is returned when the path used in the ClosestProof
+	// method does not match the size of the trie's PathHasher
+	ErrInvalidClosestPath = errors.New("invalid path does not match path hasher size")
 )
diff --git a/proofs.go b/proofs.go
index 3a2b690..12d3de8 100644
--- a/proofs.go
+++ b/proofs.go
@@ -205,6 +205,12 @@ func (proof *SparseMerkleClosestProof) GetValueHash(spec *TrieSpec) []byte {
 }
 
 func (proof *SparseMerkleClosestProof) validateBasic(spec *TrieSpec) error {
+	// ensure the proof length is the same size (in bytes) as the path
+	// hasher of the spec provided
+	if len(proof.Path) != spec.PathHasherSize() {
+		return fmt.Errorf("invalid path length: got %d, want %d", len(proof.Path), spec.PathHasherSize())
+	}
+
 	// ensure the depth of the leaf node being proven is within the path size
 	if proof.Depth < 0 || proof.Depth > spec.ph.PathSize()*8 {
 		return fmt.Errorf("invalid depth: got %d, outside of [0, %d]", proof.Depth, spec.ph.PathSize()*8)
@@ -250,6 +256,12 @@ type SparseCompactMerkleClosestProof struct {
 }
 
 func (proof *SparseCompactMerkleClosestProof) validateBasic(spec *TrieSpec) error {
+	// Ensure the proof length is the same size (in bytes) as the path
+	// hasher of the spec provided
+	if len(proof.Path) != spec.PathHasherSize() {
+		return fmt.Errorf("invalid path length: got %d, want %d", len(proof.Path), spec.PathHasherSize())
+	}
+
 	// Do a basic sanity check on the proof on the fields of the proof specific to
 	// the compact proof only.
 	//
diff --git a/smt.go b/smt.go
index f55a3f8..7928c40 100644
--- a/smt.go
+++ b/smt.go
@@ -424,6 +424,11 @@ func (smt *SMT) ProveClosest(path []byte) (
 	proof *SparseMerkleClosestProof, // proof of the key-value pair found
 	err error, // the error value encountered
 ) {
+	// Ensure the path provided is the correct length for the path hasher.
+	if len(path) != smt.Spec().PathHasherSize() {
+		return nil, ErrInvalidClosestPath
+	}
+
 	workingPath := make([]byte, len(path))
 	copy(workingPath, path)
 	var siblings []trieNode