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