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