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