From e8f819d8a0aeb12820e1278cd51815c346dfb96b Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Tue, 30 Jun 2020 22:15:00 +0200
Subject: [PATCH] Add Proof verification

---
 merkletree.go      | 61 ++++++++++++++++++++++++++++++++++++++++++++++
 merkletree_test.go |  6 +++++
 2 files changed, 67 insertions(+)

diff --git a/merkletree.go b/merkletree.go
index 8aa3d49..17e5b53 100644
--- a/merkletree.go
+++ b/merkletree.go
@@ -3,6 +3,7 @@ package merkletree
 import (
 	"bytes"
 	"errors"
+	"fmt"
 	"math/big"
 	"sync"
 
@@ -430,3 +431,63 @@ func (mt *MerkleTree) GenerateProof(k *big.Int, rootKey *Hash) (*Proof, error) {
 	}
 	return nil, ErrEntryIndexNotFound
 }
+
+// VerifyProof verifies the Merkle Proof for the entry and root.
+func VerifyProof(rootKey *Hash, proof *Proof, k, v *big.Int) bool {
+	rootFromProof, err := RootFromProof(proof, k, v)
+	if err != nil {
+		return false
+	}
+	return bytes.Equal(rootKey[:], rootFromProof[:])
+}
+
+// RootFromProof calculates the root that would correspond to a tree whose
+// siblings are the ones in the proof with the claim hashing to hIndex and
+// hValue.
+func RootFromProof(proof *Proof, k, v *big.Int) (*Hash, error) {
+	kHash := NewHashFromBigInt(k)
+	vHash := NewHashFromBigInt(v)
+	sibIdx := len(proof.Siblings) - 1
+	var err error
+	var midKey *Hash
+	if proof.Existence {
+		midKey, err = LeafKey(kHash, vHash)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		if proof.NodeAux == nil {
+			midKey = &HashZero
+		} else {
+			if bytes.Equal(kHash[:], proof.NodeAux.Key[:]) {
+				return nil, fmt.Errorf("Non-existence proof being checked against hIndex equal to nodeAux")
+			}
+			midKey, err = LeafKey(proof.NodeAux.Key, proof.NodeAux.Value)
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+	path := getPath(int(proof.depth), kHash[:])
+	var siblingKey *Hash
+	for lvl := int(proof.depth) - 1; lvl >= 0; lvl-- {
+		if common.TestBitBigEndian(proof.notempties[:], uint(lvl)) {
+			siblingKey = proof.Siblings[sibIdx]
+			sibIdx--
+		} else {
+			siblingKey = &HashZero
+		}
+		if path[lvl] {
+			midKey, err = NewNodeMiddle(siblingKey, midKey).Key()
+			if err != nil {
+				return nil, err
+			}
+		} else {
+			midKey, err = NewNodeMiddle(midKey, siblingKey).Key()
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+	return midKey, nil
+}
diff --git a/merkletree_test.go b/merkletree_test.go
index ae923fb..2a23c08 100644
--- a/merkletree_test.go
+++ b/merkletree_test.go
@@ -25,4 +25,10 @@ func TestNewTree(t *testing.T) {
 	err = mt.Add(big.NewInt(1234), big.NewInt(9876))
 	assert.Nil(t, err)
 	assert.Equal(t, "16970503620176669663662021947486532860010370357132361783766545149750777353066", mt.Root().BigInt().String())
+
+	proof, err := mt.GenerateProof(big.NewInt(33), nil)
+	assert.Nil(t, err)
+
+	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(44)))
+	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(45)))
 }