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.db.Get(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 err = tx.Put(dbKeyRoot, t.root)
70 if err != nil {
71 return nil, err
72 }
73 err = tx.Commit()
74 if 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 is optimized to
100// do some internal parallelization. Returns an array containing the indexes of
101// the keys failed to add.
102func (t *Tree) AddBatch(keys, values [][]byte) ([]int, error) {
103 return nil, fmt.Errorf("unimplemented")
104}
105
106// Add inserts the key-value into the Tree.
107// If the inputs come from a *big.Int, is expected that are represented by a
108// Little-Endian byte array (for circom compatibility).
109func (t *Tree) Add(k, v []byte) error {
110 // TODO check validity of key & value (for the Tree.HashFunction type)
111
112 keyPath := make([]byte, t.hashFunction.Len())
113 copy(keyPath[:], k)
114
115 path := getPath(t.maxLevels, keyPath)
116 // go down to the leaf
117 var siblings [][]byte
118 _, _, siblings, err := t.down(k, t.root, siblings, path, 0, false)
119 if err != nil {
120 return err
121 }
122
123 leafKey, leafValue, err := newLeafValue(t.hashFunction, k, v)
124 if err != nil {
125 return err
126 }
127
128 tx, err := t.db.NewTx()
129 if err != nil {
130 return err
131 }
132 if err := tx.Put(leafKey, leafValue); err != nil {
133 return err
134 }
135
136 // go up to the root
137 if len(siblings) == 0 {
138 t.root = leafKey
139 return tx.Commit()
140 }
141 root, err := t.up(tx, leafKey, siblings, path, len(siblings)-1)
142 if err != nil {
143 return err
144 }
145
146 t.root = root
147 // store root to db
148 return tx.Commit()
149}
150
151// down goes down to the leaf recursively
152func (t *Tree) down(newKey, currKey []byte, siblings [][]byte, path []bool, l int, getLeaf bool) (
153 []byte, []byte, [][]byte, error) {
154 if l > t.maxLevels-1 {
155 return nil, nil, nil, fmt.Errorf("max level")
156 }
157 var err error
158 var currValue []byte
159 emptyKey := make([]byte, t.hashFunction.Len())
160 if bytes.Equal(currKey, emptyKey) {
161 // empty value
162 return currKey, emptyValue, siblings, nil
163 }
164 currValue, err = t.db.Get(currKey)
165 if err != nil {
166 return nil, nil, nil, err
167 }
168
169 switch currValue[0] {
170 case PrefixValueEmpty: // empty
171 // TODO WIP WARNING should not be reached, as the 'if' above should avoid
172 // reaching this point
173 // return currKey, empty, siblings, nil
174 panic("should not be reached, as the 'if' above should avoid reaching this point") // TMP
175 case PrefixValueLeaf: // leaf
176 if bytes.Equal(newKey, currKey) {
177 return nil, nil, nil, fmt.Errorf("key already exists")
178 }
179
180 if !bytes.Equal(currValue, emptyValue) {
181 if getLeaf {
182 return currKey, currValue, siblings, nil
183 }
184 oldLeafKey, _ := readLeafValue(currValue)
185 oldLeafKeyFull := make([]byte, t.hashFunction.Len())
186 copy(oldLeafKeyFull[:], oldLeafKey)
187
188 // if currKey is already used, go down until paths diverge
189 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
190 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, l)
191 if err != nil {
192 return nil, nil, nil, err
193 }
194 }
195 return currKey, currValue, siblings, nil
196 case PrefixValueIntermediate: // intermediate
197 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
198 return nil, nil, nil,
199 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
200 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
201 }
202 // collect siblings while going down
203 if path[l] {
204 // right
205 lChild, rChild := readIntermediateChilds(currValue)
206 siblings = append(siblings, lChild)
207 return t.down(newKey, rChild, siblings, path, l+1, getLeaf)
208 }
209 // left
210 lChild, rChild := readIntermediateChilds(currValue)
211 siblings = append(siblings, rChild)
212 return t.down(newKey, lChild, siblings, path, l+1, getLeaf)
213 default:
214 return nil, nil, nil, fmt.Errorf("invalid value")
215 }
216}
217
218// downVirtually is used when in a leaf already exists, and a new leaf which
219// shares the path until the existing leaf is being added
220func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
221 newPath []bool, l int) ([][]byte, error) {
222 var err error
223 if l > t.maxLevels-1 {
224 return nil, fmt.Errorf("max virtual level %d", l)
225 }
226
227 if oldPath[l] == newPath[l] {
228 emptyKey := make([]byte, t.hashFunction.Len()) // empty
229 siblings = append(siblings, emptyKey)
230
231 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, l+1)
232 if err != nil {
233 return nil, err
234 }
235 return siblings, nil
236 }
237 // reached the divergence
238 siblings = append(siblings, oldKey)
239
240 return siblings, nil
241}
242
243// up goes up recursively updating the intermediate nodes
244func (t *Tree) up(tx db.Tx, key []byte, siblings [][]byte, path []bool, l int) ([]byte, error) {
245 var k, v []byte
246 var err error
247 if path[l] {
248 k, v, err = newIntermediate(t.hashFunction, siblings[l], key)
249 if err != nil {
250 return nil, err
251 }
252 } else {
253 k, v, err = newIntermediate(t.hashFunction, key, siblings[l])
254 if err != nil {
255 return nil, err
256 }
257 }
258 // store k-v to db
259 err = tx.Put(k, v)
260 if err != nil {
261 return nil, err
262 }
263
264 if l == 0 {
265 // reached the root
266 return k, nil
267 }
268
269 return t.up(tx, k, siblings, path, l-1)
270}
271
272func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
273 leafKey, err := hashFunc.Hash(k, v, []byte{1})
274 if err != nil {
275 return nil, nil, err
276 }
277 var leafValue []byte
278 leafValue = append(leafValue, byte(1))
279 leafValue = append(leafValue, byte(len(k)))
280 leafValue = append(leafValue, k...)
281 leafValue = append(leafValue, v...)
282 return leafKey, leafValue, nil
283}
284
285func readLeafValue(b []byte) ([]byte, []byte) {
286 if len(b) < PrefixValueLen {
287 return []byte{}, []byte{}
288 }
289
290 kLen := b[1]
291 if len(b) < PrefixValueLen+int(kLen) {
292 return []byte{}, []byte{}
293 }
294 k := b[PrefixValueLen : PrefixValueLen+kLen]
295 v := b[PrefixValueLen+kLen:]
296 return k, v
297}
298
299func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
300 b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
301 b[0] = 2
302 b[1] = byte(len(l))
303 copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
304 copy(b[PrefixValueLen+hashFunc.Len():], r)
305
306 key, err := hashFunc.Hash(l, r)
307 if err != nil {
308 return nil, nil, err
309 }
310
311 return key, b, nil
312}
313
314func readIntermediateChilds(b []byte) ([]byte, []byte) {
315 if len(b) < PrefixValueLen {
316 return []byte{}, []byte{}
317 }
318
319 lLen := b[1]
320 if len(b) < PrefixValueLen+int(lLen) {
321 return []byte{}, []byte{}
322 }
323 l := b[PrefixValueLen : PrefixValueLen+lLen]
324 r := b[PrefixValueLen+lLen:]
325 return l, r
326}
327
328func getPath(numLevels int, k []byte) []bool {
329 path := make([]bool, numLevels)
330 for n := 0; n < numLevels; n++ {
331 path[n] = k[n/8]&(1<<(n%8)) != 0
332 }
333 return path
334}
335
336// GenProof generates a MerkleTree proof for the given key. If the key exists in
337// the Tree, the proof will be of existence, if the key does not exist in the
338// tree, the proof will be of non-existence.
339func (t *Tree) GenProof(k []byte) ([]byte, error) {
340 keyPath := make([]byte, t.hashFunction.Len())
341 copy(keyPath[:], k)
342
343 path := getPath(t.maxLevels, keyPath)
344 // go down to the leaf
345 var siblings [][]byte
346 _, value, siblings, err := t.down(k, t.root, siblings, path, 0, true)
347 if err != nil {
348 return nil, err
349 }
350
351 leafK, leafV := readLeafValue(value)
352 if !bytes.Equal(k, leafK) {
353 fmt.Println("key not in Tree")
354 fmt.Println(leafK)
355 fmt.Println(leafV)
356 // TODO proof of non-existence
357 panic(fmt.Errorf("unimplemented"))
358 }
359
360 s := PackSiblings(t.hashFunction, siblings)
361 return s, nil
362}
363
364// PackSiblings packs the siblings into a byte array.
365// [ 1 byte | L bytes | 32 * N bytes ]
366// [ bitmap length (L) | bitmap | N non-zero siblings ]
367// Where the bitmap indicates if the sibling is 0 or a value from the siblings array.
368func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte {
369 var b []byte
370 var bitmap []bool
371 emptySibling := make([]byte, hashFunc.Len())
372 for i := 0; i < len(siblings); i++ {
373 if bytes.Equal(siblings[i], emptySibling) {
374 bitmap = append(bitmap, false)
375 } else {
376 bitmap = append(bitmap, true)
377 b = append(b, siblings[i]...)
378 }
379 }
380
381 bitmapBytes := bitmapToBytes(bitmap)
382 l := len(bitmapBytes)
383
384 res := make([]byte, l+1+len(b))
385 res[0] = byte(l) // set the bitmapBytes length
386 copy(res[1:1+l], bitmapBytes)
387 copy(res[1+l:], b)
388 return res
389}
390
391// UnpackSiblings unpacks the siblings from a byte array.
392func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
393 l := b[0]
394 bitmapBytes := b[1 : 1+l]
395 bitmap := bytesToBitmap(bitmapBytes)
396 siblingsBytes := b[1+l:]
397 iSibl := 0
398 emptySibl := make([]byte, hashFunc.Len())
399 var siblings [][]byte
400 for i := 0; i < len(bitmap); i++ {
401 if iSibl >= len(siblingsBytes) {
402 break
403 }
404 if bitmap[i] {
405 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
406 iSibl += hashFunc.Len()
407 } else {
408 siblings = append(siblings, emptySibl)
409 }
410 }
411 return siblings, nil
412}
413
414func bitmapToBytes(bitmap []bool) []byte {
415 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
416 b := make([]byte, bitmapBytesLen)
417 for i := 0; i < len(bitmap); i++ {
418 if bitmap[i] {
419 b[i/8] |= 1 << (i % 8)
420 }
421 }
422 return b
423}
424
425func bytesToBitmap(b []byte) []bool {
426 var bitmap []bool
427 for i := 0; i < len(b); i++ {
428 for j := 0; j < 8; j++ {
429 bitmap = append(bitmap, b[i]&(1<<j) > 0)
430 }
431 }
432 return bitmap
433}
434
435// Get returns the value for a given key
436func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
437 keyPath := make([]byte, t.hashFunction.Len())
438 copy(keyPath[:], k)
439
440 path := getPath(t.maxLevels, keyPath)
441 // go down to the leaf
442 var siblings [][]byte
443 _, value, _, err := t.down(k, t.root, siblings, path, 0, true)
444 if err != nil {
445 return nil, nil, err
446 }
447 leafK, leafV := readLeafValue(value)
448 if !bytes.Equal(k, leafK) {
449 panic(fmt.Errorf("%s != %s", BytesToBigInt(k), BytesToBigInt(leafK)))
450 }
451
452 return leafK, leafV, nil
453}
454
455// CheckProof verifies the given proof. The proof verification depends on the
456// HashFunction passed as parameter.
457func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
458 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
459 if err != nil {
460 return false, err
461 }
462
463 keyPath := make([]byte, hashFunc.Len())
464 copy(keyPath[:], k)
465
466 key, _, err := newLeafValue(hashFunc, k, v)
467 if err != nil {
468 return false, err
469 }
470
471 path := getPath(len(siblings), keyPath)
472 for i := len(siblings) - 1; i >= 0; i-- {
473 if path[i] {
474 key, _, err = newIntermediate(hashFunc, siblings[i], key)
475 if err != nil {
476 return false, err
477 }
478 } else {
479 key, _, err = newIntermediate(hashFunc, key, siblings[i])
480 if err != nil {
481 return false, err
482 }
483 }
484 }
485 if bytes.Equal(key[:], root) {
486 return true, nil
487 }
488 return false, nil
489}