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