arbo/tree.go
tree.goRaw
1/*
2Package arbo implements a Merkle Tree compatible with the circomlib
3implementation of the MerkleTree (when using the Poseidon hash function),
4following the specification from
5https://docs.iden3.io/publications/pdfs/Merkle-Tree.pdf and
6https://eprint.iacr.org/2018/955.
7
8Also allows to define which hash function to use. So for example, when working
9with zkSnarks the Poseidon hash function can be used, but when not, it can be
10used the Blake3 hash function, which improves the computation time.
11*/
12package arbo
13
14import (
15 "bytes"
16 "encoding/hex"
17 "fmt"
18 "io"
19 "math"
20 "sync"
21 "sync/atomic"
22 "time"
23
24 "github.com/iden3/go-merkletree/db"
25)
26
27const (
28 // PrefixValueLen defines the bytes-prefix length used for the Value
29 // bytes representation stored in the db
30 PrefixValueLen = 2
31
32 // PrefixValueEmpty is used for the first byte of a Value to indicate
33 // that is an Empty value
34 PrefixValueEmpty = 0
35 // PrefixValueLeaf is used for the first byte of a Value to indicate
36 // that is a Leaf value
37 PrefixValueLeaf = 1
38 // PrefixValueIntermediate is used for the first byte of a Value to
39 // indicate that is a Intermediate value
40 PrefixValueIntermediate = 2
41)
42
43var (
44 dbKeyRoot = []byte("root")
45 emptyValue = []byte{0}
46)
47
48// Tree defines the struct that implements the MerkleTree functionalities
49type Tree struct {
50 sync.RWMutex
51 tx db.Tx
52 db db.Storage
53 lastAccess int64 // in unix time
54 maxLevels int
55 root []byte
56
57 hashFunction HashFunction
58}
59
60// NewTree returns a new Tree, if there is a Tree still in the given storage, it
61// will load it.
62func NewTree(storage db.Storage, maxLevels int, hash HashFunction) (*Tree, error) {
63 t := Tree{db: storage, maxLevels: maxLevels, hashFunction: hash}
64 t.updateAccessTime()
65
66 root, err := t.dbGet(dbKeyRoot)
67 if err == db.ErrNotFound {
68 // store new root 0
69 tx, err := t.db.NewTx()
70 if err != nil {
71 return nil, err
72 }
73 t.root = make([]byte, t.hashFunction.Len()) // empty
74 if err = tx.Put(dbKeyRoot, t.root); err != nil {
75 return nil, err
76 }
77 if err = tx.Commit(); err != nil {
78 return nil, err
79 }
80 return &t, err
81 } else if err != nil {
82 return nil, err
83 }
84 t.root = root
85 return &t, nil
86}
87
88func (t *Tree) updateAccessTime() {
89 atomic.StoreInt64(&t.lastAccess, time.Now().Unix())
90}
91
92// LastAccess returns the last access timestamp in Unixtime
93func (t *Tree) LastAccess() int64 {
94 return atomic.LoadInt64(&t.lastAccess)
95}
96
97// Root returns the root of the Tree
98func (t *Tree) Root() []byte {
99 return t.root
100}
101
102// AddBatch adds a batch of key-values to the Tree. This method will be
103// optimized to do some internal parallelization. Returns an array containing
104// the indexes of the keys failed to add.
105func (t *Tree) AddBatch(keys, values [][]byte) ([]int, error) {
106 t.updateAccessTime()
107 if len(keys) != len(values) {
108 return nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)",
109 len(keys), len(values))
110 }
111
112 t.Lock()
113 defer t.Unlock()
114
115 var err error
116 t.tx, err = t.db.NewTx()
117 if err != nil {
118 return nil, err
119 }
120
121 var indexes []int
122 for i := 0; i < len(keys); i++ {
123 err = t.add(keys[i], values[i])
124 if err != nil {
125 indexes = append(indexes, i)
126 }
127 }
128 // store root to db
129 if err := t.tx.Put(dbKeyRoot, t.root); err != nil {
130 return indexes, err
131 }
132
133 if err := t.tx.Commit(); err != nil {
134 return nil, err
135 }
136 return indexes, nil
137}
138
139// Add inserts the key-value into the Tree. If the inputs come from a *big.Int,
140// is expected that are represented by a Little-Endian byte array (for circom
141// compatibility).
142func (t *Tree) Add(k, v []byte) error {
143 t.updateAccessTime()
144
145 t.Lock()
146 defer t.Unlock()
147
148 var err error
149 t.tx, err = t.db.NewTx()
150 if err != nil {
151 return err
152 }
153
154 err = t.add(k, v)
155 if err != nil {
156 return err
157 }
158 // store root to db
159 if err := t.tx.Put(dbKeyRoot, t.root); err != nil {
160 return err
161 }
162 return t.tx.Commit()
163}
164
165func (t *Tree) add(k, v []byte) error {
166 // TODO check validity of key & value (for the Tree.HashFunction type)
167
168 keyPath := make([]byte, t.hashFunction.Len())
169 copy(keyPath[:], k)
170
171 path := getPath(t.maxLevels, keyPath)
172 // go down to the leaf
173 var siblings [][]byte
174 _, _, siblings, err := t.down(k, t.root, siblings, path, 0, false)
175 if err != nil {
176 return err
177 }
178
179 leafKey, leafValue, err := newLeafValue(t.hashFunction, k, v)
180 if err != nil {
181 return err
182 }
183
184 if err := t.tx.Put(leafKey, leafValue); err != nil {
185 return err
186 }
187
188 // go up to the root
189 if len(siblings) == 0 {
190 t.root = leafKey
191 return nil
192 }
193 root, err := t.up(leafKey, siblings, path, len(siblings)-1)
194 if err != nil {
195 return err
196 }
197
198 t.root = root
199 return nil
200}
201
202// down goes down to the leaf recursively
203func (t *Tree) down(newKey, currKey []byte, siblings [][]byte,
204 path []bool, l int, getLeaf bool) (
205 []byte, []byte, [][]byte, error) {
206 if l > t.maxLevels-1 {
207 return nil, nil, nil, fmt.Errorf("max level")
208 }
209 var err error
210 var currValue []byte
211 emptyKey := make([]byte, t.hashFunction.Len())
212 if bytes.Equal(currKey, emptyKey) {
213 // empty value
214 return currKey, emptyValue, siblings, nil
215 }
216 currValue, err = t.dbGet(currKey)
217 if err != nil {
218 return nil, nil, nil, err
219 }
220
221 switch currValue[0] {
222 case PrefixValueEmpty: // empty
223 // TODO WIP WARNING should not be reached, as the 'if' above should avoid
224 // reaching this point
225 // return currKey, empty, siblings, nil
226 panic("should not be reached, as the 'if' above should avoid reaching this point") // TMP
227 case PrefixValueLeaf: // leaf
228 if bytes.Equal(newKey, currKey) {
229 return nil, nil, nil, fmt.Errorf("key already exists")
230 }
231
232 if !bytes.Equal(currValue, emptyValue) {
233 if getLeaf {
234 return currKey, currValue, siblings, nil
235 }
236 oldLeafKey, _ := readLeafValue(currValue)
237 oldLeafKeyFull := make([]byte, t.hashFunction.Len())
238 copy(oldLeafKeyFull[:], oldLeafKey)
239
240 // if currKey is already used, go down until paths diverge
241 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
242 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, l)
243 if err != nil {
244 return nil, nil, nil, err
245 }
246 }
247 return currKey, currValue, siblings, nil
248 case PrefixValueIntermediate: // intermediate
249 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
250 return nil, nil, nil,
251 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
252 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
253 }
254 // collect siblings while going down
255 if path[l] {
256 // right
257 lChild, rChild := readIntermediateChilds(currValue)
258 siblings = append(siblings, lChild)
259 return t.down(newKey, rChild, siblings, path, l+1, getLeaf)
260 }
261 // left
262 lChild, rChild := readIntermediateChilds(currValue)
263 siblings = append(siblings, rChild)
264 return t.down(newKey, lChild, siblings, path, l+1, getLeaf)
265 default:
266 return nil, nil, nil, fmt.Errorf("invalid value")
267 }
268}
269
270// downVirtually is used when in a leaf already exists, and a new leaf which
271// shares the path until the existing leaf is being added
272func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
273 newPath []bool, l int) ([][]byte, error) {
274 var err error
275 if l > t.maxLevels-1 {
276 return nil, fmt.Errorf("max virtual level %d", l)
277 }
278
279 if oldPath[l] == newPath[l] {
280 emptyKey := make([]byte, t.hashFunction.Len())
281 siblings = append(siblings, emptyKey)
282
283 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, l+1)
284 if err != nil {
285 return nil, err
286 }
287 return siblings, nil
288 }
289 // reached the divergence
290 siblings = append(siblings, oldKey)
291
292 return siblings, nil
293}
294
295// up goes up recursively updating the intermediate nodes
296func (t *Tree) up(key []byte, siblings [][]byte, path []bool, l int) ([]byte, error) {
297 var k, v []byte
298 var err error
299 if path[l] {
300 k, v, err = newIntermediate(t.hashFunction, siblings[l], key)
301 if err != nil {
302 return nil, err
303 }
304 } else {
305 k, v, err = newIntermediate(t.hashFunction, key, siblings[l])
306 if err != nil {
307 return nil, err
308 }
309 }
310 // store k-v to db
311 if err = t.tx.Put(k, v); err != nil {
312 return nil, err
313 }
314
315 if l == 0 {
316 // reached the root
317 return k, nil
318 }
319
320 return t.up(k, siblings, path, l-1)
321}
322
323func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
324 leafKey, err := hashFunc.Hash(k, v, []byte{1})
325 if err != nil {
326 return nil, nil, err
327 }
328 var leafValue []byte
329 leafValue = append(leafValue, byte(1))
330 leafValue = append(leafValue, byte(len(k)))
331 leafValue = append(leafValue, k...)
332 leafValue = append(leafValue, v...)
333 return leafKey, leafValue, nil
334}
335
336func readLeafValue(b []byte) ([]byte, []byte) {
337 if len(b) < PrefixValueLen {
338 return []byte{}, []byte{}
339 }
340
341 kLen := b[1]
342 if len(b) < PrefixValueLen+int(kLen) {
343 return []byte{}, []byte{}
344 }
345 k := b[PrefixValueLen : PrefixValueLen+kLen]
346 v := b[PrefixValueLen+kLen:]
347 return k, v
348}
349
350func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
351 b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
352 b[0] = 2
353 b[1] = byte(len(l))
354 copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
355 copy(b[PrefixValueLen+hashFunc.Len():], r)
356
357 key, err := hashFunc.Hash(l, r)
358 if err != nil {
359 return nil, nil, err
360 }
361
362 return key, b, nil
363}
364
365func readIntermediateChilds(b []byte) ([]byte, []byte) {
366 if len(b) < PrefixValueLen {
367 return []byte{}, []byte{}
368 }
369
370 lLen := b[1]
371 if len(b) < PrefixValueLen+int(lLen) {
372 return []byte{}, []byte{}
373 }
374 l := b[PrefixValueLen : PrefixValueLen+lLen]
375 r := b[PrefixValueLen+lLen:]
376 return l, r
377}
378
379func getPath(numLevels int, k []byte) []bool {
380 path := make([]bool, numLevels)
381 for n := 0; n < numLevels; n++ {
382 path[n] = k[n/8]&(1<<(n%8)) != 0
383 }
384 return path
385}
386
387// Update updates the value for a given existing key. If the given key does not
388// exist, returns an error.
389func (t *Tree) Update(k, v []byte) error {
390 t.updateAccessTime()
391
392 t.Lock()
393 defer t.Unlock()
394
395 var err error
396 t.tx, err = t.db.NewTx()
397 if err != nil {
398 return err
399 }
400
401 keyPath := make([]byte, t.hashFunction.Len())
402 copy(keyPath[:], k)
403 path := getPath(t.maxLevels, keyPath)
404
405 var siblings [][]byte
406 _, valueAtBottom, siblings, err := t.down(k, t.root, siblings, path, 0, true)
407 if err != nil {
408 return err
409 }
410 oldKey, _ := readLeafValue(valueAtBottom)
411 if !bytes.Equal(oldKey, k) {
412 return fmt.Errorf("key %s does not exist", hex.EncodeToString(k))
413 }
414
415 leafKey, leafValue, err := newLeafValue(t.hashFunction, k, v)
416 if err != nil {
417 return err
418 }
419
420 if err := t.tx.Put(leafKey, leafValue); err != nil {
421 return err
422 }
423
424 // go up to the root
425 if len(siblings) == 0 {
426 t.root = leafKey
427 return t.tx.Commit()
428 }
429 root, err := t.up(leafKey, siblings, path, len(siblings)-1)
430 if err != nil {
431 return err
432 }
433
434 t.root = root
435 // store root to db
436 if err := t.tx.Put(dbKeyRoot, t.root); err != nil {
437 return err
438 }
439 return t.tx.Commit()
440}
441
442// GenProof generates a MerkleTree proof for the given key. If the key exists in
443// the Tree, the proof will be of existence, if the key does not exist in the
444// tree, the proof will be of non-existence.
445func (t *Tree) GenProof(k []byte) ([]byte, error) {
446 t.updateAccessTime()
447 keyPath := make([]byte, t.hashFunction.Len())
448 copy(keyPath[:], k)
449
450 path := getPath(t.maxLevels, keyPath)
451 // go down to the leaf
452 var siblings [][]byte
453 _, value, siblings, err := t.down(k, t.root, siblings, path, 0, true)
454 if err != nil {
455 return nil, err
456 }
457
458 leafK, leafV := readLeafValue(value)
459 if !bytes.Equal(k, leafK) {
460 fmt.Println("key not in Tree")
461 fmt.Println(leafK)
462 fmt.Println(leafV)
463 // TODO proof of non-existence
464 panic(fmt.Errorf("unimplemented"))
465 }
466
467 s := PackSiblings(t.hashFunction, siblings)
468 return s, nil
469}
470
471// PackSiblings packs the siblings into a byte array.
472// [ 1 byte | L bytes | S * N bytes ]
473// [ bitmap length (L) | bitmap | N non-zero siblings ]
474// Where the bitmap indicates if the sibling is 0 or a value from the siblings
475// array. And S is the size of the output of the hash function used for the
476// Tree.
477func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte {
478 var b []byte
479 var bitmap []bool
480 emptySibling := make([]byte, hashFunc.Len())
481 for i := 0; i < len(siblings); i++ {
482 if bytes.Equal(siblings[i], emptySibling) {
483 bitmap = append(bitmap, false)
484 } else {
485 bitmap = append(bitmap, true)
486 b = append(b, siblings[i]...)
487 }
488 }
489
490 bitmapBytes := bitmapToBytes(bitmap)
491 l := len(bitmapBytes)
492
493 res := make([]byte, l+1+len(b))
494 res[0] = byte(l) // set the bitmapBytes length
495 copy(res[1:1+l], bitmapBytes)
496 copy(res[1+l:], b)
497 return res
498}
499
500// UnpackSiblings unpacks the siblings from a byte array.
501func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
502 l := b[0]
503 bitmapBytes := b[1 : 1+l]
504 bitmap := bytesToBitmap(bitmapBytes)
505 siblingsBytes := b[1+l:]
506 iSibl := 0
507 emptySibl := make([]byte, hashFunc.Len())
508 var siblings [][]byte
509 for i := 0; i < len(bitmap); i++ {
510 if iSibl >= len(siblingsBytes) {
511 break
512 }
513 if bitmap[i] {
514 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
515 iSibl += hashFunc.Len()
516 } else {
517 siblings = append(siblings, emptySibl)
518 }
519 }
520 return siblings, nil
521}
522
523func bitmapToBytes(bitmap []bool) []byte {
524 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
525 b := make([]byte, bitmapBytesLen)
526 for i := 0; i < len(bitmap); i++ {
527 if bitmap[i] {
528 b[i/8] |= 1 << (i % 8)
529 }
530 }
531 return b
532}
533
534func bytesToBitmap(b []byte) []bool {
535 var bitmap []bool
536 for i := 0; i < len(b); i++ {
537 for j := 0; j < 8; j++ {
538 bitmap = append(bitmap, b[i]&(1<<j) > 0)
539 }
540 }
541 return bitmap
542}
543
544// Get returns the value for a given key
545func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
546 keyPath := make([]byte, t.hashFunction.Len())
547 copy(keyPath[:], k)
548
549 path := getPath(t.maxLevels, keyPath)
550 // go down to the leaf
551 var siblings [][]byte
552 _, value, _, err := t.down(k, t.root, siblings, path, 0, true)
553 if err != nil {
554 return nil, nil, err
555 }
556 leafK, leafV := readLeafValue(value)
557 if !bytes.Equal(k, leafK) {
558 panic(fmt.Errorf("%s != %s", BytesToBigInt(k), BytesToBigInt(leafK)))
559 }
560
561 return leafK, leafV, nil
562}
563
564// CheckProof verifies the given proof. The proof verification depends on the
565// HashFunction passed as parameter.
566func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
567 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
568 if err != nil {
569 return false, err
570 }
571
572 keyPath := make([]byte, hashFunc.Len())
573 copy(keyPath[:], k)
574
575 key, _, err := newLeafValue(hashFunc, k, v)
576 if err != nil {
577 return false, err
578 }
579
580 path := getPath(len(siblings), keyPath)
581 for i := len(siblings) - 1; i >= 0; i-- {
582 if path[i] {
583 key, _, err = newIntermediate(hashFunc, siblings[i], key)
584 if err != nil {
585 return false, err
586 }
587 } else {
588 key, _, err = newIntermediate(hashFunc, key, siblings[i])
589 if err != nil {
590 return false, err
591 }
592 }
593 }
594 if bytes.Equal(key[:], root) {
595 return true, nil
596 }
597 return false, nil
598}
599
600func (t *Tree) dbGet(k []byte) ([]byte, error) {
601 v, err := t.db.Get(k)
602 if err == nil {
603 return v, nil
604 }
605 if t.tx != nil {
606 return t.tx.Get(k)
607 }
608 return nil, db.ErrNotFound
609}
610
611// Iterate iterates through the full Tree, executing the given function on each
612// node of the Tree.
613func (t *Tree) Iterate(f func([]byte, []byte)) error {
614 t.updateAccessTime()
615 return t.iter(t.root, f)
616}
617
618func (t *Tree) iter(k []byte, f func([]byte, []byte)) error {
619 v, err := t.dbGet(k)
620 if err != nil {
621 return err
622 }
623
624 switch v[0] {
625 case PrefixValueEmpty:
626 f(k, v)
627 case PrefixValueLeaf:
628 f(k, v)
629 case PrefixValueIntermediate:
630 f(k, v)
631 l, r := readIntermediateChilds(v)
632 if err = t.iter(l, f); err != nil {
633 return err
634 }
635 if err = t.iter(r, f); err != nil {
636 return err
637 }
638 default:
639 return fmt.Errorf("invalid value")
640 }
641 return nil
642}
643
644// Dump exports all the Tree leafs in a byte array of length:
645// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
646// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
647// [ len(k) | len(v) | key | value ]
648// Where S is the size of the output of the hash function used for the Tree.
649func (t *Tree) Dump() ([]byte, error) {
650 t.updateAccessTime()
651 // WARNING current encoding only supports key & values of 255 bytes each
652 // (due using only 1 byte for the length headers).
653 var b []byte
654 err := t.Iterate(func(k, v []byte) {
655 if v[0] != PrefixValueLeaf {
656 return
657 }
658 leafK, leafV := readLeafValue(v)
659 kv := make([]byte, 2+len(leafK)+len(leafV))
660 kv[0] = byte(len(leafK))
661 kv[1] = byte(len(leafV))
662 copy(kv[2:2+len(leafK)], leafK)
663 copy(kv[2+len(leafK):], leafV)
664 b = append(b, kv...)
665 })
666 return b, err
667}
668
669// ImportDump imports the leafs (that have been exported with the ExportLeafs
670// method) in the Tree.
671func (t *Tree) ImportDump(b []byte) error {
672 t.updateAccessTime()
673 r := bytes.NewReader(b)
674 for {
675 l := make([]byte, 2)
676 _, err := io.ReadFull(r, l)
677 if err == io.EOF {
678 break
679 } else if err != nil {
680 return err
681 }
682 k := make([]byte, l[0])
683 _, err = io.ReadFull(r, k)
684 if err != nil {
685 return err
686 }
687 v := make([]byte, l[1])
688 _, err = io.ReadFull(r, v)
689 if err != nil {
690 return err
691 }
692 err = t.Add(k, v)
693 if err != nil {
694 return err
695 }
696 }
697 return nil
698}