go-merkletree-iden3/merkletree.go
merkletree.goRaw
1package merkletree
2
3import (
4 "bytes"
5 "encoding/hex"
6 "errors"
7 "fmt"
8 "io"
9 "math/big"
10 "sync"
11
12 "github.com/iden3/go-iden3-core/common"
13 cryptoUtils "github.com/iden3/go-iden3-crypto/utils"
14 "github.com/iden3/go-merkletree/db"
15)
16
17const (
18 // proofFlagsLen is the byte length of the flags in the proof header (first 32
19 // bytes).
20 proofFlagsLen = 2
21 // ElemBytesLen is the length of the Hash byte array
22 ElemBytesLen = 32
23)
24
25var (
26 // ErrNodeKeyAlreadyExists is used when a node key already exists.
27 ErrNodeKeyAlreadyExists = errors.New("key already exists")
28 // ErrKeyNotFound is used when a key is not found in the MerkleTree.
29 ErrKeyNotFound = errors.New("Key not found in the MerkleTree")
30 // ErrNodeBytesBadSize is used when the data of a node has an incorrect
31 // size and can't be parsed.
32 ErrNodeBytesBadSize = errors.New("node data has incorrect size in the DB")
33 // ErrReachedMaxLevel is used when a traversal of the MT reaches the
34 // maximum level.
35 ErrReachedMaxLevel = errors.New("reached maximum level of the merkle tree")
36 // ErrInvalidNodeFound is used when an invalid node is found and can't
37 // be parsed.
38 ErrInvalidNodeFound = errors.New("found an invalid node in the DB")
39 // ErrInvalidProofBytes is used when a serialized proof is invalid.
40 ErrInvalidProofBytes = errors.New("the serialized proof is invalid")
41 // ErrInvalidDBValue is used when a value in the key value DB is
42 // invalid (for example, it doen't contain a byte header and a []byte
43 // body of at least len=1.
44 ErrInvalidDBValue = errors.New("the value in the DB is invalid")
45 // ErrEntryIndexAlreadyExists is used when the entry index already
46 // exists in the tree.
47 ErrEntryIndexAlreadyExists = errors.New("the entry index already exists in the tree")
48 // ErrNotWritable is used when the MerkleTree is not writable and a write function is called
49 ErrNotWritable = errors.New("Merkle Tree not writable")
50 rootNodeValue = []byte("currentroot")
51 // HashZero is used at Empty nodes
52 HashZero = 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}
53)
54
55// Hash is the generic type stored in the MerkleTree
56type Hash [32]byte
57
58// String returns decimal representation in string format of the Hash
59func (h Hash) String() string {
60 s := h.BigInt().String()
61 if len(s) < 8 {
62 return s
63 }
64 return s[0:8] + "..."
65}
66
67// Hex returns the hexadecimal representation of the Hash
68func (h Hash) Hex() string {
69 return hex.EncodeToString(h.BigInt().Bytes())
70}
71
72// BigInt returns the *big.Int representation of the *Hash
73func (h *Hash) BigInt() *big.Int {
74 return new(big.Int).SetBytes(common.SwapEndianness(h[:]))
75}
76
77// Bytes returns the [32]byte representation of the *Hash
78func (h *Hash) Bytes() [32]byte {
79 bi := new(big.Int).SetBytes(common.SwapEndianness(h[:])).Bytes()
80 b := [32]byte{}
81 copy(b[:], bi[:])
82 return b
83}
84
85// NewBigIntFromBytes returns a *big.Int from a byte array, swapping the endianness in the process. This is the intended method to get a *big.Int from a byte array that previously has ben generated by the Hash.Bytes() method.
86func NewBigIntFromBytes(b []byte) (*big.Int, error) {
87 if len(b) != 32 {
88 return nil, fmt.Errorf("Expected 32 bytes, found %d bytes", len(b))
89 }
90 return new(big.Int).SetBytes(common.SwapEndianness(b[:32])), nil
91}
92
93// NewHashFromBigInt returns a *Hash representation of the given *big.Int
94func NewHashFromBigInt(b *big.Int) *Hash {
95 r := &Hash{}
96 copy(r[:], common.SwapEndianness(b.Bytes()))
97 return r
98}
99
100// MerkleTree is the struct with the main elements of the MerkleTree
101type MerkleTree struct {
102 sync.RWMutex
103 db db.Storage
104 rootKey *Hash
105 writable bool
106 maxLevels int
107}
108
109// NewMerkleTree loads a new Merkletree. If in the sotrage already exists one will open that one, if not, will create a new one.
110func NewMerkleTree(storage db.Storage, maxLevels int) (*MerkleTree, error) {
111 mt := MerkleTree{db: storage, maxLevels: maxLevels, writable: true}
112
113 v, err := mt.db.Get(rootNodeValue)
114 if err != nil {
115 tx, err := mt.db.NewTx()
116 if err != nil {
117 return nil, err
118 }
119 mt.rootKey = &HashZero
120 tx.Put(rootNodeValue, mt.rootKey[:])
121 err = tx.Commit()
122 if err != nil {
123 return nil, err
124 }
125 return &mt, nil
126 }
127 mt.rootKey = &Hash{}
128 copy(mt.rootKey[:], v)
129 return &mt, nil
130}
131
132// DB returns the MerkleTree.DB()
133func (mt *MerkleTree) DB() db.Storage {
134 return mt.db
135}
136
137// Root returns the MerkleRoot
138func (mt *MerkleTree) Root() *Hash {
139 return mt.rootKey
140}
141
142func (mt *MerkleTree) Snapshot(rootKey *Hash) (*MerkleTree, error) {
143 mt.RLock()
144 defer mt.RUnlock()
145 _, err := mt.GetNode(rootKey)
146 if err != nil {
147 return nil, err
148 }
149 return &MerkleTree{db: mt.db, maxLevels: mt.maxLevels, rootKey: rootKey, writable: false}, nil
150}
151
152// Add adds a Key & Value into the MerkleTree. Where the `k` determines the path from the Root to the Leaf.
153func (mt *MerkleTree) Add(k, v *big.Int) error {
154 // verify that the MerkleTree is writable
155 if !mt.writable {
156 return ErrNotWritable
157 }
158
159 // verfy that k & v are valid and fit inside the Finite Field.
160 if !cryptoUtils.CheckBigIntInField(k) {
161 return errors.New("Key not inside the Finite Field")
162 }
163 if !cryptoUtils.CheckBigIntInField(v) {
164 return errors.New("Value not inside the Finite Field")
165 }
166
167 tx, err := mt.db.NewTx()
168 if err != nil {
169 return err
170 }
171 mt.Lock()
172 defer mt.Unlock()
173
174 kHash := NewHashFromBigInt(k)
175 vHash := NewHashFromBigInt(v)
176 newNodeLeaf := NewNodeLeaf(kHash, vHash)
177 path := getPath(mt.maxLevels, kHash[:])
178
179 newRootKey, err := mt.addLeaf(tx, newNodeLeaf, mt.rootKey, 0, path)
180 if err != nil {
181 return err
182 }
183 mt.rootKey = newRootKey
184 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
185
186 if err := tx.Commit(); err != nil {
187 return err
188 }
189
190 return nil
191}
192
193// pushLeaf recursively pushes an existing oldLeaf down until its path diverges
194// from newLeaf, at which point both leafs are stored, all while updating the
195// path.
196func (mt *MerkleTree) pushLeaf(tx db.Tx, newLeaf *Node, oldLeaf *Node,
197 lvl int, pathNewLeaf []bool, pathOldLeaf []bool) (*Hash, error) {
198 if lvl > mt.maxLevels-2 {
199 return nil, ErrReachedMaxLevel
200 }
201 var newNodeMiddle *Node
202 if pathNewLeaf[lvl] == pathOldLeaf[lvl] { // We need to go deeper!
203 nextKey, err := mt.pushLeaf(tx, newLeaf, oldLeaf, lvl+1, pathNewLeaf, pathOldLeaf)
204 if err != nil {
205 return nil, err
206 }
207 if pathNewLeaf[lvl] {
208 newNodeMiddle = NewNodeMiddle(&HashZero, nextKey) // go right
209 } else {
210 newNodeMiddle = NewNodeMiddle(nextKey, &HashZero) // go left
211 }
212 return mt.addNode(tx, newNodeMiddle)
213 } else {
214 oldLeafKey, err := oldLeaf.Key()
215 if err != nil {
216 return nil, err
217 }
218 newLeafKey, err := newLeaf.Key()
219 if err != nil {
220 return nil, err
221 }
222
223 if pathNewLeaf[lvl] {
224 newNodeMiddle = NewNodeMiddle(oldLeafKey, newLeafKey)
225 } else {
226 newNodeMiddle = NewNodeMiddle(newLeafKey, oldLeafKey)
227 }
228 // We can add newLeaf now. We don't need to add oldLeaf because it's already in the tree.
229 _, err = mt.addNode(tx, newLeaf)
230 if err != nil {
231 return nil, err
232 }
233 return mt.addNode(tx, newNodeMiddle)
234 }
235}
236
237// addLeaf recursively adds a newLeaf in the MT while updating the path.
238func (mt *MerkleTree) addLeaf(tx db.Tx, newLeaf *Node, key *Hash,
239 lvl int, path []bool) (*Hash, error) {
240 var err error
241 var nextKey *Hash
242 if lvl > mt.maxLevels-1 {
243 return nil, ErrReachedMaxLevel
244 }
245 n, err := mt.GetNode(key)
246 if err != nil {
247 return nil, err
248 }
249 switch n.Type {
250 case NodeTypeEmpty:
251 // We can add newLeaf now
252 return mt.addNode(tx, newLeaf)
253 case NodeTypeLeaf:
254 nKey := n.Entry[0]
255 // Check if leaf node found contains the leaf node we are trying to add
256 newLeafKey := newLeaf.Entry[0]
257 if bytes.Equal(nKey[:], newLeafKey[:]) {
258 return nil, ErrEntryIndexAlreadyExists
259 }
260 pathOldLeaf := getPath(mt.maxLevels, nKey[:])
261 // We need to push newLeaf down until its path diverges from n's path
262 return mt.pushLeaf(tx, newLeaf, n, lvl, path, pathOldLeaf)
263 case NodeTypeMiddle:
264 // We need to go deeper, continue traversing the tree, left or right depending on path
265 var newNodeMiddle *Node
266 if path[lvl] {
267 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildR, lvl+1, path) // go right
268 newNodeMiddle = NewNodeMiddle(n.ChildL, nextKey)
269 } else {
270 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildL, lvl+1, path) // go left
271 newNodeMiddle = NewNodeMiddle(nextKey, n.ChildR)
272 }
273 if err != nil {
274 return nil, err
275 }
276 // Update the node to reflect the modified child
277 return mt.addNode(tx, newNodeMiddle)
278 default:
279 return nil, ErrInvalidNodeFound
280 }
281}
282
283// addNode adds a node into the MT. Empty nodes are not stored in the tree;
284// they are all the same and assumed to always exist.
285func (mt *MerkleTree) addNode(tx db.Tx, n *Node) (*Hash, error) {
286 // verify that the MerkleTree is writable
287 if !mt.writable {
288 return nil, ErrNotWritable
289 }
290 if n.Type == NodeTypeEmpty {
291 return n.Key()
292 }
293 k, err := n.Key()
294 if err != nil {
295 return nil, err
296 }
297 v := n.Value()
298 // Check that the node key doesn't already exist
299 if _, err := tx.Get(k[:]); err == nil {
300 return nil, ErrNodeKeyAlreadyExists
301 }
302 tx.Put(k[:], v)
303 return k, nil
304}
305
306// Get returns the value of the leaf for the given key
307func (mt *MerkleTree) Get(k *big.Int) (*big.Int, error) {
308 // verfy that k is valid and fit inside the Finite Field.
309 if !cryptoUtils.CheckBigIntInField(k) {
310 return nil, errors.New("Key not inside the Finite Field")
311 }
312
313 kHash := NewHashFromBigInt(k)
314 path := getPath(mt.maxLevels, kHash[:])
315
316 nextKey := mt.rootKey
317 for i := 0; i < mt.maxLevels; i++ {
318 n, err := mt.GetNode(nextKey)
319 if err != nil {
320 return nil, err
321 }
322 switch n.Type {
323 case NodeTypeEmpty:
324 return nil, ErrKeyNotFound
325 case NodeTypeLeaf:
326 if bytes.Equal(kHash[:], n.Entry[0][:]) {
327 return n.Entry[1].BigInt(), nil
328 } else {
329 return nil, ErrKeyNotFound
330 }
331 case NodeTypeMiddle:
332 if path[i] {
333 nextKey = n.ChildR
334 } else {
335 nextKey = n.ChildL
336 }
337 default:
338 return nil, ErrInvalidNodeFound
339 }
340 }
341
342 return nil, ErrKeyNotFound
343}
344
345// Delete removes the specified Key from the MerkleTree and updates the path
346// from the deleted key to the Root with the new values. This method removes
347// the key from the MerkleTree, but does not remove the old nodes from the
348// key-value database; this means that if the tree is accessed by an old Root
349// where the key was not deleted yet, the key will still exist. If is desired
350// to remove the key-values from the database that are not under the current
351// Root, an option could be to dump all the leafs (using mt.DumpLeafs) and
352// import them in a new MerkleTree in a new database (using
353// mt.ImportDumpedLeafs), but this will loose all the Root history of the
354// MerkleTree
355func (mt *MerkleTree) Delete(k *big.Int) error {
356 // verify that the MerkleTree is writable
357 if !mt.writable {
358 return ErrNotWritable
359 }
360
361 // verfy that k is valid and fit inside the Finite Field.
362 if !cryptoUtils.CheckBigIntInField(k) {
363 return errors.New("Key not inside the Finite Field")
364 }
365 tx, err := mt.db.NewTx()
366 if err != nil {
367 return err
368 }
369 mt.Lock()
370 defer mt.Unlock()
371
372 kHash := NewHashFromBigInt(k)
373 path := getPath(mt.maxLevels, kHash[:])
374
375 nextKey := mt.rootKey
376 var siblings []*Hash
377 for i := 0; i < mt.maxLevels; i++ {
378 n, err := mt.GetNode(nextKey)
379 if err != nil {
380 return err
381 }
382 switch n.Type {
383 case NodeTypeEmpty:
384 return ErrKeyNotFound
385 case NodeTypeLeaf:
386 if bytes.Equal(kHash[:], n.Entry[0][:]) {
387 // remove and go up with the sibling
388 err = mt.rmAndUpload(tx, path, kHash, siblings)
389 return err
390 } else {
391 return ErrKeyNotFound
392 }
393 case NodeTypeMiddle:
394 if path[i] {
395 nextKey = n.ChildR
396 siblings = append(siblings, n.ChildL)
397 } else {
398 nextKey = n.ChildL
399 siblings = append(siblings, n.ChildR)
400 }
401 default:
402 return ErrInvalidNodeFound
403 }
404 }
405
406 return nil
407}
408
409// rmAndUpload removes the key, and goes up until the root updating all the nodes with the new values.
410func (mt *MerkleTree) rmAndUpload(tx db.Tx, path []bool, kHash *Hash, siblings []*Hash) error {
411 if len(siblings) == 0 {
412 mt.rootKey = &HashZero
413 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
414 return tx.Commit()
415 }
416
417 toUpload := siblings[len(siblings)-1]
418 if len(siblings) < 2 {
419 mt.rootKey = siblings[0]
420 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
421 return tx.Commit()
422 }
423 for i := len(siblings) - 2; i >= 0; i-- {
424 if !bytes.Equal(siblings[i][:], HashZero[:]) {
425 var newNode *Node
426 if path[i] {
427 newNode = NewNodeMiddle(siblings[i], toUpload)
428 } else {
429 newNode = NewNodeMiddle(toUpload, siblings[i])
430 }
431 _, err := mt.addNode(tx, newNode)
432 if err != ErrNodeKeyAlreadyExists && err != nil {
433 return err
434 }
435 // go up until the root
436 newRootKey, err := mt.recalculatePathUntilRoot(tx, path, newNode, siblings[:i])
437 if err != nil {
438 return err
439 }
440 mt.rootKey = newRootKey
441 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
442 break
443 }
444 // if i==0 (root position), stop and store the sibling of the deleted leaf as root
445 if i == 0 {
446 mt.rootKey = toUpload
447 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
448 break
449 }
450 }
451 if err := tx.Commit(); err != nil {
452 return err
453 }
454
455 return nil
456}
457
458// recalculatePathUntilRoot recalculates the nodes until the Root
459func (mt *MerkleTree) recalculatePathUntilRoot(tx db.Tx, path []bool, node *Node, siblings []*Hash) (*Hash, error) {
460 for i := len(siblings) - 1; i >= 0; i-- {
461 nodeKey, err := node.Key()
462 if err != nil {
463 return nil, err
464 }
465 if path[i] {
466 node = NewNodeMiddle(siblings[i], nodeKey)
467 } else {
468 node = NewNodeMiddle(nodeKey, siblings[i])
469 }
470 _, err = mt.addNode(tx, node)
471 if err != ErrNodeKeyAlreadyExists && err != nil {
472 return nil, err
473 }
474 }
475
476 // return last node added, which is the root
477 nodeKey, err := node.Key()
478 return nodeKey, err
479}
480
481// dbGet is a helper function to get the node of a key from the internal
482// storage.
483func (mt *MerkleTree) dbGet(k []byte) (NodeType, []byte, error) {
484 if bytes.Equal(k, HashZero[:]) {
485 return 0, nil, nil
486 }
487
488 value, err := mt.db.Get(k)
489 if err != nil {
490 return 0, nil, err
491 }
492
493 if len(value) < 2 {
494 return 0, nil, ErrInvalidDBValue
495 }
496 nodeType := value[0]
497 nodeBytes := value[1:]
498
499 return NodeType(nodeType), nodeBytes, nil
500}
501
502// dbInsert is a helper function to insert a node into a key in an open db
503// transaction.
504func (mt *MerkleTree) dbInsert(tx db.Tx, k []byte, t NodeType, data []byte) {
505 v := append([]byte{byte(t)}, data...)
506 tx.Put(k, v)
507}
508
509// GetNode gets a node by key from the MT. Empty nodes are not stored in the
510// tree; they are all the same and assumed to always exist.
511func (mt *MerkleTree) GetNode(key *Hash) (*Node, error) {
512 if bytes.Equal(key[:], HashZero[:]) {
513 return NewNodeEmpty(), nil
514 }
515 nBytes, err := mt.db.Get(key[:])
516 if err != nil {
517 return nil, err
518 }
519 return NewNodeFromBytes(nBytes)
520}
521
522// getPath returns the binary path, from the root to the leaf.
523func getPath(numLevels int, k []byte) []bool {
524 path := make([]bool, numLevels)
525 for n := 0; n < numLevels; n++ {
526 path[n] = common.TestBit(k[:], uint(n))
527 }
528 return path
529}
530
531// NodeAux contains the auxiliary node used in a non-existence proof.
532type NodeAux struct {
533 Key *Hash
534 Value *Hash
535}
536
537// Proof defines the required elements for a MT proof of existence or non-existence.
538type Proof struct {
539 // existence indicates wether this is a proof of existence or non-existence.
540 Existence bool
541 // depth indicates how deep in the tree the proof goes.
542 depth uint
543 // notempties is a bitmap of non-empty Siblings found in Siblings.
544 notempties [ElemBytesLen - proofFlagsLen]byte
545 // Siblings is a list of non-empty sibling keys.
546 Siblings []*Hash
547 NodeAux *NodeAux
548}
549
550// NewProofFromBytes parses a byte array into a Proof.
551func NewProofFromBytes(bs []byte) (*Proof, error) {
552 if len(bs) < ElemBytesLen {
553 return nil, ErrInvalidProofBytes
554 }
555 p := &Proof{}
556 if (bs[0] & 0x01) == 0 {
557 p.Existence = true
558 }
559 p.depth = uint(bs[1])
560 copy(p.notempties[:], bs[proofFlagsLen:ElemBytesLen])
561 siblingBytes := bs[ElemBytesLen:]
562 sibIdx := 0
563 for i := uint(0); i < p.depth; i++ {
564 if common.TestBitBigEndian(p.notempties[:], i) {
565 if len(siblingBytes) < (sibIdx+1)*ElemBytesLen {
566 return nil, ErrInvalidProofBytes
567 }
568 var sib Hash
569 copy(sib[:], siblingBytes[sibIdx*ElemBytesLen:(sibIdx+1)*ElemBytesLen])
570 p.Siblings = append(p.Siblings, &sib)
571 sibIdx++
572 }
573 }
574
575 if !p.Existence && ((bs[0] & 0x02) != 0) {
576 p.NodeAux = &NodeAux{Key: &Hash{}, Value: &Hash{}}
577 nodeAuxBytes := siblingBytes[len(p.Siblings)*ElemBytesLen:]
578 if len(nodeAuxBytes) != 2*ElemBytesLen {
579 return nil, ErrInvalidProofBytes
580 }
581 copy(p.NodeAux.Key[:], nodeAuxBytes[:ElemBytesLen])
582 copy(p.NodeAux.Value[:], nodeAuxBytes[ElemBytesLen:2*ElemBytesLen])
583 }
584 return p, nil
585}
586
587// Bytes serializes a Proof into a byte array.
588func (p *Proof) Bytes() []byte {
589 bsLen := proofFlagsLen + len(p.notempties) + ElemBytesLen*len(p.Siblings)
590 if p.NodeAux != nil {
591 bsLen += 2 * ElemBytesLen
592 }
593 bs := make([]byte, bsLen)
594
595 if !p.Existence {
596 bs[0] |= 0x01
597 }
598 bs[1] = byte(p.depth)
599 copy(bs[proofFlagsLen:len(p.notempties)+proofFlagsLen], p.notempties[:])
600 siblingsBytes := bs[len(p.notempties)+proofFlagsLen:]
601 for i, k := range p.Siblings {
602 copy(siblingsBytes[i*ElemBytesLen:(i+1)*ElemBytesLen], k[:])
603 }
604 if p.NodeAux != nil {
605 bs[0] |= 0x02
606 copy(bs[len(bs)-2*ElemBytesLen:], p.NodeAux.Key[:])
607 copy(bs[len(bs)-1*ElemBytesLen:], p.NodeAux.Value[:])
608 }
609 return bs
610}
611
612// SiblingsFromProof returns all the siblings of the proof. This function is used to generate the siblings input for the circom circuits.
613func SiblingsFromProof(proof *Proof) []*Hash {
614 sibIdx := 0
615 var siblings []*Hash
616 for lvl := 0; lvl < int(proof.depth); lvl++ {
617 if common.TestBitBigEndian(proof.notempties[:], uint(lvl)) {
618 siblings = append(siblings, proof.Siblings[sibIdx])
619 sibIdx++
620 } else {
621 siblings = append(siblings, &HashZero)
622 }
623 }
624 return siblings
625}
626
627func (p *Proof) AllSiblings() []*Hash {
628 return SiblingsFromProof(p)
629}
630
631func (p *Proof) AllSiblingsCircom(levels int) []*big.Int {
632 siblings := p.AllSiblings()
633 // Add the rest of empty levels to the siblings
634 for i := len(siblings); i < levels; i++ {
635 siblings = append(siblings, &HashZero)
636 }
637 siblings = append(siblings, &HashZero) // add extra level for circom compatibility
638 siblingsBigInt := make([]*big.Int, len(siblings))
639 for i, sibling := range siblings {
640 siblingsBigInt[i] = sibling.BigInt()
641 }
642 return siblingsBigInt
643}
644
645// GenerateProof generates the proof of existence (or non-existence) of an
646// Entry's hash Index for a Merkle Tree given the root.
647// If the rootKey is nil, the current merkletree root is used
648func (mt *MerkleTree) GenerateProof(k *big.Int, rootKey *Hash) (*Proof, error) {
649 p := &Proof{}
650 var siblingKey *Hash
651
652 kHash := NewHashFromBigInt(k)
653 path := getPath(mt.maxLevels, kHash[:])
654 if rootKey == nil {
655 rootKey = mt.Root()
656 }
657 nextKey := rootKey
658 for p.depth = 0; p.depth < uint(mt.maxLevels); p.depth++ {
659 n, err := mt.GetNode(nextKey)
660 if err != nil {
661 return nil, err
662 }
663 switch n.Type {
664 case NodeTypeEmpty:
665 return p, nil
666 case NodeTypeLeaf:
667 if bytes.Equal(kHash[:], n.Entry[0][:]) {
668 p.Existence = true
669 return p, nil
670 } else {
671 // We found a leaf whose entry didn't match hIndex
672 p.NodeAux = &NodeAux{Key: n.Entry[0], Value: n.Entry[1]}
673 return p, nil
674 }
675 case NodeTypeMiddle:
676 if path[p.depth] {
677 nextKey = n.ChildR
678 siblingKey = n.ChildL
679 } else {
680 nextKey = n.ChildL
681 siblingKey = n.ChildR
682 }
683 default:
684 return nil, ErrInvalidNodeFound
685 }
686 if !bytes.Equal(siblingKey[:], HashZero[:]) {
687 common.SetBitBigEndian(p.notempties[:], uint(p.depth))
688 p.Siblings = append(p.Siblings, siblingKey)
689 }
690 }
691 return nil, ErrKeyNotFound
692}
693
694// VerifyProof verifies the Merkle Proof for the entry and root.
695func VerifyProof(rootKey *Hash, proof *Proof, k, v *big.Int) bool {
696 rootFromProof, err := RootFromProof(proof, k, v)
697 if err != nil {
698 return false
699 }
700 return bytes.Equal(rootKey[:], rootFromProof[:])
701}
702
703// RootFromProof calculates the root that would correspond to a tree whose
704// siblings are the ones in the proof with the leaf hashing to hIndex and
705// hValue.
706func RootFromProof(proof *Proof, k, v *big.Int) (*Hash, error) {
707 kHash := NewHashFromBigInt(k)
708 vHash := NewHashFromBigInt(v)
709 sibIdx := len(proof.Siblings) - 1
710 var err error
711 var midKey *Hash
712 if proof.Existence {
713 midKey, err = LeafKey(kHash, vHash)
714 if err != nil {
715 return nil, err
716 }
717 } else {
718 if proof.NodeAux == nil {
719 midKey = &HashZero
720 } else {
721 if bytes.Equal(kHash[:], proof.NodeAux.Key[:]) {
722 return nil, fmt.Errorf("Non-existence proof being checked against hIndex equal to nodeAux")
723 }
724 midKey, err = LeafKey(proof.NodeAux.Key, proof.NodeAux.Value)
725 if err != nil {
726 return nil, err
727 }
728 }
729 }
730 path := getPath(int(proof.depth), kHash[:])
731 var siblingKey *Hash
732 for lvl := int(proof.depth) - 1; lvl >= 0; lvl-- {
733 if common.TestBitBigEndian(proof.notempties[:], uint(lvl)) {
734 siblingKey = proof.Siblings[sibIdx]
735 sibIdx--
736 } else {
737 siblingKey = &HashZero
738 }
739 if path[lvl] {
740 midKey, err = NewNodeMiddle(siblingKey, midKey).Key()
741 if err != nil {
742 return nil, err
743 }
744 } else {
745 midKey, err = NewNodeMiddle(midKey, siblingKey).Key()
746 if err != nil {
747 return nil, err
748 }
749 }
750 }
751 return midKey, nil
752}
753
754// walk is a helper recursive function to iterate over all tree branches
755func (mt *MerkleTree) walk(key *Hash, f func(*Node)) error {
756 n, err := mt.GetNode(key)
757 if err != nil {
758 return err
759 }
760 switch n.Type {
761 case NodeTypeEmpty:
762 f(n)
763 case NodeTypeLeaf:
764 f(n)
765 case NodeTypeMiddle:
766 f(n)
767 if err := mt.walk(n.ChildL, f); err != nil {
768 return err
769 }
770 if err := mt.walk(n.ChildR, f); err != nil {
771 return err
772 }
773 default:
774 return ErrInvalidNodeFound
775 }
776 return nil
777}
778
779// Walk iterates over all the branches of a MerkleTree with the given rootKey
780// if rootKey is nil, it will get the current RootKey of the current state of the MerkleTree.
781// For each node, it calls the f function given in the parameters.
782// See some examples of the Walk function usage in the merkletree.go and
783// merkletree_test.go
784func (mt *MerkleTree) Walk(rootKey *Hash, f func(*Node)) error {
785 if rootKey == nil {
786 rootKey = mt.Root()
787 }
788 err := mt.walk(rootKey, f)
789 return err
790}
791
792// GraphViz uses Walk function to generate a string GraphViz representation of the
793// tree and writes it to w
794func (mt *MerkleTree) GraphViz(w io.Writer, rootKey *Hash) error {
795 fmt.Fprintf(w, `digraph hierarchy {
796node [fontname=Monospace,fontsize=10,shape=box]
797`)
798 cnt := 0
799 var errIn error
800 err := mt.Walk(rootKey, func(n *Node) {
801 k, err := n.Key()
802 if err != nil {
803 errIn = err
804 }
805 switch n.Type {
806 case NodeTypeEmpty:
807 case NodeTypeLeaf:
808 fmt.Fprintf(w, "\"%v\" [style=filled];\n", k.String())
809 case NodeTypeMiddle:
810 lr := [2]string{n.ChildL.String(), n.ChildR.String()}
811 emptyNodes := ""
812 for i := range lr {
813 if lr[i] == "0" {
814 lr[i] = fmt.Sprintf("empty%v", cnt)
815 emptyNodes += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n", lr[i])
816 cnt++
817 }
818 }
819 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", k.String(), lr[0], lr[1])
820 fmt.Fprint(w, emptyNodes)
821 default:
822 }
823 })
824 fmt.Fprintf(w, "}\n")
825 if errIn != nil {
826 return errIn
827 }
828 return err
829}
830
831// PrintGraphViz prints directly the GraphViz() output
832func (mt *MerkleTree) PrintGraphViz(rootKey *Hash) error {
833 if rootKey == nil {
834 rootKey = mt.Root()
835 }
836 w := bytes.NewBufferString("")
837 fmt.Fprintf(w, "--------\nGraphViz of the MerkleTree with RootKey "+rootKey.BigInt().String()+"\n")
838 err := mt.GraphViz(w, nil)
839 if err != nil {
840 return err
841 }
842 fmt.Fprintf(w, "End of GraphViz of the MerkleTree with RootKey "+rootKey.BigInt().String()+"\n--------\n")
843
844 fmt.Println(w)
845 return nil
846}
847
848// DumpLeafs returns all the Leafs that exist under the given Root. If no Root
849// is given (nil), it uses the current Root of the MerkleTree.
850func (mt *MerkleTree) DumpLeafs(rootKey *Hash) ([]byte, error) {
851 var b []byte
852 err := mt.Walk(rootKey, func(n *Node) {
853 if n.Type == NodeTypeLeaf {
854 l := n.Entry[0].Bytes()
855 r := n.Entry[1].Bytes()
856 b = append(b, append(l[:], r[:]...)...)
857 }
858 })
859 return b, err
860}
861
862// ImportDumpedLeafs parses and adds to the MerkleTree the dumped list of leafs
863// from the DumpLeafs function.
864func (mt *MerkleTree) ImportDumpedLeafs(b []byte) error {
865 for i := 0; i < len(b); i += 64 {
866 lr := b[i : i+64]
867 lB, err := NewBigIntFromBytes(lr[:32])
868 if err != nil {
869 return err
870 }
871 rB, err := NewBigIntFromBytes(lr[32:])
872 if err != nil {
873 return err
874 }
875 err = mt.Add(lB, rB)
876 if err != nil {
877 return err
878 }
879 }
880 return nil
881}