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