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