From cbdf0b962e89a68064715dcb1c0e90ce98894392 Mon Sep 17 00:00:00 2001
From: h5law <h5law.curve514@passfwd.com>
Date: Wed, 20 Mar 2024 16:31:15 +0000
Subject: [PATCH] feat: export hasher sizes from the trie spec and enforce path
 length in closest proof method and verification

---
 proofs.go |  5 +++++
 smt.go    |  9 +++++++++
 types.go  | 13 +++++++++++--
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/proofs.go b/proofs.go
index 3a2b690..dabdc10 100644
--- a/proofs.go
+++ b/proofs.go
@@ -205,6 +205,11 @@ func (proof *SparseMerkleClosestProof) GetValueHash(spec *TrieSpec) []byte {
 }
 
 func (proof *SparseMerkleClosestProof) validateBasic(spec *TrieSpec) error {
+	// ensure the path used in the proof is within the path size
+	if len(proof.Path) > spec.ph.PathSize() {
+		return fmt.Errorf("invalid path length: got %d, max is %d", len(proof.Path), spec.ph.PathSize())
+	}
+
 	// 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)
diff --git a/smt.go b/smt.go
index f55a3f8..838a9e9 100644
--- a/smt.go
+++ b/smt.go
@@ -2,6 +2,7 @@ package smt
 
 import (
 	"bytes"
+	"fmt"
 	"hash"
 
 	"github.com/pokt-network/smt/kvstore"
@@ -424,6 +425,14 @@ func (smt *SMT) ProveClosest(path []byte) (
 	proof *SparseMerkleClosestProof, // proof of the key-value pair found
 	err error, // the error value encountered
 ) {
+	// ensure the path used in the proof is within the path size
+	if len(proof.Path) > smt.Spec().PathHasherSize() {
+		return nil, fmt.Errorf(
+			"invalid path length: got %d, max is %d",
+			len(proof.Path), smt.Spec().PathHasherSize(),
+		)
+	}
+
 	workingPath := make([]byte, len(path))
 	copy(workingPath, path)
 	var siblings []trieNode
diff --git a/types.go b/types.go
index 150d41f..c92f01e 100644
--- a/types.go
+++ b/types.go
@@ -82,8 +82,8 @@ type SparseMerkleSumTrie interface {
 // leaf paths and stored values, and the corresponding maximum trie depth.
 type TrieSpec struct {
 	th      trieHasher
-	ph      PathHasher
-	vh      ValueHasher
+	ph      *pathHasher
+	vh      *valueHasher
 	sumTrie bool
 }
 
@@ -98,6 +98,15 @@ func newTrieSpec(hasher hash.Hash, sumTrie bool) TrieSpec {
 // Spec returns the TrieSpec associated with the given trie
 func (spec *TrieSpec) Spec() *TrieSpec { return spec }
 
+// PathHasherSize returns the size of the path hasher
+func (spec *TrieSpec) PathHasherSize() int { return spec.ph.PathSize() }
+
+// ValueHasherSize returns the size of the value hasher
+func (spec *TrieSpec) ValueHasherSize() int { return spec.th.hashSize() }
+
+// TrieHasherSize returns the size of the trie hasher
+func (spec *TrieSpec) TrieHasherSize() int { return spec.th.hasher.Size() }
+
 func (spec *TrieSpec) depth() int { return spec.ph.PathSize() * 8 }
 func (spec *TrieSpec) digestValue(data []byte) []byte {
 	if spec.vh == nil {