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