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