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