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 maxUint8 = int(^uint8(0)) // 2**8 -1
44 maxUint16 = int(^uint16(0)) // 2**16 -1
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, int(math.Ceil(float64(t.maxLevels)/float64(8)))) //nolint:gomnd
315 copy(keyPath[:], k)
316 path := getPath(t.maxLevels, keyPath)
317
318 // go down to the leaf
319 var siblings [][]byte
320 _, _, siblings, err := t.down(wTx, k, root, siblings, path, fromLvl, false)
321 if err != nil {
322 return nil, err
323 }
324
325 leafKey, leafValue, err := t.newLeafValue(k, v)
326 if err != nil {
327 return nil, err
328 }
329
330 if err := wTx.Set(leafKey, leafValue); err != nil {
331 return nil, err
332 }
333
334 // go up to the root
335 if len(siblings) == 0 {
336 // return the leafKey as root
337 return leafKey, nil
338 }
339 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, fromLvl)
340 if err != nil {
341 return nil, err
342 }
343
344 return root, nil
345}
346
347// down goes down to the leaf recursively
348func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte,
349 path []bool, currLvl int, getLeaf bool) (
350 []byte, []byte, [][]byte, error) {
351 if currLvl > t.maxLevels {
352 return nil, nil, nil, ErrMaxLevel
353 }
354
355 var err error
356 var currValue []byte
357 if bytes.Equal(currKey, t.emptyHash) {
358 // empty value
359 return currKey, emptyValue, siblings, nil
360 }
361 currValue, err = rTx.Get(currKey)
362 if err != nil {
363 return nil, nil, nil, err
364 }
365
366 switch currValue[0] {
367 case PrefixValueEmpty: // empty
368 fmt.Printf("newKey: %s, currKey: %s, currLvl: %d, currValue: %s\n",
369 hex.EncodeToString(newKey), hex.EncodeToString(currKey),
370 currLvl, hex.EncodeToString(currValue))
371 panic("This point should not be reached, as the 'if currKey==t.emptyHash'" +
372 " above should avoid reaching this point. This panic is temporary" +
373 " for reporting purposes, will be deleted in future versions." +
374 " Please paste this log (including the previous log lines) in a" +
375 " new issue: https://github.com/vocdoni/arbo/issues/new") // TMP
376 case PrefixValueLeaf: // leaf
377 if !bytes.Equal(currValue, emptyValue) {
378 if getLeaf {
379 return currKey, currValue, siblings, nil
380 }
381 oldLeafKey, _ := ReadLeafValue(currValue)
382 if bytes.Equal(newKey, oldLeafKey) {
383 return nil, nil, nil, ErrKeyAlreadyExists
384 }
385
386 oldLeafKeyFull := make([]byte, t.hashFunction.Len())
387 // if len(oldLeafKey) > t.hashFunction.Len() { // WIP
388 // return nil, nil, nil,
389 // fmt.Errorf("len(oldLeafKey) > hashFunction.Len()")
390 // }
391 copy(oldLeafKeyFull[:], oldLeafKey)
392
393 // if currKey is already used, go down until paths diverge
394 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
395 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
396 if err != nil {
397 return nil, nil, nil, err
398 }
399 }
400 return currKey, currValue, siblings, nil
401 case PrefixValueIntermediate: // intermediate
402 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
403 return nil, nil, nil,
404 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
405 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
406 }
407 // collect siblings while going down
408 if path[currLvl] {
409 // right
410 lChild, rChild := ReadIntermediateChilds(currValue)
411 siblings = append(siblings, lChild)
412 return t.down(rTx, newKey, rChild, siblings, path, currLvl+1, getLeaf)
413 }
414 // left
415 lChild, rChild := ReadIntermediateChilds(currValue)
416 siblings = append(siblings, rChild)
417 return t.down(rTx, newKey, lChild, siblings, path, currLvl+1, getLeaf)
418 default:
419 return nil, nil, nil, ErrInvalidValuePrefix
420 }
421}
422
423// downVirtually is used when in a leaf already exists, and a new leaf which
424// shares the path until the existing leaf is being added
425func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
426 newPath []bool, currLvl int) ([][]byte, error) {
427 var err error
428 if currLvl > t.maxLevels-1 {
429 return nil, ErrMaxVirtualLevel
430 }
431
432 if oldPath[currLvl] == newPath[currLvl] {
433 siblings = append(siblings, t.emptyHash)
434
435 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
436 if err != nil {
437 return nil, err
438 }
439 return siblings, nil
440 }
441 // reached the divergence
442 siblings = append(siblings, oldKey)
443
444 return siblings, nil
445}
446
447// up goes up recursively updating the intermediate nodes
448func (t *Tree) up(wTx db.WriteTx, key []byte, siblings [][]byte, path []bool,
449 currLvl, toLvl int) ([]byte, error) {
450 var k, v []byte
451 var err error
452 if path[currLvl+toLvl] {
453 k, v, err = t.newIntermediate(siblings[currLvl], key)
454 if err != nil {
455 return nil, err
456 }
457 } else {
458 k, v, err = t.newIntermediate(key, siblings[currLvl])
459 if err != nil {
460 return nil, err
461 }
462 }
463 // store k-v to db
464 if err = wTx.Set(k, v); err != nil {
465 return nil, err
466 }
467
468 if currLvl == 0 {
469 // reached the root
470 return k, nil
471 }
472
473 return t.up(wTx, k, siblings, path, currLvl-1, toLvl)
474}
475
476func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
477 t.dbg.incHash()
478 return newLeafValue(t.hashFunction, k, v)
479}
480
481// newLeafValue takes a key & value from a leaf, and computes the leaf hash,
482// which is used as the leaf key. And the value is the concatenation of the
483// inputed key & value. The output of this function is used as key-value to
484// store the leaf in the DB.
485// [ 1 byte | 1 byte | N bytes | M bytes ]
486// [ type of node | length of key | key | value ]
487func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
488 leafKey, err := hashFunc.Hash(k, v, []byte{1})
489 if err != nil {
490 return nil, nil, err
491 }
492 var leafValue []byte
493 leafValue = append(leafValue, byte(PrefixValueLeaf))
494 if len(k) > maxUint8 {
495 return nil, nil, fmt.Errorf("newLeafValue: len(k) > %v", maxUint8)
496 }
497 leafValue = append(leafValue, byte(len(k)))
498 leafValue = append(leafValue, k...)
499 leafValue = append(leafValue, v...)
500 return leafKey, leafValue, nil
501}
502
503// ReadLeafValue reads from a byte array the leaf key & value
504func ReadLeafValue(b []byte) ([]byte, []byte) {
505 if len(b) < PrefixValueLen {
506 return []byte{}, []byte{}
507 }
508
509 kLen := b[1]
510 if len(b) < PrefixValueLen+int(kLen) {
511 return []byte{}, []byte{}
512 }
513 k := b[PrefixValueLen : PrefixValueLen+kLen]
514 v := b[PrefixValueLen+kLen:]
515 return k, v
516}
517
518func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
519 t.dbg.incHash()
520 return newIntermediate(t.hashFunction, l, r)
521}
522
523// newIntermediate takes the left & right keys of a intermediate node, and
524// computes its hash. Returns the hash of the node, which is the node key, and a
525// byte array that contains the value (which contains the left & right child
526// keys) to store in the DB.
527// [ 1 byte | 1 byte | N bytes | N bytes ]
528// [ type of node | length of left key | left key | right key ]
529func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
530 b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
531 b[0] = PrefixValueIntermediate
532 if len(l) > maxUint8 {
533 return nil, nil, fmt.Errorf("newIntermediate: len(l) > %v", maxUint8)
534 }
535 b[1] = byte(len(l))
536 copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
537 copy(b[PrefixValueLen+hashFunc.Len():], r)
538
539 key, err := hashFunc.Hash(l, r)
540 if err != nil {
541 return nil, nil, err
542 }
543
544 return key, b, nil
545}
546
547// ReadIntermediateChilds reads from a byte array the two childs keys
548func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
549 if len(b) < PrefixValueLen {
550 return []byte{}, []byte{}
551 }
552
553 lLen := b[1]
554 if len(b) < PrefixValueLen+int(lLen) {
555 return []byte{}, []byte{}
556 }
557 l := b[PrefixValueLen : PrefixValueLen+lLen]
558 r := b[PrefixValueLen+lLen:]
559 return l, r
560}
561
562func getPath(numLevels int, k []byte) []bool {
563 path := make([]bool, numLevels)
564 for n := 0; n < numLevels; n++ {
565 path[n] = k[n/8]&(1<<(n%8)) != 0
566 }
567 return path
568}
569
570// Update updates the value for a given existing key. If the given key does not
571// exist, returns an error.
572func (t *Tree) Update(k, v []byte) error {
573 wTx := t.db.WriteTx()
574 defer wTx.Discard()
575
576 if err := t.UpdateWithTx(wTx, k, v); err != nil {
577 return err
578 }
579 return wTx.Commit()
580}
581
582// UpdateWithTx does the same than the Update method, but allowing to pass the
583// db.WriteTx that is used. The db.WriteTx will not be committed inside this
584// method.
585func (t *Tree) UpdateWithTx(wTx db.WriteTx, k, v []byte) error {
586 t.Lock()
587 defer t.Unlock()
588
589 if !t.editable() {
590 return ErrSnapshotNotEditable
591 }
592
593 keyPath := make([]byte, int(math.Ceil(float64(t.maxLevels)/float64(8)))) //nolint:gomnd
594 copy(keyPath[:], k)
595 path := getPath(t.maxLevels, keyPath)
596
597 root, err := t.RootWithTx(wTx)
598 if err != nil {
599 return err
600 }
601
602 var siblings [][]byte
603 _, valueAtBottom, siblings, err := t.down(wTx, k, root, siblings, path, 0, true)
604 if err != nil {
605 return err
606 }
607 oldKey, _ := ReadLeafValue(valueAtBottom)
608 if !bytes.Equal(oldKey, k) {
609 return ErrKeyNotFound
610 }
611
612 leafKey, leafValue, err := t.newLeafValue(k, v)
613 if err != nil {
614 return err
615 }
616
617 if err := wTx.Set(leafKey, leafValue); err != nil {
618 return err
619 }
620
621 // go up to the root
622 if len(siblings) == 0 {
623 return t.setRoot(wTx, leafKey)
624 }
625 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, 0)
626 if err != nil {
627 return err
628 }
629
630 // store root to db
631 if err := t.setRoot(wTx, root); err != nil {
632 return err
633 }
634 return nil
635}
636
637// GenProof generates a MerkleTree proof for the given key. The leaf value is
638// returned, together with the packed siblings of the proof, and a boolean
639// parameter that indicates if the proof is of existence (true) or not (false).
640func (t *Tree) GenProof(k []byte) ([]byte, []byte, []byte, bool, error) {
641 rTx := t.db.ReadTx()
642 defer rTx.Discard()
643
644 return t.GenProofWithTx(rTx, k)
645}
646
647// GenProofWithTx does the same than the GenProof method, but allowing to pass
648// the db.ReadTx that is used.
649func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) {
650 keyPath := make([]byte, int(math.Ceil(float64(t.maxLevels)/float64(8)))) //nolint:gomnd
651 copy(keyPath[:], k)
652 path := getPath(t.maxLevels, keyPath)
653
654 root, err := t.RootWithTx(rTx)
655 if err != nil {
656 return nil, nil, nil, false, err
657 }
658
659 // go down to the leaf
660 var siblings [][]byte
661 _, value, siblings, err := t.down(rTx, k, root, siblings, path, 0, true)
662 if err != nil {
663 return nil, nil, nil, false, err
664 }
665
666 s, err := PackSiblings(t.hashFunction, siblings)
667 if err != nil {
668 return nil, nil, nil, false, err
669 }
670
671 leafK, leafV := ReadLeafValue(value)
672 if !bytes.Equal(k, leafK) {
673 // key not in tree, proof of non-existence
674 return leafK, leafV, s, false, nil
675 }
676
677 return leafK, leafV, s, true, nil
678}
679
680// PackSiblings packs the siblings into a byte array.
681// [ 2 byte | 2 byte | L bytes | S * N bytes ]
682// [ full length | bitmap length (L) | bitmap | N non-zero siblings ]
683// Where the bitmap indicates if the sibling is 0 or a value from the siblings
684// array. And S is the size of the output of the hash function used for the
685// Tree. The 2 2-byte that define the full length and bitmap length, are
686// encoded in little-endian.
687func PackSiblings(hashFunc HashFunction, siblings [][]byte) ([]byte, error) {
688 var b []byte
689 var bitmap []bool
690 emptySibling := make([]byte, hashFunc.Len())
691 for i := 0; i < len(siblings); i++ {
692 if bytes.Equal(siblings[i], emptySibling) {
693 bitmap = append(bitmap, false)
694 } else {
695 bitmap = append(bitmap, true)
696 b = append(b, siblings[i]...)
697 }
698 }
699
700 bitmapBytes := bitmapToBytes(bitmap)
701 l := len(bitmapBytes)
702 if l > maxUint16 {
703 return nil, fmt.Errorf("PackSiblings: bitmapBytes length > %v", maxUint16)
704 }
705
706 fullLen := 4 + l + len(b) //nolint:gomnd
707 if fullLen > maxUint16 {
708 return nil, fmt.Errorf("PackSiblings: fullLen > %v", maxUint16)
709 }
710 res := make([]byte, fullLen)
711 binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length
712 binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length
713 copy(res[4:4+l], bitmapBytes)
714 copy(res[4+l:], b)
715 return res, nil
716}
717
718// UnpackSiblings unpacks the siblings from a byte array.
719func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
720 fullLen := binary.LittleEndian.Uint16(b[0:2])
721 l := binary.LittleEndian.Uint16(b[2:4]) // bitmap bytes length
722 if len(b) != int(fullLen) {
723 return nil,
724 fmt.Errorf("expected len: %d, current len: %d",
725 fullLen, len(b))
726 }
727
728 bitmapBytes := b[4 : 4+l]
729 bitmap := bytesToBitmap(bitmapBytes)
730 siblingsBytes := b[4+l:]
731 iSibl := 0
732 emptySibl := make([]byte, hashFunc.Len())
733 var siblings [][]byte
734 for i := 0; i < len(bitmap); i++ {
735 if iSibl >= len(siblingsBytes) {
736 break
737 }
738 if bitmap[i] {
739 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
740 iSibl += hashFunc.Len()
741 } else {
742 siblings = append(siblings, emptySibl)
743 }
744 }
745 return siblings, nil
746}
747
748func bitmapToBytes(bitmap []bool) []byte {
749 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
750 b := make([]byte, bitmapBytesLen)
751 for i := 0; i < len(bitmap); i++ {
752 if bitmap[i] {
753 b[i/8] |= 1 << (i % 8)
754 }
755 }
756 return b
757}
758
759func bytesToBitmap(b []byte) []bool {
760 var bitmap []bool
761 for i := 0; i < len(b); i++ {
762 for j := 0; j < 8; j++ {
763 bitmap = append(bitmap, b[i]&(1<<j) > 0)
764 }
765 }
766 return bitmap
767}
768
769// Get returns the value in the Tree for a given key. If the key is not found,
770// will return the error ErrKeyNotFound, and in the leafK & leafV parameters
771// will be placed the data found in the tree in the leaf that was on the path
772// going to the input key.
773func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
774 rTx := t.db.ReadTx()
775 defer rTx.Discard()
776
777 return t.GetWithTx(rTx, k)
778}
779
780// GetWithTx does the same than the Get method, but allowing to pass the
781// db.ReadTx that is used. If the key is not found, will return the error
782// ErrKeyNotFound, and in the leafK & leafV parameters will be placed the data
783// found in the tree in the leaf that was on the path going to the input key.
784func (t *Tree) GetWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, error) {
785 keyPath := make([]byte, int(math.Ceil(float64(t.maxLevels)/float64(8)))) //nolint:gomnd
786 copy(keyPath[:], k)
787 path := getPath(t.maxLevels, keyPath)
788
789 root, err := t.RootWithTx(rTx)
790 if err != nil {
791 return nil, nil, err
792 }
793
794 // go down to the leaf
795 var siblings [][]byte
796 _, value, _, err := t.down(rTx, k, root, siblings, path, 0, true)
797 if err != nil {
798 return nil, nil, err
799 }
800 leafK, leafV := ReadLeafValue(value)
801 if !bytes.Equal(k, leafK) {
802 return leafK, leafV, ErrKeyNotFound
803 }
804
805 return leafK, leafV, nil
806}
807
808// CheckProof verifies the given proof. The proof verification depends on the
809// HashFunction passed as parameter.
810func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
811 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
812 if err != nil {
813 return false, err
814 }
815
816 keyPath := make([]byte, len(siblings))
817 copy(keyPath[:], k)
818
819 key, _, err := newLeafValue(hashFunc, k, v)
820 if err != nil {
821 return false, err
822 }
823
824 path := getPath(len(siblings), keyPath)
825 for i := len(siblings) - 1; i >= 0; i-- {
826 if path[i] {
827 key, _, err = newIntermediate(hashFunc, siblings[i], key)
828 if err != nil {
829 return false, err
830 }
831 } else {
832 key, _, err = newIntermediate(hashFunc, key, siblings[i])
833 if err != nil {
834 return false, err
835 }
836 }
837 }
838 if bytes.Equal(key[:], root) {
839 return true, nil
840 }
841 return false, nil
842}
843
844func (t *Tree) incNLeafs(wTx db.WriteTx, nLeafs int) error {
845 oldNLeafs, err := t.GetNLeafsWithTx(wTx)
846 if err != nil {
847 return err
848 }
849 newNLeafs := oldNLeafs + nLeafs
850 return t.setNLeafs(wTx, newNLeafs)
851}
852
853func (t *Tree) setNLeafs(wTx db.WriteTx, nLeafs int) error {
854 b := make([]byte, 8)
855 binary.LittleEndian.PutUint64(b, uint64(nLeafs))
856 if err := wTx.Set(dbKeyNLeafs, b); err != nil {
857 return err
858 }
859 return nil
860}
861
862// GetNLeafs returns the number of Leafs of the Tree.
863func (t *Tree) GetNLeafs() (int, error) {
864 rTx := t.db.ReadTx()
865 defer rTx.Discard()
866
867 return t.GetNLeafsWithTx(rTx)
868}
869
870// GetNLeafsWithTx does the same than the GetNLeafs method, but allowing to
871// pass the db.ReadTx that is used.
872func (t *Tree) GetNLeafsWithTx(rTx db.ReadTx) (int, error) {
873 b, err := rTx.Get(dbKeyNLeafs)
874 if err != nil {
875 return 0, err
876 }
877 nLeafs := binary.LittleEndian.Uint64(b)
878 return int(nLeafs), nil
879}
880
881// SetRoot sets the root to the given root
882func (t *Tree) SetRoot(root []byte) error {
883 wTx := t.db.WriteTx()
884 defer wTx.Discard()
885
886 if err := t.SetRootWithTx(wTx, root); err != nil {
887 return err
888 }
889 return wTx.Commit()
890}
891
892// SetRootWithTx sets the root to the given root using the given db.WriteTx
893func (t *Tree) SetRootWithTx(wTx db.WriteTx, root []byte) error {
894 if !t.editable() {
895 return ErrSnapshotNotEditable
896 }
897
898 if root == nil {
899 return fmt.Errorf("can not SetRoot with nil root")
900 }
901
902 // check that the root exists in the db
903 if !bytes.Equal(root, t.emptyHash) {
904 if _, err := wTx.Get(root); err == ErrKeyNotFound {
905 return fmt.Errorf("can not SetRoot with root %x, as it does not exist in the db", root)
906 } else if err != nil {
907 return err
908 }
909 }
910
911 return wTx.Set(dbKeyRoot, root)
912}
913
914// Snapshot returns a read-only copy of the Tree from the given root
915func (t *Tree) Snapshot(fromRoot []byte) (*Tree, error) {
916 // allow to define which root to use
917 if fromRoot == nil {
918 var err error
919 fromRoot, err = t.Root()
920 if err != nil {
921 return nil, err
922 }
923 }
924 rTx := t.db.ReadTx()
925 defer rTx.Discard()
926 // check that the root exists in the db
927 if !bytes.Equal(fromRoot, t.emptyHash) {
928 if _, err := rTx.Get(fromRoot); err == ErrKeyNotFound {
929 return nil,
930 fmt.Errorf("can not do a Snapshot with root %x, as it does not exist in the db",
931 fromRoot)
932 } else if err != nil {
933 return nil, err
934 }
935 }
936
937 return &Tree{
938 db: t.db,
939 maxLevels: t.maxLevels,
940 snapshotRoot: fromRoot,
941 emptyHash: t.emptyHash,
942 hashFunction: t.hashFunction,
943 dbg: t.dbg,
944 }, nil
945}
946
947// Iterate iterates through the full Tree, executing the given function on each
948// node of the Tree.
949func (t *Tree) Iterate(fromRoot []byte, f func([]byte, []byte)) error {
950 rTx := t.db.ReadTx()
951 defer rTx.Discard()
952
953 return t.IterateWithTx(rTx, fromRoot, f)
954}
955
956// IterateWithTx does the same than the Iterate method, but allowing to pass
957// the db.ReadTx that is used.
958func (t *Tree) IterateWithTx(rTx db.ReadTx, fromRoot []byte, f func([]byte, []byte)) error {
959 // allow to define which root to use
960 if fromRoot == nil {
961 var err error
962 fromRoot, err = t.RootWithTx(rTx)
963 if err != nil {
964 return err
965 }
966 }
967 return t.iter(rTx, fromRoot, f)
968}
969
970// IterateWithStop does the same than Iterate, but with int for the current
971// level, and a boolean parameter used by the passed function, is to indicate to
972// stop iterating on the branch when the method returns 'true'.
973func (t *Tree) IterateWithStop(fromRoot []byte, f func(int, []byte, []byte) bool) error {
974 rTx := t.db.ReadTx()
975 defer rTx.Discard()
976
977 // allow to define which root to use
978 if fromRoot == nil {
979 var err error
980 fromRoot, err = t.RootWithTx(rTx)
981 if err != nil {
982 return err
983 }
984 }
985 return t.iterWithStop(rTx, fromRoot, 0, f)
986}
987
988// IterateWithStopWithTx does the same than the IterateWithStop method, but
989// allowing to pass the db.ReadTx that is used.
990func (t *Tree) IterateWithStopWithTx(rTx db.ReadTx, fromRoot []byte,
991 f func(int, []byte, []byte) bool) error {
992 // allow to define which root to use
993 if fromRoot == nil {
994 var err error
995 fromRoot, err = t.RootWithTx(rTx)
996 if err != nil {
997 return err
998 }
999 }
1000 return t.iterWithStop(rTx, fromRoot, 0, f)
1001}
1002
1003func (t *Tree) iterWithStop(rTx db.ReadTx, k []byte, currLevel int,
1004 f func(int, []byte, []byte) bool) error {
1005 var v []byte
1006 var err error
1007 if bytes.Equal(k, t.emptyHash) {
1008 v = t.emptyHash
1009 } else {
1010 v, err = rTx.Get(k)
1011 if err != nil {
1012 return err
1013 }
1014 }
1015 currLevel++
1016
1017 switch v[0] {
1018 case PrefixValueEmpty:
1019 f(currLevel, k, v)
1020 case PrefixValueLeaf:
1021 f(currLevel, k, v)
1022 case PrefixValueIntermediate:
1023 stop := f(currLevel, k, v)
1024 if stop {
1025 return nil
1026 }
1027 l, r := ReadIntermediateChilds(v)
1028 if err = t.iterWithStop(rTx, l, currLevel, f); err != nil {
1029 return err
1030 }
1031 if err = t.iterWithStop(rTx, r, currLevel, f); err != nil {
1032 return err
1033 }
1034 default:
1035 return ErrInvalidValuePrefix
1036 }
1037 return nil
1038}
1039
1040func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
1041 f2 := func(currLvl int, k, v []byte) bool {
1042 f(k, v)
1043 return false
1044 }
1045 return t.iterWithStop(rTx, k, 0, f2)
1046}
1047
1048// Dump exports all the Tree leafs in a byte array of length:
1049// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
1050// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
1051// [ len(k) | len(v) | key | value ]
1052// Where S is the size of the output of the hash function used for the Tree.
1053func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
1054 // allow to define which root to use
1055 if fromRoot == nil {
1056 var err error
1057 fromRoot, err = t.Root()
1058 if err != nil {
1059 return nil, err
1060 }
1061 }
1062
1063 // WARNING current encoding only supports key & values of 255 bytes each
1064 // (due using only 1 byte for the length headers).
1065 var b []byte
1066 var callbackErr error
1067 err := t.IterateWithStop(fromRoot, func(_ int, k, v []byte) bool {
1068 if v[0] != PrefixValueLeaf {
1069 return false
1070 }
1071 leafK, leafV := ReadLeafValue(v)
1072 kv := make([]byte, 2+len(leafK)+len(leafV))
1073 if len(leafK) > maxUint8 {
1074 callbackErr = fmt.Errorf("len(leafK) > %v", maxUint8)
1075 return true
1076 }
1077 kv[0] = byte(len(leafK))
1078 if len(leafV) > maxUint8 {
1079 callbackErr = fmt.Errorf("len(leafV) > %v", maxUint8)
1080 return true
1081 }
1082 kv[1] = byte(len(leafV))
1083 copy(kv[2:2+len(leafK)], leafK)
1084 copy(kv[2+len(leafK):], leafV)
1085 b = append(b, kv...)
1086 return false
1087 })
1088 if callbackErr != nil {
1089 return nil, callbackErr
1090 }
1091 return b, err
1092}
1093
1094// ImportDump imports the leafs (that have been exported with the Dump method)
1095// in the Tree.
1096func (t *Tree) ImportDump(b []byte) error {
1097 if !t.editable() {
1098 return ErrSnapshotNotEditable
1099 }
1100 root, err := t.Root()
1101 if err != nil {
1102 return err
1103 }
1104 if !bytes.Equal(root, t.emptyHash) {
1105 return ErrTreeNotEmpty
1106 }
1107
1108 r := bytes.NewReader(b)
1109 var keys, values [][]byte
1110 for {
1111 l := make([]byte, 2)
1112 _, err = io.ReadFull(r, l)
1113 if err == io.EOF {
1114 break
1115 } else if err != nil {
1116 return err
1117 }
1118 k := make([]byte, l[0])
1119 _, err = io.ReadFull(r, k)
1120 if err != nil {
1121 return err
1122 }
1123 v := make([]byte, l[1])
1124 _, err = io.ReadFull(r, v)
1125 if err != nil {
1126 return err
1127 }
1128 keys = append(keys, k)
1129 values = append(values, v)
1130 }
1131 if _, err = t.AddBatch(keys, values); err != nil {
1132 return err
1133 }
1134 return nil
1135}
1136
1137// Graphviz iterates across the full tree to generate a string Graphviz
1138// representation of the tree and writes it to w
1139func (t *Tree) Graphviz(w io.Writer, fromRoot []byte) error {
1140 return t.GraphvizFirstNLevels(w, fromRoot, t.maxLevels)
1141}
1142
1143// GraphvizFirstNLevels iterates across the first NLevels of the tree to
1144// generate a string Graphviz representation of the first NLevels of the tree
1145// and writes it to w
1146func (t *Tree) GraphvizFirstNLevels(w io.Writer, fromRoot []byte, untilLvl int) error {
1147 fmt.Fprintf(w, `digraph hierarchy {
1148node [fontname=Monospace,fontsize=10,shape=box]
1149`)
1150
1151 rTx := t.db.ReadTx()
1152 defer rTx.Discard()
1153
1154 if fromRoot == nil {
1155 var err error
1156 fromRoot, err = t.RootWithTx(rTx)
1157 if err != nil {
1158 return err
1159 }
1160 }
1161
1162 nEmpties := 0
1163 err := t.iterWithStop(rTx, fromRoot, 0, func(currLvl int, k, v []byte) bool {
1164 if currLvl == untilLvl {
1165 return true // to stop the iter from going down
1166 }
1167 switch v[0] {
1168 case PrefixValueEmpty:
1169 case PrefixValueLeaf:
1170 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
1171 // key & value from the leaf
1172 kB, vB := ReadLeafValue(v)
1173 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
1174 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
1175 hex.EncodeToString(vB[:nChars]))
1176 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
1177 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
1178 case PrefixValueIntermediate:
1179 l, r := ReadIntermediateChilds(v)
1180 lStr := hex.EncodeToString(l[:nChars])
1181 rStr := hex.EncodeToString(r[:nChars])
1182 eStr := ""
1183 if bytes.Equal(l, t.emptyHash) {
1184 lStr = fmt.Sprintf("empty%v", nEmpties)
1185 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1186 lStr)
1187 nEmpties++
1188 }
1189 if bytes.Equal(r, t.emptyHash) {
1190 rStr = fmt.Sprintf("empty%v", nEmpties)
1191 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1192 rStr)
1193 nEmpties++
1194 }
1195 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
1196 lStr, rStr)
1197 fmt.Fprint(w, eStr)
1198 default:
1199 }
1200 return false
1201 })
1202 fmt.Fprintf(w, "}\n")
1203 return err
1204}
1205
1206// PrintGraphviz prints the output of Tree.Graphviz
1207func (t *Tree) PrintGraphviz(fromRoot []byte) error {
1208 if fromRoot == nil {
1209 var err error
1210 fromRoot, err = t.Root()
1211 if err != nil {
1212 return err
1213 }
1214 }
1215 return t.PrintGraphvizFirstNLevels(fromRoot, t.maxLevels)
1216}
1217
1218// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
1219func (t *Tree) PrintGraphvizFirstNLevels(fromRoot []byte, untilLvl int) error {
1220 if fromRoot == nil {
1221 var err error
1222 fromRoot, err = t.Root()
1223 if err != nil {
1224 return err
1225 }
1226 }
1227 w := bytes.NewBufferString("")
1228 fmt.Fprintf(w,
1229 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+":\n")
1230 err := t.GraphvizFirstNLevels(w, fromRoot, untilLvl)
1231 if err != nil {
1232 fmt.Println(w)
1233 return err
1234 }
1235 fmt.Fprintf(w,
1236 "End of Graphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+"\n--------\n")
1237
1238 fmt.Println(w)
1239 return nil
1240}
1241
1242// TODO circom proofs
1243// TODO data structure for proofs (including root, key, value, siblings,
1244// hashFunction) + method to verify that data structure