package merkletree
|
|
|
|
import (
|
|
"bytes"
|
|
"errors"
|
|
|
|
"github.com/syndtr/goleveldb/leveldb"
|
|
)
|
|
|
|
const (
|
|
// EmptyNodeType indicates the type of an EmptyNodeValue Node
|
|
EmptyNodeType = 00
|
|
// NormalNodeType indicates the type of a middle Node
|
|
normalNodeType = 01
|
|
// FinalNodeType indicates the type of middle Node that is in an optimized branch, then in the value contains the value of the final leaf node of that branch
|
|
finalNodeType = 02
|
|
// ValueNodeType indicates the type of a value Node
|
|
valueNodeType = 03
|
|
// RootNodeType indicates the type of a root Node
|
|
rootNodeType = 04
|
|
)
|
|
|
|
var (
|
|
// ErrNodeAlreadyExists is an error that indicates that a node already exists in the merkletree database
|
|
ErrNodeAlreadyExists = errors.New("node already exists")
|
|
rootNodeValue = HashBytes([]byte("root"))
|
|
// EmptyNodeValue is a [32]byte EmptyNodeValue array, all to zero
|
|
EmptyNodeValue = Hash{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
|
|
)
|
|
|
|
// Hash used in this tree, is the [32]byte keccak()
|
|
type Hash [32]byte
|
|
|
|
// Value is the interface of a generic leaf, a key value object stored in the leveldb
|
|
type Value interface {
|
|
IndexLength() uint32 // returns the index length value
|
|
Bytes() []byte // returns the value in byte array representation
|
|
}
|
|
|
|
//MerkleTree struct with the main elements of the Merkle Tree
|
|
type MerkleTree struct {
|
|
// sync.RWMutex
|
|
storage *leveldb.DB
|
|
root Hash
|
|
numLevels int // Height of the Merkle Tree, number of levels
|
|
}
|
|
|
|
// New generates a new Merkle Tree
|
|
func New(storage *leveldb.DB, numLevels int) (*MerkleTree, error) {
|
|
var mt MerkleTree
|
|
mt.storage = storage
|
|
mt.numLevels = numLevels
|
|
var err error
|
|
_, _, rootHash, err := mt.Get(rootNodeValue)
|
|
if err != nil {
|
|
mt.root = EmptyNodeValue
|
|
err = mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
copy(mt.root[:], rootHash)
|
|
return &mt, nil
|
|
}
|
|
|
|
// Root returns the merkletree.Root
|
|
func (mt *MerkleTree) Root() Hash {
|
|
return mt.root
|
|
}
|
|
|
|
// NumLevels returns the merkletree.NumLevels
|
|
func (mt *MerkleTree) NumLevels() int {
|
|
return mt.numLevels
|
|
}
|
|
|
|
// Add adds the leaf to the MT
|
|
func (mt *MerkleTree) Add(v Value) error {
|
|
// add the leaf that we are adding
|
|
mt.Insert(HashBytes(v.Bytes()), valueNodeType, v.IndexLength(), v.Bytes())
|
|
|
|
hi := HashBytes(v.Bytes()[:v.IndexLength()])
|
|
path := getPath(mt.numLevels, hi)
|
|
|
|
nodeHash := mt.root
|
|
var siblings []Hash
|
|
for i := mt.numLevels - 2; i >= 0; i-- {
|
|
nodeType, indexLength, nodeBytes, err := mt.Get(nodeHash)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if nodeType == byte(finalNodeType) {
|
|
hiChild := HashBytes(nodeBytes[:indexLength])
|
|
pathChild := getPath(mt.numLevels, hiChild)
|
|
posDiff := comparePaths(pathChild, path)
|
|
if posDiff == -1 {
|
|
return ErrNodeAlreadyExists
|
|
}
|
|
finalNode1Hash := calcHashFromLeafAndLevel(posDiff, pathChild, HashBytes(nodeBytes))
|
|
mt.Insert(finalNode1Hash, finalNodeType, indexLength, nodeBytes)
|
|
finalNode2Hash := calcHashFromLeafAndLevel(posDiff, path, HashBytes(v.Bytes()))
|
|
mt.Insert(finalNode2Hash, finalNodeType, v.IndexLength(), v.Bytes())
|
|
// now the parent
|
|
var parentNode treeNode
|
|
if path[posDiff] {
|
|
parentNode = treeNode{
|
|
ChildL: finalNode1Hash,
|
|
ChildR: finalNode2Hash,
|
|
}
|
|
} else {
|
|
parentNode = treeNode{
|
|
ChildL: finalNode2Hash,
|
|
ChildR: finalNode1Hash,
|
|
}
|
|
}
|
|
siblings = append(siblings, getEmptiesBetweenIAndPosHash(mt, i, posDiff+1)...)
|
|
|
|
if mt.root, err = mt.replaceLeaf(siblings, path[posDiff+1:], parentNode.Ht(), normalNodeType, 0, parentNode.Bytes()); err != nil {
|
|
return err
|
|
}
|
|
mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
|
|
return nil
|
|
}
|
|
node := parseNodeBytes(nodeBytes)
|
|
var sibling Hash
|
|
if !path[i] {
|
|
nodeHash = node.ChildL
|
|
sibling = node.ChildR
|
|
} else {
|
|
nodeHash = node.ChildR
|
|
sibling = node.ChildL
|
|
}
|
|
siblings = append(siblings, sibling)
|
|
|
|
if bytes.Equal(nodeHash[:], EmptyNodeValue[:]) {
|
|
// if the node is EmptyNodeValue, the leaf data will go directly at that height, as a Final Node
|
|
if i == mt.numLevels-2 && bytes.Equal(siblings[len(siblings)-1][:], EmptyNodeValue[:]) {
|
|
// if the pt node is the unique in the tree, just put it into the root node
|
|
// this means to be in i==mt.NumLevels-2 && nodeHash==EmptyNodeValue
|
|
finalNodeHash := calcHashFromLeafAndLevel(i+1, path, HashBytes(v.Bytes()))
|
|
mt.Insert(finalNodeHash, finalNodeType, v.IndexLength(), v.Bytes())
|
|
mt.root = finalNodeHash
|
|
mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
|
|
return nil
|
|
}
|
|
finalNodeHash := calcHashFromLeafAndLevel(i, path, HashBytes(v.Bytes()))
|
|
if mt.root, err = mt.replaceLeaf(siblings, path[i:], finalNodeHash, finalNodeType, v.IndexLength(), v.Bytes()); err != nil {
|
|
return err
|
|
}
|
|
mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
|
|
return nil
|
|
}
|
|
}
|
|
|
|
var err error
|
|
mt.root, err = mt.replaceLeaf(siblings, path, HashBytes(v.Bytes()), valueNodeType, v.IndexLength(), v.Bytes())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
|
|
return nil
|
|
}
|
|
|
|
// GenerateProof generates the Merkle Proof from a given leafHash for the current root
|
|
func (mt *MerkleTree) GenerateProof(hi Hash) ([]byte, error) {
|
|
var empties [32]byte
|
|
|
|
path := getPath(mt.numLevels, hi)
|
|
var siblings []Hash
|
|
nodeHash := mt.root
|
|
|
|
for level := 0; level < mt.numLevels-1; level++ {
|
|
nodeType, indexLength, nodeBytes, err := mt.Get(nodeHash)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if nodeType == byte(finalNodeType) {
|
|
realValueInPos, err := mt.GetValueInPos(hi)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if bytes.Equal(realValueInPos[:], EmptyNodeValue[:]) {
|
|
// go until the path is different, then get the nodes between this FinalNode and the node in the diffPath, they will be the siblings of the merkle proof
|
|
leafHi := HashBytes(nodeBytes[:indexLength]) // hi of element that was in the end of the branch (the finalNode)
|
|
pathChild := getPath(mt.numLevels, leafHi)
|
|
|
|
// get the position where the path is different
|
|
posDiff := comparePaths(pathChild, path)
|
|
if posDiff == -1 {
|
|
return nil, ErrNodeAlreadyExists
|
|
}
|
|
|
|
if posDiff != mt.NumLevels()-1-level {
|
|
sibling := calcHashFromLeafAndLevel(posDiff, pathChild, HashBytes(nodeBytes))
|
|
setbitmap(empties[:], uint(mt.NumLevels()-2-posDiff))
|
|
siblings = append([]Hash{sibling}, siblings...)
|
|
}
|
|
|
|
}
|
|
break
|
|
}
|
|
node := parseNodeBytes(nodeBytes)
|
|
|
|
var sibling Hash
|
|
if !path[mt.numLevels-level-2] {
|
|
nodeHash = node.ChildL
|
|
sibling = node.ChildR
|
|
} else {
|
|
nodeHash = node.ChildR
|
|
sibling = node.ChildL
|
|
}
|
|
if !bytes.Equal(sibling[:], EmptyNodeValue[:]) {
|
|
setbitmap(empties[:], uint(level))
|
|
siblings = append([]Hash{sibling}, siblings...)
|
|
}
|
|
}
|
|
// merge empties and siblings
|
|
var mp []byte
|
|
mp = append(mp, empties[:]...)
|
|
for k := range siblings {
|
|
mp = append(mp, siblings[k][:]...)
|
|
}
|
|
return mp, nil
|
|
}
|
|
|
|
// GetValueInPos returns the merkletree value in the position of the Hash of the Index (Hi)
|
|
func (mt *MerkleTree) GetValueInPos(hi Hash) ([]byte, error) {
|
|
path := getPath(mt.numLevels, hi)
|
|
nodeHash := mt.root
|
|
for i := mt.numLevels - 2; i >= 0; i-- {
|
|
nodeType, indexLength, nodeBytes, err := mt.Get(nodeHash)
|
|
if err != nil {
|
|
return nodeBytes, err
|
|
}
|
|
if nodeType == byte(finalNodeType) {
|
|
// check if nodeBytes path is different of hi
|
|
index := nodeBytes[:indexLength]
|
|
nodeHi := HashBytes(index)
|
|
nodePath := getPath(mt.numLevels, nodeHi)
|
|
posDiff := comparePaths(path, nodePath)
|
|
// if is different, return an EmptyNodeValue, else return the nodeBytes
|
|
if posDiff != -1 {
|
|
return EmptyNodeValue[:], nil
|
|
}
|
|
return nodeBytes, nil
|
|
}
|
|
node := parseNodeBytes(nodeBytes)
|
|
if !path[i] {
|
|
nodeHash = node.ChildL
|
|
} else {
|
|
nodeHash = node.ChildR
|
|
}
|
|
}
|
|
_, _, valueBytes, err := mt.Get(nodeHash)
|
|
if err != nil {
|
|
return valueBytes, err
|
|
}
|
|
return valueBytes, nil
|
|
}
|
|
|
|
func calcHashFromLeafAndLevel(untilLevel int, path []bool, leafHash Hash) Hash {
|
|
nodeCurrLevel := leafHash
|
|
for i := 0; i < untilLevel; i++ {
|
|
if path[i] {
|
|
node := treeNode{
|
|
ChildL: EmptyNodeValue,
|
|
ChildR: nodeCurrLevel,
|
|
}
|
|
nodeCurrLevel = node.Ht()
|
|
} else {
|
|
node := treeNode{
|
|
ChildL: nodeCurrLevel,
|
|
ChildR: EmptyNodeValue,
|
|
}
|
|
nodeCurrLevel = node.Ht()
|
|
}
|
|
}
|
|
return nodeCurrLevel
|
|
}
|
|
|
|
func (mt *MerkleTree) replaceLeaf(siblings []Hash, path []bool, newLeafHash Hash, nodetype byte, indexLength uint32, newLeafValue []byte) (Hash, error) {
|
|
// add the new leaf
|
|
mt.Insert(newLeafHash, nodetype, indexLength, newLeafValue)
|
|
currNode := newLeafHash
|
|
// here the path is only the path[posDiff+1]
|
|
for i := 0; i < len(siblings); i++ {
|
|
if !path[i] {
|
|
node := treeNode{
|
|
ChildL: currNode,
|
|
ChildR: siblings[len(siblings)-1-i],
|
|
}
|
|
mt.Insert(node.Ht(), normalNodeType, 0, node.Bytes())
|
|
currNode = node.Ht()
|
|
} else {
|
|
|
|
node := treeNode{
|
|
ChildL: siblings[len(siblings)-1-i],
|
|
ChildR: currNode,
|
|
}
|
|
mt.Insert(node.Ht(), normalNodeType, 0, node.Bytes())
|
|
currNode = node.Ht()
|
|
}
|
|
}
|
|
|
|
return currNode, nil // currNode = root
|
|
}
|
|
|
|
// CheckProof validates the Merkle Proof for the leafHash and root
|
|
func CheckProof(root Hash, proof []byte, hi Hash, ht Hash, numLevels int) bool {
|
|
var empties [32]byte
|
|
copy(empties[:], proof[:len(empties)])
|
|
hashLen := len(EmptyNodeValue)
|
|
|
|
var siblings []Hash
|
|
for i := len(empties); i < len(proof); i += hashLen {
|
|
var siblingHash Hash
|
|
copy(siblingHash[:], proof[i:i+hashLen])
|
|
siblings = append(siblings, siblingHash)
|
|
}
|
|
|
|
path := getPath(numLevels, hi)
|
|
nodeHash := ht
|
|
siblingUsedPos := 0
|
|
|
|
for level := numLevels - 2; level >= 0; level-- {
|
|
var sibling Hash
|
|
if testbitmap(empties[:], uint(level)) {
|
|
sibling = siblings[siblingUsedPos]
|
|
siblingUsedPos++
|
|
} else {
|
|
sibling = EmptyNodeValue
|
|
}
|
|
// calculate the nodeHash with the current nodeHash and the sibling
|
|
var node treeNode
|
|
if path[numLevels-level-2] {
|
|
node = treeNode{
|
|
ChildL: sibling,
|
|
ChildR: nodeHash,
|
|
}
|
|
} else {
|
|
node = treeNode{
|
|
ChildL: nodeHash,
|
|
ChildR: sibling,
|
|
}
|
|
}
|
|
// if both childs are EmptyNodeValue, the parent will be EmptyNodeValue
|
|
if bytes.Equal(nodeHash[:], EmptyNodeValue[:]) && bytes.Equal(sibling[:], EmptyNodeValue[:]) {
|
|
nodeHash = EmptyNodeValue
|
|
} else {
|
|
nodeHash = node.Ht()
|
|
}
|
|
}
|
|
return bytes.Equal(nodeHash[:], root[:])
|
|
}
|