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