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