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