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