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