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