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