arbo/tree_test.go
tree_test.goRaw
1package arbo
2
3import (
4 "encoding/hex"
5 "math"
6 "math/big"
7 "runtime"
8 "testing"
9 "time"
10
11 qt "github.com/frankban/quicktest"
12 "go.vocdoni.io/dvote/db"
13 "go.vocdoni.io/dvote/db/badgerdb"
14 "go.vocdoni.io/dvote/db/pebbledb"
15)
16
17func checkRootBIString(c *qt.C, tree *Tree, expected string) {
18 root, err := tree.Root()
19 c.Assert(err, qt.IsNil)
20 rootBI := BytesToBigInt(root)
21 c.Check(rootBI.String(), qt.Equals, expected)
22}
23
24func TestDBTx(t *testing.T) {
25 c := qt.New(t)
26
27 dbBadger, err := badgerdb.New(db.Options{Path: c.TempDir()})
28 c.Assert(err, qt.IsNil)
29 testDBTx(c, dbBadger)
30
31 dbPebble, err := pebbledb.New(db.Options{Path: c.TempDir()})
32 c.Assert(err, qt.IsNil)
33 testDBTx(c, dbPebble)
34}
35
36func testDBTx(c *qt.C, database db.Database) {
37 wTx := database.WriteTx()
38 defer wTx.Discard()
39
40 _, err := wTx.Get([]byte("a"))
41 c.Assert(err, qt.Equals, db.ErrKeyNotFound)
42
43 err = wTx.Set([]byte("a"), []byte("b"))
44 c.Assert(err, qt.IsNil)
45
46 v, err := wTx.Get([]byte("a"))
47 c.Assert(err, qt.IsNil)
48 c.Assert(v, qt.DeepEquals, []byte("b"))
49}
50
51func TestAddTestVectors(t *testing.T) {
52 c := qt.New(t)
53
54 // Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
55 testVectorsPoseidon := []string{
56 "0000000000000000000000000000000000000000000000000000000000000000",
57 "13578938674299138072471463694055224830892726234048532520316387704878000008795",
58 "5412393676474193513566895793055462193090331607895808993925969873307089394741",
59 "14204494359367183802864593755198662203838502594566452929175967972147978322084",
60 }
61 testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
62
63 testVectorsSha256 := []string{
64 "0000000000000000000000000000000000000000000000000000000000000000",
65 "46910109172468462938850740851377282682950237270676610513794735904325820156367",
66 "59481735341404520835410489183267411392292882901306595567679529387376287440550",
67 "20573794434149960984975763118181266662429997821552560184909083010514790081771",
68 }
69 testAdd(c, HashFunctionSha256, testVectorsSha256)
70}
71
72func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
73 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
74 c.Assert(err, qt.IsNil)
75 tree, err := NewTree(database, 256, hashFunc)
76 c.Assert(err, qt.IsNil)
77 defer tree.db.Close() //nolint:errcheck
78
79 root, err := tree.Root()
80 c.Assert(err, qt.IsNil)
81 c.Check(hex.EncodeToString(root), qt.Equals, testVectors[0])
82
83 bLen := 32
84 err = tree.Add(
85 BigIntToBytes(bLen, big.NewInt(1)),
86 BigIntToBytes(bLen, big.NewInt(2)))
87 c.Assert(err, qt.IsNil)
88 checkRootBIString(c, tree, testVectors[1])
89
90 err = tree.Add(
91 BigIntToBytes(bLen, big.NewInt(33)),
92 BigIntToBytes(bLen, big.NewInt(44)))
93 c.Assert(err, qt.IsNil)
94 checkRootBIString(c, tree, testVectors[2])
95
96 err = tree.Add(
97 BigIntToBytes(bLen, big.NewInt(1234)),
98 BigIntToBytes(bLen, big.NewInt(9876)))
99 c.Assert(err, qt.IsNil)
100 checkRootBIString(c, tree, testVectors[3])
101}
102
103func TestAddBatch(t *testing.T) {
104 c := qt.New(t)
105 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
106 c.Assert(err, qt.IsNil)
107 tree, err := NewTree(database, 256, HashFunctionPoseidon)
108 c.Assert(err, qt.IsNil)
109 defer tree.db.Close() //nolint:errcheck
110
111 bLen := 32
112 for i := 0; i < 1000; i++ {
113 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
114 v := BigIntToBytes(bLen, big.NewInt(0))
115 if err := tree.Add(k, v); err != nil {
116 t.Fatal(err)
117 }
118 }
119
120 checkRootBIString(c, tree,
121 "296519252211642170490407814696803112091039265640052570497930797516015811235")
122
123 database, err = badgerdb.New(db.Options{Path: c.TempDir()})
124 c.Assert(err, qt.IsNil)
125 tree2, err := NewTree(database, 256, HashFunctionPoseidon)
126 c.Assert(err, qt.IsNil)
127 defer tree2.db.Close() //nolint:errcheck
128
129 var keys, values [][]byte
130 for i := 0; i < 1000; i++ {
131 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
132 v := BigIntToBytes(bLen, big.NewInt(0))
133 keys = append(keys, k)
134 values = append(values, v)
135 }
136 indexes, err := tree2.AddBatch(keys, values)
137 c.Assert(err, qt.IsNil)
138 c.Check(len(indexes), qt.Equals, 0)
139
140 checkRootBIString(c, tree2,
141 "296519252211642170490407814696803112091039265640052570497930797516015811235")
142}
143
144func TestAddDifferentOrder(t *testing.T) {
145 c := qt.New(t)
146 database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
147 c.Assert(err, qt.IsNil)
148 tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
149 c.Assert(err, qt.IsNil)
150 defer tree1.db.Close() //nolint:errcheck
151
152 bLen := 32
153 for i := 0; i < 16; i++ {
154 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
155 v := BigIntToBytes(bLen, big.NewInt(0))
156 if err := tree1.Add(k, v); err != nil {
157 t.Fatal(err)
158 }
159 }
160
161 database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
162 c.Assert(err, qt.IsNil)
163 tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
164 c.Assert(err, qt.IsNil)
165 defer tree2.db.Close() //nolint:errcheck
166
167 for i := 16 - 1; i >= 0; i-- {
168 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
169 v := BigIntToBytes(bLen, big.NewInt(0))
170 if err := tree2.Add(k, v); err != nil {
171 t.Fatal(err)
172 }
173 }
174
175 root1, err := tree1.Root()
176 c.Assert(err, qt.IsNil)
177 root2, err := tree2.Root()
178 c.Assert(err, qt.IsNil)
179 c.Check(hex.EncodeToString(root2), qt.Equals, hex.EncodeToString(root1))
180 c.Check(hex.EncodeToString(root1), qt.Equals,
181 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
182}
183
184func TestAddRepeatedIndex(t *testing.T) {
185 c := qt.New(t)
186 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
187 c.Assert(err, qt.IsNil)
188 tree, err := NewTree(database, 256, HashFunctionPoseidon)
189 c.Assert(err, qt.IsNil)
190 defer tree.db.Close() //nolint:errcheck
191
192 bLen := 32
193 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
194 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
195
196 err = tree.Add(k, v)
197 c.Assert(err, qt.IsNil)
198 err = tree.Add(k, v) // repeating same key-value
199 c.Check(err, qt.Equals, ErrKeyAlreadyExists)
200}
201
202func TestUpdate(t *testing.T) {
203 c := qt.New(t)
204 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
205 c.Assert(err, qt.IsNil)
206 tree, err := NewTree(database, 256, HashFunctionPoseidon)
207 c.Assert(err, qt.IsNil)
208 defer tree.db.Close() //nolint:errcheck
209
210 bLen := 32
211 k := BigIntToBytes(bLen, big.NewInt(int64(20)))
212 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
213 if err := tree.Add(k, v); err != nil {
214 t.Fatal(err)
215 }
216
217 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
218 err = tree.Update(k, v)
219 c.Assert(err, qt.IsNil)
220
221 gettedKey, gettedValue, err := tree.Get(k)
222 c.Assert(err, qt.IsNil)
223 c.Check(gettedKey, qt.DeepEquals, k)
224 c.Check(gettedValue, qt.DeepEquals, v)
225
226 // add more leafs to the tree to do another test
227 for i := 0; i < 16; i++ {
228 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
229 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
230 if err := tree.Add(k, v); err != nil {
231 t.Fatal(err)
232 }
233 }
234
235 k = BigIntToBytes(bLen, big.NewInt(int64(3)))
236 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
237 // check that before the Update, value for 3 is !=11
238 gettedKey, gettedValue, err = tree.Get(k)
239 c.Assert(err, qt.IsNil)
240 c.Check(gettedKey, qt.DeepEquals, k)
241 c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
242 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
243
244 err = tree.Update(k, v)
245 c.Assert(err, qt.IsNil)
246
247 // check that after Update, the value for 3 is ==11
248 gettedKey, gettedValue, err = tree.Get(k)
249 c.Assert(err, qt.IsNil)
250 c.Check(gettedKey, qt.DeepEquals, k)
251 c.Check(gettedValue, qt.DeepEquals, v)
252 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
253}
254
255func TestAux(t *testing.T) { // TODO split in proper tests
256 c := qt.New(t)
257 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
258 c.Assert(err, qt.IsNil)
259 tree, err := NewTree(database, 256, HashFunctionPoseidon)
260 c.Assert(err, qt.IsNil)
261 defer tree.db.Close() //nolint:errcheck
262
263 bLen := 32
264 k := BigIntToBytes(bLen, big.NewInt(int64(1)))
265 v := BigIntToBytes(bLen, big.NewInt(int64(0)))
266 err = tree.Add(k, v)
267 c.Assert(err, qt.IsNil)
268 k = BigIntToBytes(bLen, big.NewInt(int64(256)))
269 err = tree.Add(k, v)
270 c.Assert(err, qt.IsNil)
271
272 k = BigIntToBytes(bLen, big.NewInt(int64(257)))
273 err = tree.Add(k, v)
274 c.Assert(err, qt.IsNil)
275
276 k = BigIntToBytes(bLen, big.NewInt(int64(515)))
277 err = tree.Add(k, v)
278 c.Assert(err, qt.IsNil)
279 k = BigIntToBytes(bLen, big.NewInt(int64(770)))
280 err = tree.Add(k, v)
281 c.Assert(err, qt.IsNil)
282
283 k = BigIntToBytes(bLen, big.NewInt(int64(388)))
284 err = tree.Add(k, v)
285 c.Assert(err, qt.IsNil)
286 k = BigIntToBytes(bLen, big.NewInt(int64(900)))
287 err = tree.Add(k, v)
288 c.Assert(err, qt.IsNil)
289 //
290 // err = tree.PrintGraphviz(nil)
291 // c.Assert(err, qt.IsNil)
292}
293
294func TestGet(t *testing.T) {
295 c := qt.New(t)
296 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
297 c.Assert(err, qt.IsNil)
298 tree, err := NewTree(database, 256, HashFunctionPoseidon)
299 c.Assert(err, qt.IsNil)
300 defer tree.db.Close() //nolint:errcheck
301
302 bLen := 32
303 for i := 0; i < 10; i++ {
304 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
305 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
306 if err := tree.Add(k, v); err != nil {
307 t.Fatal(err)
308 }
309 }
310
311 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
312 gettedKey, gettedValue, err := tree.Get(k)
313 c.Assert(err, qt.IsNil)
314 c.Check(gettedKey, qt.DeepEquals, k)
315 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
316}
317
318func TestGenProofAndVerify(t *testing.T) {
319 c := qt.New(t)
320 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
321 c.Assert(err, qt.IsNil)
322 tree, err := NewTree(database, 256, HashFunctionPoseidon)
323 c.Assert(err, qt.IsNil)
324 defer tree.db.Close() //nolint:errcheck
325
326 bLen := 32
327 for i := 0; i < 10; i++ {
328 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
329 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
330 if err := tree.Add(k, v); err != nil {
331 t.Fatal(err)
332 }
333 }
334
335 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
336 v := BigIntToBytes(bLen, big.NewInt(int64(14)))
337 kAux, proofV, siblings, existence, err := tree.GenProof(k)
338 c.Assert(err, qt.IsNil)
339 c.Assert(proofV, qt.DeepEquals, v)
340 c.Assert(k, qt.DeepEquals, kAux)
341 c.Assert(existence, qt.IsTrue)
342
343 root, err := tree.Root()
344 c.Assert(err, qt.IsNil)
345 verif, err := CheckProof(tree.hashFunction, k, v, root, siblings)
346 c.Assert(err, qt.IsNil)
347 c.Check(verif, qt.IsTrue)
348}
349
350func TestDumpAndImportDump(t *testing.T) {
351 c := qt.New(t)
352 database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
353 c.Assert(err, qt.IsNil)
354 tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
355 c.Assert(err, qt.IsNil)
356 defer tree1.db.Close() //nolint:errcheck
357
358 bLen := 32
359 for i := 0; i < 16; i++ {
360 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
361 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
362 if err := tree1.Add(k, v); err != nil {
363 t.Fatal(err)
364 }
365 }
366
367 e, err := tree1.Dump(nil)
368 c.Assert(err, qt.IsNil)
369
370 database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
371 c.Assert(err, qt.IsNil)
372 tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
373 c.Assert(err, qt.IsNil)
374 defer tree2.db.Close() //nolint:errcheck
375 err = tree2.ImportDump(e)
376 c.Assert(err, qt.IsNil)
377
378 root1, err := tree1.Root()
379 c.Assert(err, qt.IsNil)
380 root2, err := tree2.Root()
381 c.Assert(err, qt.IsNil)
382 c.Check(root2, qt.DeepEquals, root1)
383 c.Check(hex.EncodeToString(root2), qt.Equals,
384 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
385}
386
387func TestRWMutex(t *testing.T) {
388 c := qt.New(t)
389 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
390 c.Assert(err, qt.IsNil)
391 tree, err := NewTree(database, 256, HashFunctionPoseidon)
392 c.Assert(err, qt.IsNil)
393 defer tree.db.Close() //nolint:errcheck
394
395 bLen := 32
396 var keys, values [][]byte
397 for i := 0; i < 1000; i++ {
398 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
399 v := BigIntToBytes(bLen, big.NewInt(0))
400 keys = append(keys, k)
401 values = append(values, v)
402 }
403 go func() {
404 _, err = tree.AddBatch(keys, values)
405 if err != nil {
406 panic(err)
407 }
408 }()
409
410 time.Sleep(500 * time.Millisecond)
411 k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
412 v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
413 if err := tree.Add(k, v); err != nil {
414 t.Fatal(err)
415 }
416}
417
418// TODO UPDATE
419// func TestSetGetNLeafs(t *testing.T) {
420// c := qt.New(t)
421// database, err := badgerdb.New(db.Options{Path: c.TempDir()})
422// c.Assert(err, qt.IsNil)
423// tree, err := NewTree(database, 100, HashFunctionPoseidon)
424// c.Assert(err, qt.IsNil)
425//
426// // 0
427// tree.dbBatch = tree.db.NewBatch()
428//
429// err = tree.setNLeafs(0)
430// c.Assert(err, qt.IsNil)
431//
432// err = tree.dbBatch.Write()
433// c.Assert(err, qt.IsNil)
434//
435// n, err := tree.GetNLeafs()
436// c.Assert(err, qt.IsNil)
437// c.Assert(n, qt.Equals, 0)
438//
439// // 1024
440// tree.dbBatch = tree.db.NewBatch()
441//
442// err = tree.setNLeafs(1024)
443// c.Assert(err, qt.IsNil)
444//
445// err = tree.dbBatch.Write()
446// c.Assert(err, qt.IsNil)
447//
448// n, err = tree.GetNLeafs()
449// c.Assert(err, qt.IsNil)
450// c.Assert(n, qt.Equals, 1024)
451//
452// // 2**64 -1
453// tree.dbBatch = tree.db.NewBatch()
454//
455// maxUint := ^uint(0)
456// maxInt := int(maxUint >> 1)
457//
458// err = tree.setNLeafs(maxInt)
459// c.Assert(err, qt.IsNil)
460//
461// err = tree.dbBatch.Write()
462// c.Assert(err, qt.IsNil)
463//
464// n, err = tree.GetNLeafs()
465// c.Assert(err, qt.IsNil)
466// c.Assert(n, qt.Equals, maxInt)
467// }
468
469func TestAddBatchFullyUsed(t *testing.T) {
470 c := qt.New(t)
471
472 database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
473 c.Assert(err, qt.IsNil)
474 tree1, err := NewTree(database1, 4, HashFunctionPoseidon)
475 c.Assert(err, qt.IsNil)
476
477 database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
478 c.Assert(err, qt.IsNil)
479 tree2, err := NewTree(database2, 4, HashFunctionPoseidon)
480 c.Assert(err, qt.IsNil)
481
482 var keys, values [][]byte
483 for i := 0; i < 16; i++ {
484 k := BigIntToBytes(1, big.NewInt(int64(i)))
485 v := k
486
487 keys = append(keys, k)
488 values = append(values, v)
489
490 // add one by one expecting no error
491 err := tree1.Add(k, v)
492 c.Assert(err, qt.IsNil)
493 }
494
495 invalids, err := tree2.AddBatch(keys, values)
496 c.Assert(err, qt.IsNil)
497 c.Assert(0, qt.Equals, len(invalids))
498
499 root1, err := tree1.Root()
500 c.Assert(err, qt.IsNil)
501 root2, err := tree2.Root()
502 c.Assert(err, qt.IsNil)
503 c.Assert(root1, qt.DeepEquals, root2)
504
505 // get all key-values and check that are equal between both trees
506 for i := 0; i < 16; i++ {
507 auxK1, auxV1, err := tree1.Get(BigIntToBytes(1, big.NewInt(int64(i))))
508 c.Assert(err, qt.IsNil)
509
510 auxK2, auxV2, err := tree2.Get(BigIntToBytes(1, big.NewInt(int64(i))))
511 c.Assert(err, qt.IsNil)
512
513 c.Assert(auxK1, qt.DeepEquals, auxK2)
514 c.Assert(auxV1, qt.DeepEquals, auxV2)
515 }
516
517 // try adding one more key to both trees (through Add & AddBatch) and
518 // expect not being added due the tree is already full
519 k := BigIntToBytes(1, big.NewInt(int64(16)))
520 v := k
521 err = tree1.Add(k, v)
522 c.Assert(err, qt.Equals, ErrMaxVirtualLevel)
523
524 invalids, err = tree2.AddBatch([][]byte{k}, [][]byte{v})
525 c.Assert(err, qt.IsNil)
526 c.Assert(1, qt.Equals, len(invalids))
527}
528
529func TestSetRoot(t *testing.T) {
530 c := qt.New(t)
531 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
532 c.Assert(err, qt.IsNil)
533 tree, err := NewTree(database, 256, HashFunctionPoseidon)
534 c.Assert(err, qt.IsNil)
535
536 expectedRoot := "13742386369878513332697380582061714160370929283209286127733983161245560237407"
537
538 // fill the tree
539 bLen := 32
540 var keys, values [][]byte
541 for i := 0; i < 1000; i++ {
542 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
543 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
544 keys = append(keys, k)
545 values = append(values, v)
546 }
547 indexes, err := tree.AddBatch(keys, values)
548 c.Assert(err, qt.IsNil)
549 c.Check(len(indexes), qt.Equals, 0)
550 checkRootBIString(c, tree,
551 expectedRoot)
552
553 // add one more k-v
554 k := BigIntToBytes(bLen, big.NewInt(1000))
555 v := BigIntToBytes(bLen, big.NewInt(1000))
556 err = tree.Add(k, v)
557 c.Assert(err, qt.IsNil)
558 checkRootBIString(c, tree,
559 "10747149055773881257049574592162159501044114324358186833013814735296193179713")
560
561 // do a SetRoot, and expect the same root than the original tree
562 pastRootBI, ok := new(big.Int).SetString(expectedRoot, 10)
563 c.Assert(ok, qt.IsTrue)
564 pastRoot := BigIntToBytes(32, pastRootBI)
565
566 err = tree.SetRoot(pastRoot)
567 c.Assert(err, qt.IsNil)
568 checkRootBIString(c, tree, expectedRoot)
569
570 // check that the tree can be updated
571 err = tree.Add([]byte("test"), []byte("test"))
572 c.Assert(err, qt.IsNil)
573 err = tree.Update([]byte("test"), []byte("test"))
574 c.Assert(err, qt.IsNil)
575
576 // check that the k-v '1000' does not exist in the new tree
577 _, _, err = tree.Get(k)
578 c.Assert(err, qt.Equals, ErrKeyNotFound)
579
580 // check that can be set an empty root
581 err = tree.SetRoot(tree.emptyHash)
582 c.Assert(err, qt.IsNil)
583}
584
585func TestSnapshot(t *testing.T) {
586 c := qt.New(t)
587 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
588 c.Assert(err, qt.IsNil)
589 tree, err := NewTree(database, 256, HashFunctionPoseidon)
590 c.Assert(err, qt.IsNil)
591
592 // fill the tree
593 bLen := 32
594 var keys, values [][]byte
595 for i := 0; i < 1000; i++ {
596 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
597 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
598 keys = append(keys, k)
599 values = append(values, v)
600 }
601 indexes, err := tree.AddBatch(keys, values)
602 c.Assert(err, qt.IsNil)
603 c.Check(len(indexes), qt.Equals, 0)
604 checkRootBIString(c, tree,
605 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
606
607 // do a snapshot, and expect the same root than the original tree
608 snapshotTree, err := tree.Snapshot(nil)
609 c.Assert(err, qt.IsNil)
610 checkRootBIString(c, snapshotTree,
611 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
612
613 // check that the snapshotTree can not be updated
614 _, err = snapshotTree.AddBatch(keys, values)
615 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
616 err = snapshotTree.Add([]byte("test"), []byte("test"))
617 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
618 err = snapshotTree.Update([]byte("test"), []byte("test"))
619 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
620 err = snapshotTree.ImportDump(nil)
621 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
622
623 // update the original tree by adding a new key-value, and check that
624 // snapshotTree still has the old root, but the original tree has a new
625 // root
626 err = tree.Add([]byte("test"), []byte("test"))
627 c.Assert(err, qt.IsNil)
628 checkRootBIString(c, snapshotTree,
629 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
630 checkRootBIString(c, tree,
631 "1025190963769001718196479367844646783678188389989148142691917685159698888868")
632}
633
634func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
635 c := qt.New(t)
636
637 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
638 c.Assert(err, qt.IsNil)
639 tree, err := NewTree(database, 256, HashFunctionPoseidon)
640 c.Assert(err, qt.IsNil)
641 defer tree.db.Close() //nolint:errcheck
642
643 bLen := 32
644 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
645
646 root, err := tree.Root()
647 c.Assert(err, qt.IsNil)
648 tree, err = tree.Snapshot(root)
649 c.Assert(err, qt.IsNil)
650
651 _, _, err = tree.Get(k)
652 c.Assert(err, qt.Equals, ErrKeyNotFound) // and not equal to db.ErrKeyNotFound
653}
654
655func TestKeyLen(t *testing.T) {
656 c := qt.New(t)
657 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
658 c.Assert(err, qt.IsNil)
659 // maxLevels is 100, keyPath length = ceil(maxLevels/8) = 13
660 maxLevels := 100
661 tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
662 c.Assert(err, qt.IsNil)
663
664 // expect no errors when adding a key of only 4 bytes (when the
665 // required length of keyPath for 100 levels would be 13 bytes)
666 bLen := 4
667 k := BigIntToBytes(bLen, big.NewInt(1))
668 v := BigIntToBytes(bLen, big.NewInt(1))
669
670 err = tree.Add(k, v)
671 c.Assert(err, qt.IsNil)
672
673 err = tree.Update(k, v)
674 c.Assert(err, qt.IsNil)
675
676 _, _, _, _, err = tree.GenProof(k)
677 c.Assert(err, qt.IsNil)
678
679 _, _, err = tree.Get(k)
680 c.Assert(err, qt.IsNil)
681
682 k = BigIntToBytes(bLen, big.NewInt(2))
683 v = BigIntToBytes(bLen, big.NewInt(2))
684 invalids, err := tree.AddBatch([][]byte{k}, [][]byte{v})
685 c.Assert(err, qt.IsNil)
686 c.Assert(len(invalids), qt.Equals, 0)
687
688 // expect errors when adding a key bigger than maximum capacity of the
689 // tree: ceil(maxLevels/8)
690 maxLevels = 32
691 database, err = badgerdb.New(db.Options{Path: c.TempDir()})
692 c.Assert(err, qt.IsNil)
693 tree, err = NewTree(database, maxLevels, HashFunctionBlake2b)
694 c.Assert(err, qt.IsNil)
695
696 maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
697 k = BigIntToBytes(maxKeyLen+1, big.NewInt(1))
698 v = BigIntToBytes(maxKeyLen+1, big.NewInt(1))
699
700 expectedErrMsg := "len(k) can not be bigger than ceil(maxLevels/8)," +
701 " where len(k): 5, maxLevels: 32, max key len=ceil(maxLevels/8): 4." +
702 " Might need a bigger tree depth (maxLevels>=40) in order to input" +
703 " keys of length 5"
704
705 err = tree.Add(k, v)
706 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
707
708 err = tree.Update(k, v)
709 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
710
711 _, _, _, _, err = tree.GenProof(k)
712 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
713
714 _, _, err = tree.Get(k)
715 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
716
717 // check AddBatch with few key-values
718 invalids, err = tree.AddBatch([][]byte{k}, [][]byte{v})
719 c.Assert(err, qt.IsNil)
720 c.Assert(len(invalids), qt.Equals, 1)
721
722 // check AddBatch with many key-values
723 nCPU := flp2(runtime.NumCPU())
724 nKVs := nCPU + 1
725 var ks, vs [][]byte
726 for i := 0; i < nKVs; i++ {
727 ks = append(ks, BigIntToBytes(maxKeyLen+1, big.NewInt(1)))
728 vs = append(vs, BigIntToBytes(maxKeyLen+1, big.NewInt(1)))
729 }
730 invalids, err = tree.AddBatch(ks, vs)
731 c.Assert(err, qt.IsNil)
732 c.Assert(len(invalids), qt.Equals, nKVs)
733
734 // check that with maxKeyLen it can be added
735 k = BigIntToBytes(maxKeyLen, big.NewInt(1))
736 err = tree.Add(k, v)
737 c.Assert(err, qt.IsNil)
738
739 // check CheckProof check with key longer
740 kAux, vAux, packedSiblings, existence, err := tree.GenProof(k)
741 c.Assert(err, qt.IsNil)
742 c.Assert(existence, qt.IsTrue)
743
744 root, err := tree.Root()
745 c.Assert(err, qt.IsNil)
746 verif, err := CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings)
747 c.Assert(err, qt.IsNil)
748 c.Assert(verif, qt.IsTrue)
749
750 // use a similar key but with one zero, expect that CheckProof fails on
751 // the verification
752 kAux = append(kAux, 0)
753 verif, err = CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings)
754 c.Assert(err, qt.IsNil)
755 c.Assert(verif, qt.IsFalse)
756}
757
758func TestKeyLenBiggerThan32(t *testing.T) {
759 c := qt.New(t)
760 maxLevels := 264
761 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
762 c.Assert(err, qt.IsNil)
763 tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
764 c.Assert(err, qt.IsNil)
765
766 bLen := 33
767 err = tree.Add(
768 randomBytes(bLen),
769 randomBytes(bLen))
770 c.Assert(err, qt.IsNil)
771
772 // 2nd key that we add, will find a node with len(key)==32 (due the
773 // hash output size, expect that next Add does not give any error, as
774 // internally it will use a keyPath of size corresponent to the
775 // maxLevels size of the tree
776 err = tree.Add(
777 randomBytes(bLen),
778 randomBytes(bLen))
779 c.Assert(err, qt.IsNil)
780}
781
782func BenchmarkAdd(b *testing.B) {
783 bLen := 32 // for both Poseidon & Sha256
784 // prepare inputs
785 var ks, vs [][]byte
786 for i := 0; i < 1000; i++ {
787 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
788 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
789 ks = append(ks, k)
790 vs = append(vs, v)
791 }
792
793 b.Run("Poseidon", func(b *testing.B) {
794 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
795 })
796 b.Run("Sha256", func(b *testing.B) {
797 benchmarkAdd(b, HashFunctionSha256, ks, vs)
798 })
799}
800
801func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
802 c := qt.New(b)
803 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
804 c.Assert(err, qt.IsNil)
805 tree, err := NewTree(database, 140, hashFunc)
806 c.Assert(err, qt.IsNil)
807 defer tree.db.Close() //nolint:errcheck
808
809 for i := 0; i < len(ks); i++ {
810 if err := tree.Add(ks[i], vs[i]); err != nil {
811 b.Fatal(err)
812 }
813 }
814}