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