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