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