From 6fa739f145d5b80cc845b390229bc8c50689029d Mon Sep 17 00:00:00 2001
From: zjubfd <296179868@qq.com>
Date: Sat, 8 Oct 2022 13:18:51 +0800
Subject: [PATCH]  fix: add ProofOp verifier to enhancement the check

---
 crypto/merkle/proof.go      | 26 +++++++++++++++++---------
 crypto/merkle/proof_test.go | 14 ++++++++++++++
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/crypto/merkle/proof.go b/crypto/merkle/proof.go
index 5e2a3ab12bc..b4fca168a9f 100644
--- a/crypto/merkle/proof.go
+++ b/crypto/merkle/proof.go
@@ -30,17 +30,25 @@ type ProofOperator interface {
 // and the last Merkle root will be verified with already known data
 type ProofOperators []ProofOperator
 
-func (poz ProofOperators) VerifyValue(root []byte, keypath string, value []byte) (err error) {
-	return poz.Verify(root, keypath, [][]byte{value})
+type ProofOpVerifier func(ProofOperator) error
+
+func (poz ProofOperators) VerifyValue(root []byte, keypath string, value []byte, verifiers ...ProofOpVerifier) (err error) {
+	return poz.Verify(root, keypath, [][]byte{value}, verifiers...)
 }
 
-func (poz ProofOperators) Verify(root []byte, keypath string, args [][]byte) (err error) {
+func (poz ProofOperators) Verify(root []byte, keypath string, args [][]byte, verifiers ...ProofOpVerifier) (err error) {
 	keys, err := KeyPathToKeys(keypath)
 	if err != nil {
 		return
 	}
 
 	for i, op := range poz {
+		for _, verifier := range verifiers {
+			if err := verifier(op); err != nil {
+				return err
+			}
+		}
+
 		key := op.GetKey()
 		if len(key) != 0 {
 			if len(keys) == 0 {
@@ -109,22 +117,22 @@ func (prt *ProofRuntime) DecodeProof(proof *Proof) (ProofOperators, error) {
 	return poz, nil
 }
 
-func (prt *ProofRuntime) VerifyValue(proof *Proof, root []byte, keypath string, value []byte) (err error) {
-	return prt.Verify(proof, root, keypath, [][]byte{value})
+func (prt *ProofRuntime) VerifyValue(proof *Proof, root []byte, keypath string, value []byte, verifiers ...ProofOpVerifier) (err error) {
+	return prt.Verify(proof, root, keypath, [][]byte{value}, verifiers...)
 }
 
 // TODO In the long run we'll need a method of classifcation of ops,
 // whether existence or absence or perhaps a third?
-func (prt *ProofRuntime) VerifyAbsence(proof *Proof, root []byte, keypath string) (err error) {
-	return prt.Verify(proof, root, keypath, nil)
+func (prt *ProofRuntime) VerifyAbsence(proof *Proof, root []byte, keypath string, verifiers ...ProofOpVerifier) (err error) {
+	return prt.Verify(proof, root, keypath, nil, verifiers...)
 }
 
-func (prt *ProofRuntime) Verify(proof *Proof, root []byte, keypath string, args [][]byte) (err error) {
+func (prt *ProofRuntime) Verify(proof *Proof, root []byte, keypath string, args [][]byte, verifiers ...ProofOpVerifier) (err error) {
 	poz, err := prt.DecodeProof(proof)
 	if err != nil {
 		return cmn.ErrorWrap(err, "decoding proof")
 	}
-	return poz.Verify(root, keypath, args)
+	return poz.Verify(root, keypath, args, verifiers...)
 }
 
 // DefaultProofRuntime only knows about Simple value
diff --git a/crypto/merkle/proof_test.go b/crypto/merkle/proof_test.go
index 4de3246f1a1..15debfb3949 100644
--- a/crypto/merkle/proof_test.go
+++ b/crypto/merkle/proof_test.go
@@ -1,6 +1,7 @@
 package merkle
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -134,6 +135,19 @@ func TestProofOperators(t *testing.T) {
 	popz = []ProofOperator{}
 	err = popz.Verify(bz("OUTPUT4"), "/KEY4/KEY2/KEY1", [][]byte{bz("INPUT1")})
 	assert.NotNil(t, err)
+
+	// Test OP Verifier 1
+	popz = []ProofOperator{op1, op2, op3, op4}
+	err = popz.VerifyValue(bz("OUTPUT4"), "/KEY4/KEY2/KEY1", bz("INPUT1"), func(operator ProofOperator) error {
+		return nil
+	})
+	assert.Nil(t, err)
+
+	err = popz.VerifyValue(bz("OUTPUT4"), "/KEY4/KEY2/KEY1", bz("INPUT1"), func(operator ProofOperator) error {
+		return fmt.Errorf("suspend")
+	})
+	assert.EqualError(t, err, "suspend")
+
 }
 
 func bz(s string) []byte {