arbo/tree.go
tree.goRaw
1/*
2Package arbo implements a Merkle Tree compatible with the circomlib
3implementation of the MerkleTree, following the specification from
4https://docs.iden3.io/publications/pdfs/Merkle-Tree.pdf and
5https://eprint.iacr.org/2018/955.
6
7Allows to define which hash function to use. So for example, when working with
8zkSnarks the Poseidon hash function can be used, but when not, it can be used
9the Blake2b hash function, which has much faster computation time.
10*/
11package arbo
12
13import (
14 "bytes"
15 "encoding/binary"
16 "encoding/hex"
17 "fmt"
18 "io"
19 "math"
20 "sync"
21
22 "go.vocdoni.io/dvote/db"
23)
24
25const (
26 // PrefixValueLen defines the bytes-prefix length used for the Value
27 // bytes representation stored in the db
28 PrefixValueLen = 2
29
30 // PrefixValueEmpty is used for the first byte of a Value to indicate
31 // that is an Empty value
32 PrefixValueEmpty = 0
33 // PrefixValueLeaf is used for the first byte of a Value to indicate
34 // that is a Leaf value
35 PrefixValueLeaf = 1
36 // PrefixValueIntermediate is used for the first byte of a Value to
37 // indicate that is a Intermediate value
38 PrefixValueIntermediate = 2
39
40 // nChars is used to crop the Graphviz nodes labels
41 nChars = 4
42
43 maxKeyLen = 255
44 maxSiblingsFieldLen = 65535
45)
46
47var (
48 dbKeyRoot = []byte("root")
49 dbKeyNLeafs = []byte("nleafs")
50 emptyValue = []byte{0}
51
52 // ErrKeyNotFound is used when a key is not found in the db neither in
53 // the current db Batch.
54 ErrKeyNotFound = fmt.Errorf("key not found")
55 // ErrKeyAlreadyExists is used when trying to add a key as leaf to the
56 // tree that already exists.
57 ErrKeyAlreadyExists = fmt.Errorf("key already exists")
58 // ErrInvalidValuePrefix is used when going down into the tree, a value
59 // is read from the db and has an unrecognized prefix.
60 ErrInvalidValuePrefix = fmt.Errorf("invalid value prefix")
61 // ErrDBNoTx is used when trying to use Tree.dbPut but Tree.dbBatch==nil
62 ErrDBNoTx = fmt.Errorf("dbPut error: no db Batch")
63 // ErrMaxLevel indicates when going down into the tree, the max level is
64 // reached
65 ErrMaxLevel = fmt.Errorf("max level reached")
66 // ErrMaxVirtualLevel indicates when going down into the tree, the max
67 // virtual level is reached
68 ErrMaxVirtualLevel = fmt.Errorf("max virtual level reached")
69 // ErrSnapshotNotEditable indicates when the tree is a non writable
70 // snapshot, thus can not be modified
71 ErrSnapshotNotEditable = fmt.Errorf("snapshot tree can not be edited")
72 // ErrTreeNotEmpty indicates when the tree was expected to be empty and
73 // it is not
74 ErrTreeNotEmpty = fmt.Errorf("tree is not empty")
75)
76
77// Tree defines the struct that implements the MerkleTree functionalities
78type Tree struct {
79 sync.Mutex
80
81 db db.Database
82 maxLevels int
83 snapshotRoot []byte
84
85 hashFunction HashFunction
86 // TODO in the methods that use it, check if emptyHash param is len>0
87 // (check if it has been initialized)
88 emptyHash []byte
89
90 dbg *dbgStats
91}
92
93// NewTree returns a new Tree, if there is a Tree still in the given database, it
94// will load it.
95func NewTree(database db.Database, maxLevels int, hash HashFunction) (*Tree, error) {
96 wTx := database.WriteTx()
97 defer wTx.Discard()
98
99 t, err := NewTreeWithTx(wTx, database, maxLevels, hash)
100 if err != nil {
101 return nil, err
102 }
103
104 if err = wTx.Commit(); err != nil {
105 return nil, err
106 }
107 return t, nil
108}
109
110// NewTreeWithTx returns a new Tree using the given db.WriteTx, which will not
111// be ccommited inside this method, if there is a Tree still in the given
112// database, it will load it.
113func NewTreeWithTx(wTx db.WriteTx, database db.Database,
114 maxLevels int, hash HashFunction) (*Tree, error) {
115 t := Tree{db: database, maxLevels: maxLevels, hashFunction: hash}
116 t.emptyHash = make([]byte, t.hashFunction.Len()) // empty
117
118 _, err := wTx.Get(dbKeyRoot)
119 if err == db.ErrKeyNotFound {
120 // store new root 0 (empty)
121 if err = wTx.Set(dbKeyRoot, t.emptyHash); err != nil {
122 return nil, err
123 }
124 if err = t.setNLeafs(wTx, 0); err != nil {
125 return nil, err
126 }
127 return &t, nil
128 } else if err != nil {
129 return nil, err
130 }
131 return &t, nil
132}
133
134// Root returns the root of the Tree
135func (t *Tree) Root() ([]byte, error) {
136 rTx := t.db.ReadTx()
137 defer rTx.Discard()
138 return t.RootWithTx(rTx)
139}
140
141// RootWithTx returns the root of the Tree using the given db.ReadTx
142func (t *Tree) RootWithTx(rTx db.ReadTx) ([]byte, error) {
143 // if snapshotRoot is defined, means that the tree is a snapshot, and
144 // the root is not obtained from the db, but from the snapshotRoot
145 // parameter
146 if t.snapshotRoot != nil {
147 return t.snapshotRoot, nil
148 }
149 // get db root
150 return rTx.Get(dbKeyRoot)
151}
152
153func (t *Tree) setRoot(wTx db.WriteTx, root []byte) error {
154 return wTx.Set(dbKeyRoot, root)
155}
156
157// HashFunction returns Tree.hashFunction
158func (t *Tree) HashFunction() HashFunction {
159 return t.hashFunction
160}
161
162// editable returns true if the tree is editable, and false when is not
163// editable (because is a snapshot tree)
164func (t *Tree) editable() bool {
165 return t.snapshotRoot == nil
166}
167
168// AddBatch adds a batch of key-values to the Tree. Returns an array containing
169// the indexes of the keys failed to add. Supports empty values as input
170// parameters, which is equivalent to 0 valued byte array.
171func (t *Tree) AddBatch(keys, values [][]byte) ([]int, error) {
172 wTx := t.db.WriteTx()
173 defer wTx.Discard()
174
175 invalids, err := t.AddBatchWithTx(wTx, keys, values)
176 if err != nil {
177 return invalids, err
178 }
179 return invalids, wTx.Commit()
180}
181
182// AddBatchWithTx does the same than the AddBatch method, but allowing to pass
183// the db.WriteTx that is used. The db.WriteTx will not be committed inside
184// this method.
185func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]int, error) {
186 t.Lock()
187 defer t.Unlock()
188
189 if !t.editable() {
190 return nil, ErrSnapshotNotEditable
191 }
192
193 vt, err := t.loadVT(wTx)
194 if err != nil {
195 return nil, err
196 }
197
198 e := []byte{}
199 // equal the number of keys & values
200 if len(keys) > len(values) {
201 // add missing values
202 for i := len(values); i < len(keys); i++ {
203 values = append(values, e)
204 }
205 } else if len(keys) < len(values) {
206 // crop extra values
207 values = values[:len(keys)]
208 }
209
210 invalids, err := vt.addBatch(keys, values)
211 if err != nil {
212 return nil, err
213 }
214
215 // once the VirtualTree is build, compute the hashes
216 pairs, err := vt.computeHashes()
217 if err != nil {
218 // currently invalids in computeHashes are not counted,
219 // but should not be needed, as if there is an error there is
220 // nothing stored in the db and the error is returned
221 return nil, err
222 }
223
224 // store pairs in db
225 for i := 0; i < len(pairs); i++ {
226 if err := wTx.Set(pairs[i][0], pairs[i][1]); err != nil {
227 return nil, err
228 }
229 }
230
231 // store root (from the vt) to db
232 if err := wTx.Set(dbKeyRoot, vt.root.h); err != nil {
233 return nil, err
234 }
235
236 // update nLeafs
237 if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
238 return nil, err
239 }
240
241 return invalids, nil
242}
243
244// loadVT loads a new virtual tree (vt) from the current Tree, which contains
245// the same leafs.
246func (t *Tree) loadVT(rTx db.ReadTx) (vt, error) {
247 vt := newVT(t.maxLevels, t.hashFunction)
248 vt.params.dbg = t.dbg
249 var callbackErr error
250 err := t.IterateWithStopWithTx(rTx, nil, func(_ int, k, v []byte) bool {
251 if v[0] != PrefixValueLeaf {
252 return false
253 }
254 leafK, leafV := ReadLeafValue(v)
255 if err := vt.add(0, leafK, leafV); err != nil {
256 callbackErr = err
257 return true
258 }
259 return false
260 })
261 if callbackErr != nil {
262 return vt, callbackErr
263 }
264
265 return vt, err
266}
267
268// Add inserts the key-value into the Tree. If the inputs come from a
269// *big.Int, is expected that are represented by a Little-Endian byte array
270// (for circom compatibility).
271func (t *Tree) Add(k, v []byte) error {
272 wTx := t.db.WriteTx()
273 defer wTx.Discard()
274
275 if err := t.AddWithTx(wTx, k, v); err != nil {
276 return err
277 }
278
279 return wTx.Commit()
280}
281
282// AddWithTx does the same than the Add method, but allowing to pass the
283// db.WriteTx that is used. The db.WriteTx will not be committed inside this
284// method.
285func (t *Tree) AddWithTx(wTx db.WriteTx, k, v []byte) error {
286 t.Lock()
287 defer t.Unlock()
288
289 if !t.editable() {
290 return ErrSnapshotNotEditable
291 }
292
293 root, err := t.RootWithTx(wTx)
294 if err != nil {
295 return err
296 }
297
298 root, err = t.add(wTx, root, 0, k, v) // add from level 0
299 if err != nil {
300 return err
301 }
302 // store root to db
303 if err := t.setRoot(wTx, root); err != nil {
304 return err
305 }
306 // update nLeafs
307 if err = t.incNLeafs(wTx, 1); err != nil {
308 return err
309 }
310 return nil
311}
312
313func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) {
314 keyPath := make([]byte, t.hashFunction.Len())
315 if len(k) > len(keyPath) {
316 return nil, fmt.Errorf("len(k) > len(keyPath)")
317 }
318 copy(keyPath[:], k)
319
320 path := getPath(t.maxLevels, keyPath)
321 // go down to the leaf
322 var siblings [][]byte
323 _, _, siblings, err := t.down(wTx, k, root, siblings, path, fromLvl, false)
324 if err != nil {
325 return nil, err
326 }
327
328 leafKey, leafValue, err := t.newLeafValue(k, v)
329 if err != nil {
330 return nil, err
331 }
332
333 if err := wTx.Set(leafKey, leafValue); err != nil {
334 return nil, err
335 }
336
337 // go up to the root
338 if len(siblings) == 0 {
339 // return the leafKey as root
340 return leafKey, nil
341 }
342 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, fromLvl)
343 if err != nil {
344 return nil, err
345 }
346
347 return root, nil
348}
349
350// down goes down to the leaf recursively
351func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte,
352 path []bool, currLvl int, getLeaf bool) (
353 []byte, []byte, [][]byte, error) {
354 if currLvl > t.maxLevels {
355 return nil, nil, nil, ErrMaxLevel
356 }
357
358 var err error
359 var currValue []byte
360 if bytes.Equal(currKey, t.emptyHash) {
361 // empty value
362 return currKey, emptyValue, siblings, nil
363 }
364 currValue, err = rTx.Get(currKey)
365 if err != nil {
366 return nil, nil, nil, err
367 }
368
369 switch currValue[0] {
370 case PrefixValueEmpty: // empty
371 fmt.Printf("newKey: %s, currKey: %s, currLvl: %d, currValue: %s\n",
372 hex.EncodeToString(newKey), hex.EncodeToString(currKey),
373 currLvl, hex.EncodeToString(currValue))
374 panic("This point should not be reached, as the 'if' above" +
375 " should avoid reaching this point. This panic is temporary" +
376 " for reporting purposes, will be deleted in future versions." +
377 " Please paste this log (including the previous lines) in a" +
378 " new issue: https://github.com/vocdoni/arbo/issues/new") // TMP
379 case PrefixValueLeaf: // leaf
380 if !bytes.Equal(currValue, emptyValue) {
381 if getLeaf {
382 return currKey, currValue, siblings, nil
383 }
384 oldLeafKey, _ := ReadLeafValue(currValue)
385 if bytes.Equal(newKey, oldLeafKey) {
386 return nil, nil, nil, ErrKeyAlreadyExists
387 }
388
389 oldLeafKeyFull := make([]byte, t.hashFunction.Len())
390 if len(oldLeafKey) > len(oldLeafKeyFull) {
391 return nil, nil, nil,
392 fmt.Errorf("len(oldLeafKey) > len(oldLeafKeyFull)")
393 }
394 copy(oldLeafKeyFull[:], oldLeafKey)
395
396 // if currKey is already used, go down until paths diverge
397 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
398 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
399 if err != nil {
400 return nil, nil, nil, err
401 }
402 }
403 return currKey, currValue, siblings, nil
404 case PrefixValueIntermediate: // intermediate
405 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
406 return nil, nil, nil,
407 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
408 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
409 }
410 // collect siblings while going down
411 if path[currLvl] {
412 // right
413 lChild, rChild := ReadIntermediateChilds(currValue)
414 siblings = append(siblings, lChild)
415 return t.down(rTx, newKey, rChild, siblings, path, currLvl+1, getLeaf)
416 }
417 // left
418 lChild, rChild := ReadIntermediateChilds(currValue)
419 siblings = append(siblings, rChild)
420 return t.down(rTx, newKey, lChild, siblings, path, currLvl+1, getLeaf)
421 default:
422 return nil, nil, nil, ErrInvalidValuePrefix
423 }
424}
425
426// downVirtually is used when in a leaf already exists, and a new leaf which
427// shares the path until the existing leaf is being added
428func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
429 newPath []bool, currLvl int) ([][]byte, error) {
430 var err error
431 if currLvl > t.maxLevels-1 {
432 return nil, ErrMaxVirtualLevel
433 }
434
435 if oldPath[currLvl] == newPath[currLvl] {
436 siblings = append(siblings, t.emptyHash)
437
438 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
439 if err != nil {
440 return nil, err
441 }
442 return siblings, nil
443 }
444 // reached the divergence
445 siblings = append(siblings, oldKey)
446
447 return siblings, nil
448}
449
450// up goes up recursively updating the intermediate nodes
451func (t *Tree) up(wTx db.WriteTx, key []byte, siblings [][]byte, path []bool,
452 currLvl, toLvl int) ([]byte, error) {
453 var k, v []byte
454 var err error
455 if path[currLvl+toLvl] {
456 k, v, err = t.newIntermediate(siblings[currLvl], key)
457 if err != nil {
458 return nil, err
459 }
460 } else {
461 k, v, err = t.newIntermediate(key, siblings[currLvl])
462 if err != nil {
463 return nil, err
464 }
465 }
466 // store k-v to db
467 if err = wTx.Set(k, v); err != nil {
468 return nil, err
469 }
470
471 if currLvl == 0 {
472 // reached the root
473 return k, nil
474 }
475
476 return t.up(wTx, k, siblings, path, currLvl-1, toLvl)
477}
478
479func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
480 t.dbg.incHash()
481 return newLeafValue(t.hashFunction, k, v)
482}
483
484// newLeafValue takes a key & value from a leaf, and computes the leaf hash,
485// which is used as the leaf key. And the value is the concatenation of the
486// inputed key & value. The output of this function is used as key-value to
487// store the leaf in the DB.
488// [ 1 byte | 1 byte | N bytes | M bytes ]
489// [ type of node | length of key | key | value ]
490func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
491 leafKey, err := hashFunc.Hash(k, v, []byte{1})
492 if err != nil {
493 return nil, nil, err
494 }
495 var leafValue []byte
496 leafValue = append(leafValue, byte(PrefixValueLeaf))
497 if len(k) > maxKeyLen {
498 return nil, nil, fmt.Errorf("len(k) > %v", maxKeyLen)
499 }
500 leafValue = append(leafValue, byte(len(k)))
501 leafValue = append(leafValue, k...)
502 leafValue = append(leafValue, v...)
503 return leafKey, leafValue, nil
504}
505
506// ReadLeafValue reads from a byte array the leaf key & value
507func ReadLeafValue(b []byte) ([]byte, []byte) {
508 if len(b) < PrefixValueLen {
509 return []byte{}, []byte{}
510 }
511
512 kLen := b[1]
513 if len(b) < PrefixValueLen+int(kLen) {
514 return []byte{}, []byte{}
515 }
516 k := b[PrefixValueLen : PrefixValueLen+kLen]
517 v := b[PrefixValueLen+kLen:]
518 return k, v
519}
520
521func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
522 t.dbg.incHash()
523 return newIntermediate(t.hashFunction, l, r)
524}
525
526// newIntermediate takes the left & right keys of a intermediate node, and
527// computes its hash. Returns the hash of the node, which is the node key, and a
528// byte array that contains the value (which contains the left & right child
529// keys) to store in the DB.
530// [ 1 byte | 1 byte | N bytes | N bytes ]
531// [ type of node | length of left key | left key | right key ]
532func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
533 b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
534 b[0] = PrefixValueIntermediate
535 if len(l) > maxKeyLen {
536 return nil, nil, fmt.Errorf("len(k) > %v", maxKeyLen)
537 }
538 b[1] = byte(len(l))
539 copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
540 copy(b[PrefixValueLen+hashFunc.Len():], r)
541
542 key, err := hashFunc.Hash(l, r)
543 if err != nil {
544 return nil, nil, err
545 }
546
547 return key, b, nil
548}
549
550// ReadIntermediateChilds reads from a byte array the two childs keys
551func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
552 if len(b) < PrefixValueLen {
553 return []byte{}, []byte{}
554 }
555
556 lLen := b[1]
557 if len(b) < PrefixValueLen+int(lLen) {
558 return []byte{}, []byte{}
559 }
560 l := b[PrefixValueLen : PrefixValueLen+lLen]
561 r := b[PrefixValueLen+lLen:]
562 return l, r
563}
564
565func getPath(numLevels int, k []byte) []bool {
566 path := make([]bool, numLevels)
567 for n := 0; n < numLevels; n++ {
568 path[n] = k[n/8]&(1<<(n%8)) != 0
569 }
570 return path
571}
572
573// Update updates the value for a given existing key. If the given key does not
574// exist, returns an error.
575func (t *Tree) Update(k, v []byte) error {
576 wTx := t.db.WriteTx()
577 defer wTx.Discard()
578
579 if err := t.UpdateWithTx(wTx, k, v); err != nil {
580 return err
581 }
582 return wTx.Commit()
583}
584
585// UpdateWithTx does the same than the Update method, but allowing to pass the
586// db.WriteTx that is used. The db.WriteTx will not be committed inside this
587// method.
588func (t *Tree) UpdateWithTx(wTx db.WriteTx, k, v []byte) error {
589 t.Lock()
590 defer t.Unlock()
591
592 if !t.editable() {
593 return ErrSnapshotNotEditable
594 }
595
596 var err error
597
598 keyPath := make([]byte, t.hashFunction.Len())
599 if len(k) > len(keyPath) {
600 return fmt.Errorf("len(k) > len(keyPath)")
601 }
602 copy(keyPath[:], k)
603 path := getPath(t.maxLevels, keyPath)
604
605 root, err := t.RootWithTx(wTx)
606 if err != nil {
607 return err
608 }
609
610 var siblings [][]byte
611 _, valueAtBottom, siblings, err := t.down(wTx, k, root, siblings, path, 0, true)
612 if err != nil {
613 return err
614 }
615 oldKey, _ := ReadLeafValue(valueAtBottom)
616 if !bytes.Equal(oldKey, k) {
617 return ErrKeyNotFound
618 }
619
620 leafKey, leafValue, err := t.newLeafValue(k, v)
621 if err != nil {
622 return err
623 }
624
625 if err := wTx.Set(leafKey, leafValue); err != nil {
626 return err
627 }
628
629 // go up to the root
630 if len(siblings) == 0 {
631 return t.setRoot(wTx, leafKey)
632 }
633 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, 0)
634 if err != nil {
635 return err
636 }
637
638 // store root to db
639 if err := t.setRoot(wTx, root); err != nil {
640 return err
641 }
642 return nil
643}
644
645// GenProof generates a MerkleTree proof for the given key. The leaf value is
646// returned, together with the packed siblings of the proof, and a boolean
647// parameter that indicates if the proof is of existence (true) or not (false).
648func (t *Tree) GenProof(k []byte) ([]byte, []byte, []byte, bool, error) {
649 rTx := t.db.ReadTx()
650 defer rTx.Discard()
651
652 return t.GenProofWithTx(rTx, k)
653}
654
655// GenProofWithTx does the same than the GenProof method, but allowing to pass
656// the db.ReadTx that is used.
657func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) {
658 keyPath := make([]byte, t.hashFunction.Len())
659 if len(k) > len(keyPath) {
660 return nil, nil, nil, false, fmt.Errorf("len(k) > len(keyPath)")
661 }
662 copy(keyPath[:], k)
663
664 root, err := t.RootWithTx(rTx)
665 if err != nil {
666 return nil, nil, nil, false, err
667 }
668
669 path := getPath(t.maxLevels, keyPath)
670 // go down to the leaf
671 var siblings [][]byte
672 _, value, siblings, err := t.down(rTx, k, root, siblings, path, 0, true)
673 if err != nil {
674 return nil, nil, nil, false, err
675 }
676
677 s, err := PackSiblings(t.hashFunction, siblings)
678 if err != nil {
679 return nil, nil, nil, false, err
680 }
681
682 leafK, leafV := ReadLeafValue(value)
683 if !bytes.Equal(k, leafK) {
684 // key not in tree, proof of non-existence
685 return leafK, leafV, s, false, nil
686 }
687
688 return leafK, leafV, s, true, nil
689}
690
691// PackSiblings packs the siblings into a byte array.
692// [ 2 byte | 2 byte | L bytes | S * N bytes ]
693// [ full length | bitmap length (L) | bitmap | N non-zero siblings ]
694// Where the bitmap indicates if the sibling is 0 or a value from the siblings
695// array. And S is the size of the output of the hash function used for the
696// Tree. The 2 2-byte that define the full length and bitmap length, are
697// encoded in little-endian.
698func PackSiblings(hashFunc HashFunction, siblings [][]byte) ([]byte, error) {
699 var b []byte
700 var bitmap []bool
701 emptySibling := make([]byte, hashFunc.Len())
702 for i := 0; i < len(siblings); i++ {
703 if bytes.Equal(siblings[i], emptySibling) {
704 bitmap = append(bitmap, false)
705 } else {
706 bitmap = append(bitmap, true)
707 b = append(b, siblings[i]...)
708 }
709 }
710
711 bitmapBytes := bitmapToBytes(bitmap)
712 l := len(bitmapBytes)
713
714 fullLen := 4 + l + len(b) //nolint:gomnd
715 res := make([]byte, fullLen)
716 if fullLen > maxSiblingsFieldLen {
717 return nil, fmt.Errorf("fullLen > %v", maxSiblingsFieldLen)
718 }
719 binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length
720 if l > maxSiblingsFieldLen {
721 return nil, fmt.Errorf("l > %v", maxSiblingsFieldLen)
722 }
723 binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length
724 copy(res[4:4+l], bitmapBytes)
725 copy(res[4+l:], b)
726 return res, nil
727}
728
729// UnpackSiblings unpacks the siblings from a byte array.
730func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
731 fullLen := binary.LittleEndian.Uint16(b[0:2])
732 l := binary.LittleEndian.Uint16(b[2:4]) // bitmap bytes length
733 if len(b) != int(fullLen) {
734 return nil,
735 fmt.Errorf("error unpacking siblings. Expected len: %d, current len: %d",
736 fullLen, len(b))
737 }
738
739 bitmapBytes := b[4 : 4+l]
740 bitmap := bytesToBitmap(bitmapBytes)
741 siblingsBytes := b[4+l:]
742 iSibl := 0
743 emptySibl := make([]byte, hashFunc.Len())
744 var siblings [][]byte
745 for i := 0; i < len(bitmap); i++ {
746 if iSibl >= len(siblingsBytes) {
747 break
748 }
749 if bitmap[i] {
750 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
751 iSibl += hashFunc.Len()
752 } else {
753 siblings = append(siblings, emptySibl)
754 }
755 }
756 return siblings, nil
757}
758
759func bitmapToBytes(bitmap []bool) []byte {
760 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
761 b := make([]byte, bitmapBytesLen)
762 for i := 0; i < len(bitmap); i++ {
763 if bitmap[i] {
764 b[i/8] |= 1 << (i % 8)
765 }
766 }
767 return b
768}
769
770func bytesToBitmap(b []byte) []bool {
771 var bitmap []bool
772 for i := 0; i < len(b); i++ {
773 for j := 0; j < 8; j++ {
774 bitmap = append(bitmap, b[i]&(1<<j) > 0)
775 }
776 }
777 return bitmap
778}
779
780// Get returns the value in the Tree for a given key. If the key is not found,
781// will return the error ErrKeyNotFound, and in the leafK & leafV parameters
782// will be placed the data found in the tree in the leaf that was on the path
783// going to the input key.
784func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
785 rTx := t.db.ReadTx()
786 defer rTx.Discard()
787
788 return t.GetWithTx(rTx, k)
789}
790
791// GetWithTx does the same than the Get method, but allowing to pass the
792// db.ReadTx that is used. If the key is not found, will return the error
793// ErrKeyNotFound, and in the leafK & leafV parameters will be placed the data
794// found in the tree in the leaf that was on the path going to the input key.
795func (t *Tree) GetWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, error) {
796 keyPath := make([]byte, t.hashFunction.Len())
797 copy(keyPath[:], k)
798
799 root, err := t.RootWithTx(rTx)
800 if err != nil {
801 return nil, nil, err
802 }
803
804 path := getPath(t.maxLevels, keyPath)
805 // go down to the leaf
806 var siblings [][]byte
807 _, value, _, err := t.down(rTx, k, root, siblings, path, 0, true)
808 if err != nil {
809 return nil, nil, err
810 }
811 leafK, leafV := ReadLeafValue(value)
812 if !bytes.Equal(k, leafK) {
813 return leafK, leafV, ErrKeyNotFound
814 }
815
816 return leafK, leafV, nil
817}
818
819// CheckProof verifies the given proof. The proof verification depends on the
820// HashFunction passed as parameter.
821func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
822 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
823 if err != nil {
824 return false, err
825 }
826
827 keyPath := make([]byte, hashFunc.Len())
828 copy(keyPath[:], k)
829
830 key, _, err := newLeafValue(hashFunc, k, v)
831 if err != nil {
832 return false, err
833 }
834
835 path := getPath(len(siblings), keyPath)
836 for i := len(siblings) - 1; i >= 0; i-- {
837 if path[i] {
838 key, _, err = newIntermediate(hashFunc, siblings[i], key)
839 if err != nil {
840 return false, err
841 }
842 } else {
843 key, _, err = newIntermediate(hashFunc, key, siblings[i])
844 if err != nil {
845 return false, err
846 }
847 }
848 }
849 if bytes.Equal(key[:], root) {
850 return true, nil
851 }
852 return false, nil
853}
854
855func (t *Tree) incNLeafs(wTx db.WriteTx, nLeafs int) error {
856 oldNLeafs, err := t.GetNLeafsWithTx(wTx)
857 if err != nil {
858 return err
859 }
860 newNLeafs := oldNLeafs + nLeafs
861 return t.setNLeafs(wTx, newNLeafs)
862}
863
864func (t *Tree) setNLeafs(wTx db.WriteTx, nLeafs int) error {
865 b := make([]byte, 8)
866 binary.LittleEndian.PutUint64(b, uint64(nLeafs))
867 if err := wTx.Set(dbKeyNLeafs, b); err != nil {
868 return err
869 }
870 return nil
871}
872
873// GetNLeafs returns the number of Leafs of the Tree.
874func (t *Tree) GetNLeafs() (int, error) {
875 rTx := t.db.ReadTx()
876 defer rTx.Discard()
877
878 return t.GetNLeafsWithTx(rTx)
879}
880
881// GetNLeafsWithTx does the same than the GetNLeafs method, but allowing to
882// pass the db.ReadTx that is used.
883func (t *Tree) GetNLeafsWithTx(rTx db.ReadTx) (int, error) {
884 b, err := rTx.Get(dbKeyNLeafs)
885 if err != nil {
886 return 0, err
887 }
888 nLeafs := binary.LittleEndian.Uint64(b)
889 return int(nLeafs), nil
890}
891
892// SetRoot sets the root to the given root
893func (t *Tree) SetRoot(root []byte) error {
894 wTx := t.db.WriteTx()
895 defer wTx.Discard()
896
897 if err := t.SetRootWithTx(wTx, root); err != nil {
898 return err
899 }
900 return wTx.Commit()
901}
902
903// SetRootWithTx sets the root to the given root using the given db.WriteTx
904func (t *Tree) SetRootWithTx(wTx db.WriteTx, root []byte) error {
905 if !t.editable() {
906 return ErrSnapshotNotEditable
907 }
908
909 if root == nil {
910 return fmt.Errorf("can not SetRoot with nil root")
911 }
912
913 // check that the root exists in the db
914 if !bytes.Equal(root, t.emptyHash) {
915 if _, err := wTx.Get(root); err == ErrKeyNotFound {
916 return fmt.Errorf("can not SetRoot with root %x, as it does not exist in the db", root)
917 } else if err != nil {
918 return err
919 }
920 }
921
922 return wTx.Set(dbKeyRoot, root)
923}
924
925// Snapshot returns a read-only copy of the Tree from the given root
926func (t *Tree) Snapshot(fromRoot []byte) (*Tree, error) {
927 // allow to define which root to use
928 if fromRoot == nil {
929 var err error
930 fromRoot, err = t.Root()
931 if err != nil {
932 return nil, err
933 }
934 }
935 rTx := t.db.ReadTx()
936 defer rTx.Discard()
937 // check that the root exists in the db
938 if !bytes.Equal(fromRoot, t.emptyHash) {
939 if _, err := rTx.Get(fromRoot); err == ErrKeyNotFound {
940 return nil,
941 fmt.Errorf("can not do a Snapshot with root %x, as it does not exist in the db",
942 fromRoot)
943 } else if err != nil {
944 return nil, err
945 }
946 }
947
948 return &Tree{
949 db: t.db,
950 maxLevels: t.maxLevels,
951 snapshotRoot: fromRoot,
952 emptyHash: t.emptyHash,
953 hashFunction: t.hashFunction,
954 dbg: t.dbg,
955 }, nil
956}
957
958// Iterate iterates through the full Tree, executing the given function on each
959// node of the Tree.
960func (t *Tree) Iterate(fromRoot []byte, f func([]byte, []byte)) error {
961 rTx := t.db.ReadTx()
962 defer rTx.Discard()
963
964 return t.IterateWithTx(rTx, fromRoot, f)
965}
966
967// IterateWithTx does the same than the Iterate method, but allowing to pass
968// the db.ReadTx that is used.
969func (t *Tree) IterateWithTx(rTx db.ReadTx, fromRoot []byte, f func([]byte, []byte)) error {
970 // allow to define which root to use
971 if fromRoot == nil {
972 var err error
973 fromRoot, err = t.RootWithTx(rTx)
974 if err != nil {
975 return err
976 }
977 }
978 return t.iter(rTx, fromRoot, f)
979}
980
981// IterateWithStop does the same than Iterate, but with int for the current
982// level, and a boolean parameter used by the passed function, is to indicate to
983// stop iterating on the branch when the method returns 'true'.
984func (t *Tree) IterateWithStop(fromRoot []byte, f func(int, []byte, []byte) bool) error {
985 rTx := t.db.ReadTx()
986 defer rTx.Discard()
987
988 // allow to define which root to use
989 if fromRoot == nil {
990 var err error
991 fromRoot, err = t.RootWithTx(rTx)
992 if err != nil {
993 return err
994 }
995 }
996 return t.iterWithStop(rTx, fromRoot, 0, f)
997}
998
999// IterateWithStopWithTx does the same than the IterateWithStop method, but
1000// allowing to pass the db.ReadTx that is used.
1001func (t *Tree) IterateWithStopWithTx(rTx db.ReadTx, fromRoot []byte,
1002 f func(int, []byte, []byte) bool) error {
1003 // allow to define which root to use
1004 if fromRoot == nil {
1005 var err error
1006 fromRoot, err = t.RootWithTx(rTx)
1007 if err != nil {
1008 return err
1009 }
1010 }
1011 return t.iterWithStop(rTx, fromRoot, 0, f)
1012}
1013
1014func (t *Tree) iterWithStop(rTx db.ReadTx, k []byte, currLevel int,
1015 f func(int, []byte, []byte) bool) error {
1016 var v []byte
1017 var err error
1018 if bytes.Equal(k, t.emptyHash) {
1019 v = t.emptyHash
1020 } else {
1021 v, err = rTx.Get(k)
1022 if err != nil {
1023 return err
1024 }
1025 }
1026 currLevel++
1027
1028 switch v[0] {
1029 case PrefixValueEmpty:
1030 f(currLevel, k, v)
1031 case PrefixValueLeaf:
1032 f(currLevel, k, v)
1033 case PrefixValueIntermediate:
1034 stop := f(currLevel, k, v)
1035 if stop {
1036 return nil
1037 }
1038 l, r := ReadIntermediateChilds(v)
1039 if err = t.iterWithStop(rTx, l, currLevel, f); err != nil {
1040 return err
1041 }
1042 if err = t.iterWithStop(rTx, r, currLevel, f); err != nil {
1043 return err
1044 }
1045 default:
1046 return ErrInvalidValuePrefix
1047 }
1048 return nil
1049}
1050
1051func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
1052 f2 := func(currLvl int, k, v []byte) bool {
1053 f(k, v)
1054 return false
1055 }
1056 return t.iterWithStop(rTx, k, 0, f2)
1057}
1058
1059// Dump exports all the Tree leafs in a byte array of length:
1060// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
1061// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
1062// [ len(k) | len(v) | key | value ]
1063// Where S is the size of the output of the hash function used for the Tree.
1064func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
1065 // allow to define which root to use
1066 if fromRoot == nil {
1067 var err error
1068 fromRoot, err = t.Root()
1069 if err != nil {
1070 return nil, err
1071 }
1072 }
1073
1074 // WARNING current encoding only supports key & values of 255 bytes each
1075 // (due using only 1 byte for the length headers).
1076 var b []byte
1077 var callbackErr error
1078 err := t.IterateWithStop(fromRoot, func(_ int, k, v []byte) bool {
1079 if v[0] != PrefixValueLeaf {
1080 return false
1081 }
1082 leafK, leafV := ReadLeafValue(v)
1083 kv := make([]byte, 2+len(leafK)+len(leafV))
1084 if len(leafK) > maxKeyLen {
1085 callbackErr = fmt.Errorf("len(leafK) > %v", maxKeyLen)
1086 return true
1087 }
1088 kv[0] = byte(len(leafK))
1089 if len(leafV) > maxKeyLen {
1090 callbackErr = fmt.Errorf("len(leafV) > %v", maxKeyLen)
1091 return true
1092 }
1093 kv[1] = byte(len(leafV))
1094 copy(kv[2:2+len(leafK)], leafK)
1095 copy(kv[2+len(leafK):], leafV)
1096 b = append(b, kv...)
1097 return false
1098 })
1099 if callbackErr != nil {
1100 return nil, callbackErr
1101 }
1102 return b, err
1103}
1104
1105// ImportDump imports the leafs (that have been exported with the Dump method)
1106// in the Tree.
1107func (t *Tree) ImportDump(b []byte) error {
1108 if !t.editable() {
1109 return ErrSnapshotNotEditable
1110 }
1111 root, err := t.Root()
1112 if err != nil {
1113 return err
1114 }
1115 if !bytes.Equal(root, t.emptyHash) {
1116 return ErrTreeNotEmpty
1117 }
1118
1119 r := bytes.NewReader(b)
1120 var keys, values [][]byte
1121 for {
1122 l := make([]byte, 2)
1123 _, err = io.ReadFull(r, l)
1124 if err == io.EOF {
1125 break
1126 } else if err != nil {
1127 return err
1128 }
1129 k := make([]byte, l[0])
1130 _, err = io.ReadFull(r, k)
1131 if err != nil {
1132 return err
1133 }
1134 v := make([]byte, l[1])
1135 _, err = io.ReadFull(r, v)
1136 if err != nil {
1137 return err
1138 }
1139 keys = append(keys, k)
1140 values = append(values, v)
1141 }
1142 if _, err = t.AddBatch(keys, values); err != nil {
1143 return err
1144 }
1145 return nil
1146}
1147
1148// Graphviz iterates across the full tree to generate a string Graphviz
1149// representation of the tree and writes it to w
1150func (t *Tree) Graphviz(w io.Writer, fromRoot []byte) error {
1151 return t.GraphvizFirstNLevels(w, fromRoot, t.maxLevels)
1152}
1153
1154// GraphvizFirstNLevels iterates across the first NLevels of the tree to
1155// generate a string Graphviz representation of the first NLevels of the tree
1156// and writes it to w
1157func (t *Tree) GraphvizFirstNLevels(w io.Writer, fromRoot []byte, untilLvl int) error {
1158 fmt.Fprintf(w, `digraph hierarchy {
1159node [fontname=Monospace,fontsize=10,shape=box]
1160`)
1161
1162 rTx := t.db.ReadTx()
1163 defer rTx.Discard()
1164
1165 if fromRoot == nil {
1166 var err error
1167 fromRoot, err = t.RootWithTx(rTx)
1168 if err != nil {
1169 return err
1170 }
1171 }
1172
1173 nEmpties := 0
1174 err := t.iterWithStop(rTx, fromRoot, 0, func(currLvl int, k, v []byte) bool {
1175 if currLvl == untilLvl {
1176 return true // to stop the iter from going down
1177 }
1178 switch v[0] {
1179 case PrefixValueEmpty:
1180 case PrefixValueLeaf:
1181 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
1182 // key & value from the leaf
1183 kB, vB := ReadLeafValue(v)
1184 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
1185 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
1186 hex.EncodeToString(vB[:nChars]))
1187 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
1188 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
1189 case PrefixValueIntermediate:
1190 l, r := ReadIntermediateChilds(v)
1191 lStr := hex.EncodeToString(l[:nChars])
1192 rStr := hex.EncodeToString(r[:nChars])
1193 eStr := ""
1194 if bytes.Equal(l, t.emptyHash) {
1195 lStr = fmt.Sprintf("empty%v", nEmpties)
1196 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1197 lStr)
1198 nEmpties++
1199 }
1200 if bytes.Equal(r, t.emptyHash) {
1201 rStr = fmt.Sprintf("empty%v", nEmpties)
1202 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1203 rStr)
1204 nEmpties++
1205 }
1206 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
1207 lStr, rStr)
1208 fmt.Fprint(w, eStr)
1209 default:
1210 }
1211 return false
1212 })
1213 fmt.Fprintf(w, "}\n")
1214 return err
1215}
1216
1217// PrintGraphviz prints the output of Tree.Graphviz
1218func (t *Tree) PrintGraphviz(fromRoot []byte) error {
1219 if fromRoot == nil {
1220 var err error
1221 fromRoot, err = t.Root()
1222 if err != nil {
1223 return err
1224 }
1225 }
1226 return t.PrintGraphvizFirstNLevels(fromRoot, t.maxLevels)
1227}
1228
1229// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
1230func (t *Tree) PrintGraphvizFirstNLevels(fromRoot []byte, untilLvl int) error {
1231 if fromRoot == nil {
1232 var err error
1233 fromRoot, err = t.Root()
1234 if err != nil {
1235 return err
1236 }
1237 }
1238 w := bytes.NewBufferString("")
1239 fmt.Fprintf(w,
1240 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+":\n")
1241 err := t.GraphvizFirstNLevels(w, fromRoot, untilLvl)
1242 if err != nil {
1243 fmt.Println(w)
1244 return err
1245 }
1246 fmt.Fprintf(w,
1247 "End of Graphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+"\n--------\n")
1248
1249 fmt.Println(w)
1250 return nil
1251}
1252
1253// TODO circom proofs
1254// TODO data structure for proofs (including root, key, value, siblings,
1255// hashFunction) + method to verify that data structure