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