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