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