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