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