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 cp, err := mt.GenerateSCVerifierProof(k, rootKey)
926 if err != nil {
927 return nil, err
928 }
929 cp.Siblings = CircomSiblingsFromSiblings(cp.Siblings, mt.maxLevels)
930 return cp, nil
931}
932
933// GenerateSCVerifierProof returns the CircomVerifierProof for a certain key in
934// the MerkleTree with the Siblings without the extra 0 needed at the circom
935// circuits, which makes it straight forward to verifiy inside a Smart
936// Contract. If the rootKey is nil, the current merkletree root is used.
937func (mt *MerkleTree) GenerateSCVerifierProof(k *big.Int,
938 rootKey *Hash) (*CircomVerifierProof, error) {
939 if rootKey == nil {
940 rootKey = mt.Root()
941 }
942 p, v, err := mt.GenerateProof(k, rootKey)
943 if err != nil && err != ErrKeyNotFound {
944 return nil, err
945 }
946 var cp CircomVerifierProof
947 cp.Root = rootKey
948 cp.Siblings = p.AllSiblings()
949 cp.OldKey = &HashZero
950 cp.OldValue = &HashZero
951 cp.Key = NewHashFromBigInt(k)
952 cp.Value = NewHashFromBigInt(v)
953 if p.Existence {
954 cp.Fnc = 0 // inclusion
955 } else {
956 cp.Fnc = 1 // non inclusion
957 }
958
959 return &cp, nil
960}
961
962// GenerateProof generates the proof of existence (or non-existence) of an
963// Entry's hash Index for a Merkle Tree given the root.
964// If the rootKey is nil, the current merkletree root is used
965func (mt *MerkleTree) GenerateProof(k *big.Int, rootKey *Hash) (*Proof,
966 *big.Int, error) {
967 p := &Proof{}
968 var siblingKey *Hash
969
970 kHash := NewHashFromBigInt(k)
971 path := getPath(mt.maxLevels, kHash[:])
972 if rootKey == nil {
973 rootKey = mt.Root()
974 }
975 nextKey := rootKey
976 for p.depth = 0; p.depth < uint(mt.maxLevels); p.depth++ {
977 n, err := mt.GetNode(nextKey)
978 if err != nil {
979 return nil, nil, err
980 }
981 switch n.Type {
982 case NodeTypeEmpty:
983 return p, big.NewInt(0), nil
984 case NodeTypeLeaf:
985 if bytes.Equal(kHash[:], n.Entry[0][:]) {
986 p.Existence = true
987 return p, n.Entry[1].BigInt(), nil
988 }
989 // We found a leaf whose entry didn't match hIndex
990 p.NodeAux = &NodeAux{Key: n.Entry[0], Value: n.Entry[1]}
991 return p, n.Entry[1].BigInt(), nil
992 case NodeTypeMiddle:
993 if path[p.depth] {
994 nextKey = n.ChildR
995 siblingKey = n.ChildL
996 } else {
997 nextKey = n.ChildL
998 siblingKey = n.ChildR
999 }
1000 default:
1001 return nil, nil, ErrInvalidNodeFound
1002 }
1003 if !bytes.Equal(siblingKey[:], HashZero[:]) {
1004 SetBitBigEndian(p.notempties[:], uint(p.depth))
1005 p.Siblings = append(p.Siblings, siblingKey)
1006 }
1007 }
1008 return nil, nil, ErrKeyNotFound
1009}
1010
1011// VerifyProof verifies the Merkle Proof for the entry and root.
1012func VerifyProof(rootKey *Hash, proof *Proof, k, v *big.Int) bool {
1013 rootFromProof, err := RootFromProof(proof, k, v)
1014 if err != nil {
1015 return false
1016 }
1017 return bytes.Equal(rootKey[:], rootFromProof[:])
1018}
1019
1020// RootFromProof calculates the root that would correspond to a tree whose
1021// siblings are the ones in the proof with the leaf hashing to hIndex and
1022// hValue.
1023func RootFromProof(proof *Proof, k, v *big.Int) (*Hash, error) {
1024 kHash := NewHashFromBigInt(k)
1025 vHash := NewHashFromBigInt(v)
1026 sibIdx := len(proof.Siblings) - 1
1027 var err error
1028 var midKey *Hash
1029 if proof.Existence {
1030 midKey, err = LeafKey(kHash, vHash)
1031 if err != nil {
1032 return nil, err
1033 }
1034 } else {
1035 if proof.NodeAux == nil {
1036 midKey = &HashZero
1037 } else {
1038 if bytes.Equal(kHash[:], proof.NodeAux.Key[:]) {
1039 return nil,
1040 fmt.Errorf("Non-existence proof being checked against hIndex equal to nodeAux")
1041 }
1042 midKey, err = LeafKey(proof.NodeAux.Key, proof.NodeAux.Value)
1043 if err != nil {
1044 return nil, err
1045 }
1046 }
1047 }
1048 path := getPath(int(proof.depth), kHash[:])
1049 var siblingKey *Hash
1050 for lvl := int(proof.depth) - 1; lvl >= 0; lvl-- {
1051 if TestBitBigEndian(proof.notempties[:], uint(lvl)) {
1052 siblingKey = proof.Siblings[sibIdx]
1053 sibIdx--
1054 } else {
1055 siblingKey = &HashZero
1056 }
1057 if path[lvl] {
1058 midKey, err = NewNodeMiddle(siblingKey, midKey).Key()
1059 if err != nil {
1060 return nil, err
1061 }
1062 } else {
1063 midKey, err = NewNodeMiddle(midKey, siblingKey).Key()
1064 if err != nil {
1065 return nil, err
1066 }
1067 }
1068 }
1069 return midKey, nil
1070}
1071
1072// walk is a helper recursive function to iterate over all tree branches
1073func (mt *MerkleTree) walk(key *Hash, f func(*Node)) error {
1074 n, err := mt.GetNode(key)
1075 if err != nil {
1076 return err
1077 }
1078 switch n.Type {
1079 case NodeTypeEmpty:
1080 f(n)
1081 case NodeTypeLeaf:
1082 f(n)
1083 case NodeTypeMiddle:
1084 f(n)
1085 if err := mt.walk(n.ChildL, f); err != nil {
1086 return err
1087 }
1088 if err := mt.walk(n.ChildR, f); err != nil {
1089 return err
1090 }
1091 default:
1092 return ErrInvalidNodeFound
1093 }
1094 return nil
1095}
1096
1097// Walk iterates over all the branches of a MerkleTree with the given rootKey
1098// if rootKey is nil, it will get the current RootKey of the current state of
1099// the MerkleTree. For each node, it calls the f function given in the
1100// parameters. See some examples of the Walk function usage in the
1101// merkletree.go and merkletree_test.go
1102func (mt *MerkleTree) Walk(rootKey *Hash, f func(*Node)) error {
1103 if rootKey == nil {
1104 rootKey = mt.Root()
1105 }
1106 err := mt.walk(rootKey, f)
1107 return err
1108}
1109
1110// GraphViz uses Walk function to generate a string GraphViz representation of
1111// the tree and writes it to w
1112func (mt *MerkleTree) GraphViz(w io.Writer, rootKey *Hash) error {
1113 fmt.Fprintf(w, `digraph hierarchy {
1114node [fontname=Monospace,fontsize=10,shape=box]
1115`)
1116 cnt := 0
1117 var errIn error
1118 err := mt.Walk(rootKey, func(n *Node) {
1119 k, err := n.Key()
1120 if err != nil {
1121 errIn = err
1122 }
1123 switch n.Type {
1124 case NodeTypeEmpty:
1125 case NodeTypeLeaf:
1126 fmt.Fprintf(w, "\"%v\" [style=filled];\n", k.String())
1127 case NodeTypeMiddle:
1128 lr := [2]string{n.ChildL.String(), n.ChildR.String()}
1129 emptyNodes := ""
1130 for i := range lr {
1131 if lr[i] == "0" {
1132 lr[i] = fmt.Sprintf("empty%v", cnt)
1133 emptyNodes += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n", lr[i])
1134 cnt++
1135 }
1136 }
1137 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", k.String(), lr[0], lr[1])
1138 fmt.Fprint(w, emptyNodes)
1139 default:
1140 }
1141 })
1142 fmt.Fprintf(w, "}\n")
1143 if errIn != nil {
1144 return errIn
1145 }
1146 return err
1147}
1148
1149// PrintGraphViz prints directly the GraphViz() output
1150func (mt *MerkleTree) PrintGraphViz(rootKey *Hash) error {
1151 if rootKey == nil {
1152 rootKey = mt.Root()
1153 }
1154 w := bytes.NewBufferString("")
1155 fmt.Fprintf(w,
1156 "--------\nGraphViz of the MerkleTree with RootKey "+rootKey.BigInt().String()+"\n")
1157 err := mt.GraphViz(w, nil)
1158 if err != nil {
1159 return err
1160 }
1161 fmt.Fprintf(w,
1162 "End of GraphViz of the MerkleTree with RootKey "+rootKey.BigInt().String()+"\n--------\n")
1163
1164 fmt.Println(w)
1165 return nil
1166}
1167
1168// DumpLeafs returns all the Leafs that exist under the given Root. If no Root
1169// is given (nil), it uses the current Root of the MerkleTree.
1170func (mt *MerkleTree) DumpLeafs(rootKey *Hash) ([]byte, error) {
1171 var b []byte
1172 err := mt.Walk(rootKey, func(n *Node) {
1173 if n.Type == NodeTypeLeaf {
1174 l := n.Entry[0].Bytes()
1175 r := n.Entry[1].Bytes()
1176 b = append(b, append(l[:], r[:]...)...)
1177 }
1178 })
1179 return b, err
1180}
1181
1182// ImportDumpedLeafs parses and adds to the MerkleTree the dumped list of leafs
1183// from the DumpLeafs function.
1184func (mt *MerkleTree) ImportDumpedLeafs(b []byte) error {
1185 for i := 0; i < len(b); i += 64 {
1186 lr := b[i : i+64]
1187 lB, err := NewBigIntFromHashBytes(lr[:32])
1188 if err != nil {
1189 return err
1190 }
1191 rB, err := NewBigIntFromHashBytes(lr[32:])
1192 if err != nil {
1193 return err
1194 }
1195 err = mt.Add(lB, rB)
1196 if err != nil {
1197 return err
1198 }
1199 }
1200 return nil
1201}