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