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