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