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 "github.com/iden3/go-iden3-core/db"
14 cryptoUtils "github.com/iden3/go-iden3-crypto/utils"
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("node already exists")
28 // ErrEntryIndexNotFound is used when no entry is found for an index.
29 ErrEntryIndexNotFound = errors.New("node index not found in the DB")
30 // ErrNodeDataBadSize is used when the data of a node has an incorrect
31 // size and can't be parsed.
32 ErrNodeDataBadSize = 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
58func (h Hash) String() string {
59 return new(big.Int).SetBytes(h[:]).String()
60}
61func (h Hash) Hex() string {
62 return hex.EncodeToString(h[:])
63}
64
65// BigInt returns the *big.Int representation of the *Hash
66func (h *Hash) BigInt() *big.Int {
67 return new(big.Int).SetBytes(common.SwapEndianness(h[:]))
68}
69
70// NewHashFromBigInt returns a *Hash representation of the given *big.Int
71func NewHashFromBigInt(b *big.Int) *Hash {
72 r := &Hash{}
73 copy(r[:], common.SwapEndianness(b.Bytes()))
74 return r
75}
76
77// MerkleTree is the struct with the main elements of the MerkleTree
78type MerkleTree struct {
79 sync.RWMutex
80 db db.Storage
81 rootKey *Hash
82 writable bool
83 maxLevels int
84}
85
86// NewMerkleTree loads a new Merkletree. If in the sotrage already exists one will open that one, if not, will create a new one.
87func NewMerkleTree(storage db.Storage, maxLevels int) (*MerkleTree, error) {
88 mt := MerkleTree{db: storage, maxLevels: maxLevels, writable: true}
89
90 v, err := mt.db.Get(rootNodeValue)
91 if err != nil {
92 tx, err := mt.db.NewTx()
93 if err != nil {
94 return nil, err
95 }
96 mt.rootKey = &HashZero
97 tx.Put(rootNodeValue, mt.rootKey[:])
98 err = tx.Commit()
99 if err != nil {
100 return nil, err
101 }
102 return &mt, nil
103 }
104 mt.rootKey = &Hash{}
105 copy(mt.rootKey[:], v)
106 return &mt, nil
107}
108
109// DB returns the MerkleTree.DB()
110func (mt *MerkleTree) DB() db.Storage {
111 return mt.db
112}
113
114// Root returns the MerkleRoot
115func (mt *MerkleTree) Root() *Hash {
116 return mt.rootKey
117}
118
119// Add adds a Key & Value into the MerkleTree. Where the `k` determines the path from the Root to the Leaf.
120func (mt *MerkleTree) Add(k, v *big.Int) error {
121 // verify that the MerkleTree is writable
122 if !mt.writable {
123 return ErrNotWritable
124 }
125
126 // verfy that the ElemBytes are valid and fit inside the Finite Field.
127 if !cryptoUtils.CheckBigIntInField(k) {
128 return errors.New("Key not inside the Finite Field")
129 }
130 if !cryptoUtils.CheckBigIntInField(v) {
131 return errors.New("Value not inside the Finite Field")
132 }
133
134 tx, err := mt.db.NewTx()
135 if err != nil {
136 return err
137 }
138 mt.Lock()
139 defer mt.Unlock()
140
141 kHash := NewHashFromBigInt(k)
142 vHash := NewHashFromBigInt(v)
143 newNodeLeaf := NewNodeLeaf(kHash, vHash)
144 path := getPath(mt.maxLevels, kHash[:])
145
146 newRootKey, err := mt.addLeaf(tx, newNodeLeaf, mt.rootKey, 0, path)
147 if err != nil {
148 return err
149 }
150 mt.rootKey = newRootKey
151 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
152
153 if err := tx.Commit(); err != nil {
154 return err
155 }
156
157 return nil
158}
159
160// pushLeaf recursively pushes an existing oldLeaf down until its path diverges
161// from newLeaf, at which point both leafs are stored, all while updating the
162// path.
163func (mt *MerkleTree) pushLeaf(tx db.Tx, newLeaf *Node, oldLeaf *Node,
164 lvl int, pathNewLeaf []bool, pathOldLeaf []bool) (*Hash, error) {
165 if lvl > mt.maxLevels-2 {
166 return nil, ErrReachedMaxLevel
167 }
168 var newNodeMiddle *Node
169 if pathNewLeaf[lvl] == pathOldLeaf[lvl] { // We need to go deeper!
170 nextKey, err := mt.pushLeaf(tx, newLeaf, oldLeaf, lvl+1, pathNewLeaf, pathOldLeaf)
171 if err != nil {
172 return nil, err
173 }
174 if pathNewLeaf[lvl] {
175 newNodeMiddle = NewNodeMiddle(&HashZero, nextKey) // go right
176 } else {
177 newNodeMiddle = NewNodeMiddle(nextKey, &HashZero) // go left
178 }
179 return mt.addNode(tx, newNodeMiddle)
180 } else {
181 oldLeafKey, err := oldLeaf.Key()
182 if err != nil {
183 return nil, err
184 }
185 newLeafKey, err := newLeaf.Key()
186 if err != nil {
187 return nil, err
188 }
189
190 if pathNewLeaf[lvl] {
191 newNodeMiddle = NewNodeMiddle(oldLeafKey, newLeafKey)
192 } else {
193 newNodeMiddle = NewNodeMiddle(newLeafKey, oldLeafKey)
194 }
195 // We can add newLeaf now. We don't need to add oldLeaf because it's already in the tree.
196 _, err = mt.addNode(tx, newLeaf)
197 if err != nil {
198 return nil, err
199 }
200 return mt.addNode(tx, newNodeMiddle)
201 }
202}
203
204// addLeaf recursively adds a newLeaf in the MT while updating the path.
205func (mt *MerkleTree) addLeaf(tx db.Tx, newLeaf *Node, key *Hash,
206 lvl int, path []bool) (*Hash, error) {
207 var err error
208 var nextKey *Hash
209 if lvl > mt.maxLevels-1 {
210 return nil, ErrReachedMaxLevel
211 }
212 n, err := mt.GetNode(key)
213 if err != nil {
214 return nil, err
215 }
216 switch n.Type {
217 case NodeTypeEmpty:
218 // We can add newLeaf now
219 return mt.addNode(tx, newLeaf)
220 case NodeTypeLeaf:
221 nKey := n.Entry[0]
222 // Check if leaf node found contains the leaf node we are trying to add
223 newLeafKey := newLeaf.Entry[0]
224 if bytes.Equal(nKey[:], newLeafKey[:]) {
225 return nil, ErrEntryIndexAlreadyExists
226 }
227 pathOldLeaf := getPath(mt.maxLevels, nKey[:])
228 // We need to push newLeaf down until its path diverges from n's path
229 return mt.pushLeaf(tx, newLeaf, n, lvl, path, pathOldLeaf)
230 case NodeTypeMiddle:
231 // We need to go deeper, continue traversing the tree, left or right depending on path
232 var newNodeMiddle *Node
233 if path[lvl] {
234 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildR, lvl+1, path) // go right
235 newNodeMiddle = NewNodeMiddle(n.ChildL, nextKey)
236 } else {
237 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildL, lvl+1, path) // go left
238 newNodeMiddle = NewNodeMiddle(nextKey, n.ChildR)
239 }
240 if err != nil {
241 return nil, err
242 }
243 // Update the node to reflect the modified child
244 return mt.addNode(tx, newNodeMiddle)
245 default:
246 return nil, ErrInvalidNodeFound
247 }
248}
249
250// addNode adds a node into the MT. Empty nodes are not stored in the tree;
251// they are all the same and assumed to always exist.
252func (mt *MerkleTree) addNode(tx db.Tx, n *Node) (*Hash, error) {
253 // verify that the MerkleTree is writable
254 if !mt.writable {
255 return nil, ErrNotWritable
256 }
257 if n.Type == NodeTypeEmpty {
258 return n.Key()
259 }
260 k, err := n.Key()
261 if err != nil {
262 return nil, err
263 }
264 v := n.Value()
265 // Check that the node key doesn't already exist
266 if _, err := tx.Get(k[:]); err == nil {
267 return nil, ErrNodeKeyAlreadyExists
268 }
269 tx.Put(k[:], v)
270 return k, nil
271}
272
273// dbGet is a helper function to get the node of a key from the internal
274// storage.
275func (mt *MerkleTree) dbGet(k []byte) (NodeType, []byte, error) {
276 if bytes.Equal(k, HashZero[:]) {
277 return 0, nil, nil
278 }
279
280 value, err := mt.db.Get(k)
281 if err != nil {
282 return 0, nil, err
283 }
284
285 if len(value) < 2 {
286 return 0, nil, ErrInvalidDBValue
287 }
288 nodeType := value[0]
289 nodeBytes := value[1:]
290
291 return NodeType(nodeType), nodeBytes, nil
292}
293
294// dbInsert is a helper function to insert a node into a key in an open db
295// transaction.
296func (mt *MerkleTree) dbInsert(tx db.Tx, k []byte, t NodeType, data []byte) {
297 v := append([]byte{byte(t)}, data...)
298 tx.Put(k, v)
299}
300
301// GetNode gets a node by key from the MT. Empty nodes are not stored in the
302// tree; they are all the same and assumed to always exist.
303func (mt *MerkleTree) GetNode(key *Hash) (*Node, error) {
304 if bytes.Equal(key[:], HashZero[:]) {
305 return NewNodeEmpty(), nil
306 }
307 nBytes, err := mt.db.Get(key[:])
308 if err != nil {
309 return nil, err
310 }
311 return NewNodeFromBytes(nBytes)
312}
313
314// getPath returns the binary path, from the root to the leaf.
315func getPath(numLevels int, k []byte) []bool {
316 path := make([]bool, numLevels)
317 for n := 0; n < numLevels; n++ {
318 path[n] = common.TestBit(k[:], uint(n))
319 }
320 return path
321}
322
323// NodeAux contains the auxiliary node used in a non-existence proof.
324type NodeAux struct {
325 Key *Hash
326 Value *Hash
327}
328
329// Proof defines the required elements for a MT proof of existence or non-existence.
330type Proof struct {
331 // existence indicates wether this is a proof of existence or non-existence.
332 Existence bool
333 // depth indicates how deep in the tree the proof goes.
334 depth uint
335 // notempties is a bitmap of non-empty Siblings found in Siblings.
336 notempties [ElemBytesLen - proofFlagsLen]byte
337 // Siblings is a list of non-empty sibling keys.
338 Siblings []*Hash
339 NodeAux *NodeAux
340}
341
342// NewProofFromBytes parses a byte array into a Proof.
343func NewProofFromBytes(bs []byte) (*Proof, error) {
344 if len(bs) < ElemBytesLen {
345 return nil, ErrInvalidProofBytes
346 }
347 p := &Proof{}
348 if (bs[0] & 0x01) == 0 {
349 p.Existence = true
350 }
351 p.depth = uint(bs[1])
352 copy(p.notempties[:], bs[proofFlagsLen:ElemBytesLen])
353 siblingBytes := bs[ElemBytesLen:]
354 sibIdx := 0
355 for i := uint(0); i < p.depth; i++ {
356 if common.TestBitBigEndian(p.notempties[:], i) {
357 if len(siblingBytes) < (sibIdx+1)*ElemBytesLen {
358 return nil, ErrInvalidProofBytes
359 }
360 var sib Hash
361 copy(sib[:], siblingBytes[sibIdx*ElemBytesLen:(sibIdx+1)*ElemBytesLen])
362 p.Siblings = append(p.Siblings, &sib)
363 sibIdx++
364 }
365 }
366
367 if !p.Existence && ((bs[0] & 0x02) != 0) {
368 p.NodeAux = &NodeAux{Key: &Hash{}, Value: &Hash{}}
369 nodeAuxBytes := siblingBytes[len(p.Siblings)*ElemBytesLen:]
370 if len(nodeAuxBytes) != 2*ElemBytesLen {
371 return nil, ErrInvalidProofBytes
372 }
373 copy(p.NodeAux.Key[:], nodeAuxBytes[:ElemBytesLen])
374 copy(p.NodeAux.Value[:], nodeAuxBytes[ElemBytesLen:2*ElemBytesLen])
375 }
376 return p, nil
377}
378
379// Bytes serializes a Proof into a byte array.
380func (p *Proof) Bytes() []byte {
381 bsLen := proofFlagsLen + len(p.notempties) + ElemBytesLen*len(p.Siblings)
382 if p.NodeAux != nil {
383 bsLen += 2 * ElemBytesLen
384 }
385 bs := make([]byte, bsLen)
386
387 if !p.Existence {
388 bs[0] |= 0x01
389 }
390 bs[1] = byte(p.depth)
391 copy(bs[proofFlagsLen:len(p.notempties)+proofFlagsLen], p.notempties[:])
392 siblingsBytes := bs[len(p.notempties)+proofFlagsLen:]
393 for i, k := range p.Siblings {
394 copy(siblingsBytes[i*ElemBytesLen:(i+1)*ElemBytesLen], k[:])
395 }
396 if p.NodeAux != nil {
397 bs[0] |= 0x02
398 copy(bs[len(bs)-2*ElemBytesLen:], p.NodeAux.Key[:])
399 copy(bs[len(bs)-1*ElemBytesLen:], p.NodeAux.Value[:])
400 }
401 return bs
402}
403
404// SiblingsFromProof returns all the siblings of the proof. This function is used to generate the siblings input for the circom circuits.
405func SiblingsFromProof(proof *Proof) []*Hash {
406 sibIdx := 0
407 var siblings []*Hash
408 for lvl := 0; lvl < int(proof.depth); lvl++ {
409 if common.TestBitBigEndian(proof.notempties[:], uint(lvl)) {
410 siblings = append(siblings, proof.Siblings[sibIdx])
411 sibIdx++
412 } else {
413 siblings = append(siblings, &HashZero)
414 }
415 }
416 return siblings
417}
418
419func (p *Proof) AllSiblings() []*Hash {
420 return SiblingsFromProof(p)
421}
422
423func (p *Proof) AllSiblingsCircom(levels int) []*big.Int {
424 siblings := p.AllSiblings()
425 // Add the rest of empty levels to the siblings
426 for i := len(siblings); i < levels; i++ {
427 siblings = append(siblings, &HashZero)
428 }
429 siblings = append(siblings, &HashZero) // add extra level for circom compatibility
430 siblingsBigInt := make([]*big.Int, len(siblings))
431 for i, sibling := range siblings {
432 siblingsBigInt[i] = sibling.BigInt()
433 }
434 return siblingsBigInt
435}
436
437// GenerateProof generates the proof of existence (or non-existence) of an
438// Entry's hash Index for a Merkle Tree given the root.
439// If the rootKey is nil, the current merkletree root is used
440func (mt *MerkleTree) GenerateProof(k *big.Int, rootKey *Hash) (*Proof, error) {
441 p := &Proof{}
442 var siblingKey *Hash
443
444 kHash := NewHashFromBigInt(k)
445 path := getPath(mt.maxLevels, kHash[:])
446 if rootKey == nil {
447 rootKey = mt.Root()
448 }
449 nextKey := rootKey
450 for p.depth = 0; p.depth < uint(mt.maxLevels); p.depth++ {
451 n, err := mt.GetNode(nextKey)
452 if err != nil {
453 return nil, err
454 }
455 switch n.Type {
456 case NodeTypeEmpty:
457 return p, nil
458 case NodeTypeLeaf:
459 if bytes.Equal(kHash[:], n.Entry[0][:]) {
460 p.Existence = true
461 return p, nil
462 } else {
463 // We found a leaf whose entry didn't match hIndex
464 p.NodeAux = &NodeAux{Key: n.Entry[0], Value: n.Entry[1]}
465 return p, nil
466 }
467 case NodeTypeMiddle:
468 if path[p.depth] {
469 nextKey = n.ChildR
470 siblingKey = n.ChildL
471 } else {
472 nextKey = n.ChildL
473 siblingKey = n.ChildR
474 }
475 default:
476 return nil, ErrInvalidNodeFound
477 }
478 if !bytes.Equal(siblingKey[:], HashZero[:]) {
479 common.SetBitBigEndian(p.notempties[:], uint(p.depth))
480 p.Siblings = append(p.Siblings, siblingKey)
481 }
482 }
483 return nil, ErrEntryIndexNotFound
484}
485
486// VerifyProof verifies the Merkle Proof for the entry and root.
487func VerifyProof(rootKey *Hash, proof *Proof, k, v *big.Int) bool {
488 rootFromProof, err := RootFromProof(proof, k, v)
489 if err != nil {
490 return false
491 }
492 return bytes.Equal(rootKey[:], rootFromProof[:])
493}
494
495// RootFromProof calculates the root that would correspond to a tree whose
496// siblings are the ones in the proof with the claim hashing to hIndex and
497// hValue.
498func RootFromProof(proof *Proof, k, v *big.Int) (*Hash, error) {
499 kHash := NewHashFromBigInt(k)
500 vHash := NewHashFromBigInt(v)
501 sibIdx := len(proof.Siblings) - 1
502 var err error
503 var midKey *Hash
504 if proof.Existence {
505 midKey, err = LeafKey(kHash, vHash)
506 if err != nil {
507 return nil, err
508 }
509 } else {
510 if proof.NodeAux == nil {
511 midKey = &HashZero
512 } else {
513 if bytes.Equal(kHash[:], proof.NodeAux.Key[:]) {
514 return nil, fmt.Errorf("Non-existence proof being checked against hIndex equal to nodeAux")
515 }
516 midKey, err = LeafKey(proof.NodeAux.Key, proof.NodeAux.Value)
517 if err != nil {
518 return nil, err
519 }
520 }
521 }
522 path := getPath(int(proof.depth), kHash[:])
523 var siblingKey *Hash
524 for lvl := int(proof.depth) - 1; lvl >= 0; lvl-- {
525 if common.TestBitBigEndian(proof.notempties[:], uint(lvl)) {
526 siblingKey = proof.Siblings[sibIdx]
527 sibIdx--
528 } else {
529 siblingKey = &HashZero
530 }
531 if path[lvl] {
532 midKey, err = NewNodeMiddle(siblingKey, midKey).Key()
533 if err != nil {
534 return nil, err
535 }
536 } else {
537 midKey, err = NewNodeMiddle(midKey, siblingKey).Key()
538 if err != nil {
539 return nil, err
540 }
541 }
542 }
543 return midKey, nil
544}
545
546// walk is a helper recursive function to iterate over all tree branches
547func (mt *MerkleTree) walk(key *Hash, f func(*Node)) error {
548 n, err := mt.GetNode(key)
549 if err != nil {
550 return err
551 }
552 switch n.Type {
553 case NodeTypeEmpty:
554 f(n)
555 case NodeTypeLeaf:
556 f(n)
557 case NodeTypeMiddle:
558 f(n)
559 if err := mt.walk(n.ChildL, f); err != nil {
560 return err
561 }
562 if err := mt.walk(n.ChildR, f); err != nil {
563 return err
564 }
565 default:
566 return ErrInvalidNodeFound
567 }
568 return nil
569}
570
571// Walk iterates over all the branches of a MerkleTree with the given rootKey
572// if rootKey is nil, it will get the current RootKey of the current state of the MerkleTree.
573// For each node, it calls the f function given in the parameters.
574// See some examples of the Walk function usage in the merkletree_test.go
575// test functions: TestMTWalk, TestMTWalkGraphViz, TestMTWalkDumpClaims
576func (mt *MerkleTree) Walk(rootKey *Hash, f func(*Node)) error {
577 if rootKey == nil {
578 rootKey = mt.Root()
579 }
580 err := mt.walk(rootKey, f)
581 return err
582}
583
584// GraphViz uses Walk function to generate a string GraphViz representation of the
585// tree and writes it to w
586func (mt *MerkleTree) GraphViz(w io.Writer, rootKey *Hash) error {
587 fmt.Fprintf(w, `digraph hierarchy {
588node [fontname=Monospace,fontsize=10,shape=box]
589`)
590 cnt := 0
591 var errIn error
592 err := mt.Walk(rootKey, func(n *Node) {
593 k, err := n.Key()
594 if err != nil {
595 errIn = err
596 }
597 switch n.Type {
598 case NodeTypeEmpty:
599 case NodeTypeLeaf:
600 fmt.Fprintf(w, "\"%v\" [style=filled];\n", k.BigInt().String())
601 case NodeTypeMiddle:
602 lr := [2]string{n.ChildL.BigInt().String(), n.ChildR.BigInt().String()}
603 for i := range lr {
604 if lr[i] == "0" {
605 lr[i] = fmt.Sprintf("empty%v", cnt)
606 fmt.Fprintf(w, "\"%v\" [style=dashed,label=0];\n", lr[i])
607 cnt++
608 }
609 }
610 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", k.BigInt().String(), lr[0], lr[1])
611 default:
612 }
613 })
614 fmt.Fprintf(w, "}\n")
615 if errIn != nil {
616 return errIn
617 }
618 return err
619}
620
621// PrintGraphViz prints directly the GraphViz() output
622func (mt *MerkleTree) PrintGraphViz(rootKey *Hash) error {
623 if rootKey == nil {
624 rootKey = mt.Root()
625 }
626 w := bytes.NewBufferString("")
627 fmt.Fprintf(w, "--------\nGraphViz of the MerkleTree with RootKey "+rootKey.BigInt().String()+"\n")
628 err := mt.GraphViz(w, nil)
629 if err != nil {
630 return err
631 }
632 fmt.Fprintf(w, "End of GraphViz of the MerkleTree with RootKey "+rootKey.BigInt().String()+"\n--------\n")
633
634 fmt.Println(w)
635 return nil
636}