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