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