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