asmt/asmt_test.go
asmt_test.goRaw
1package asmtree
2
3import (
4 "bytes"
5 "crypto/rand"
6 "fmt"
7 "testing"
8 "time"
9)
10
11func TestTree(t *testing.T) {
12 censusSize := 1000
13 storage := t.TempDir()
14 tr1 := &Tree{}
15 err := tr1.Init("test1", storage)
16 if err != nil {
17 t.Fatal(err)
18 }
19 for i := 0; i < censusSize; i++ {
20 if err = tr1.Add([]byte(fmt.Sprintf("number %d", i)),
21 []byte(fmt.Sprintf("number %d value", i))); err != nil {
22 t.Fatal(err)
23 }
24 }
25
26 root1 := tr1.Root()
27 data, err := tr1.Dump(root1)
28 if err != nil {
29 t.Fatal(err)
30 }
31 t.Logf("dumped data size is: %d bytes", len(data))
32
33 tr2 := &Tree{}
34 err = tr2.Init("test2", storage)
35 if err != nil {
36 t.Fatal(err)
37 }
38 if err = tr2.ImportDump(data); err != nil {
39 t.Fatal(err)
40 }
41 root2 := tr2.Root()
42 if !bytes.Equal(root1, root2) {
43 t.Errorf("roots are different but they should be equal (%x != %x)", root1, root2)
44 }
45
46 // Try closing the storage and creating the tree again
47 tr2.Close()
48 err = tr2.Init("test2", storage)
49 if err != nil {
50 t.Fatal(err)
51 }
52
53 // Get the size
54 s, err := tr2.Size(nil)
55 if err != nil {
56 t.Errorf("cannot get te size of the tree after reopen: (%s)", err)
57 }
58 if s != int64(censusSize) {
59 t.Errorf("Size is wrong (have %d, expexted %d)", s, censusSize)
60 }
61
62 // Check Root is still the same
63 if !bytes.Equal(tr2.Root(), root2) {
64 t.Fatalf("after closing and opening the tree, the root is different")
65 }
66
67 // Generate a proof on tr1 and check validity on snapshot and tr2
68 proof1, err := tr1.GenProof([]byte("number 5"), []byte("number 5 value"))
69 if err != nil {
70 t.Error(err)
71 }
72 t.Logf("Proof Length: %d", len(proof1))
73 tr1s, err := tr1.Snapshot(root1)
74 if err != nil {
75 t.Fatal(err)
76 }
77 valid, err := tr1s.CheckProof([]byte("number 5"), []byte("number 5 value"), root1, proof1)
78 if err != nil {
79 t.Error(err)
80 }
81 if !valid {
82 t.Errorf("proof is invalid on snapshot")
83 }
84 valid, err = tr2.CheckProof([]byte("number 5"), []byte("number 5 value"), nil, proof1)
85 if err != nil {
86 t.Error(err)
87 }
88 if !valid {
89 t.Errorf("proof is invalid on tree2")
90 }
91}
92
93func TestProofs(t *testing.T) {
94 censusSize := 10000
95 storage := t.TempDir()
96 tr1 := &Tree{}
97 err := tr1.Init("test1", storage)
98 if err != nil {
99 t.Fatal(err)
100 }
101
102 var keys, values [][]byte
103 for i := 0; i < censusSize; i++ {
104 keys = append(keys, RandomBytes(32))
105 values = append(values, RandomBytes(32))
106 }
107
108 i := 0
109 for i < censusSize-200 {
110 if fail, err := tr1.AddBatch(keys[i:i+200], values[i:i+200]); err != nil {
111 t.Fatal(err)
112 } else if len(fail) > 0 {
113 t.Fatalf("some keys failed to add on addBatch: %v", fail)
114 }
115 i += 200
116 }
117 // Add remaining claims (if size%200 != 0)
118 if i < censusSize {
119 if fail, err := tr1.AddBatch(keys[i:], values[i:]); err != nil {
120 t.Fatal(err)
121 } else if len(fail) > 0 {
122 t.Fatalf("some keys failed to add on addBatch: %v", fail)
123 }
124 }
125
126 root1 := tr1.Root()
127 data, err := tr1.Dump(root1)
128 if err != nil {
129 t.Fatal(err)
130 }
131 t.Logf("dumped data size is: %d bytes", len(data))
132
133 // Get the size
134 s, err := tr1.Size(nil)
135 if err != nil {
136 t.Errorf("cannot get te size: %v", err)
137 }
138 if s != int64(censusSize) {
139 t.Errorf("size is wrong (have %d, expexted %d)", s, censusSize)
140 }
141
142 // Generate a proofs
143 time.Sleep(5 * time.Second)
144 proofs := [][]byte{}
145 for i := 0; i < censusSize; i++ {
146 p, err := tr1.GenProof(keys[i], values[i])
147 if err != nil {
148 t.Fatal(err)
149 }
150 if len(p) == 0 {
151 t.Fatal("proof not generated")
152 }
153 proofs = append(proofs, p)
154 }
155
156 // Check proofs
157 for i := 0; i < censusSize; i++ {
158 valid, err := tr1.CheckProof(keys[i], values[i], root1, proofs[i])
159 if err != nil {
160 t.Fatal(err)
161 }
162 if !valid {
163 t.Errorf("proof %d is invalid", i)
164 }
165 }
166
167}
168
169// go test -v -run=- -bench=Tree -benchtime=30s .
170func BenchmarkTree(b *testing.B) {
171 b.ReportAllocs()
172
173 b.ResetTimer()
174 b.RunParallel(func(pb *testing.PB) {
175 // Create websocket client
176 for pb.Next() {
177 benchProofs(b, 100000)
178 }
179 })
180}
181
182func benchProofs(b *testing.B, censusSize int) {
183 totalTimer := time.Now()
184 storage := b.TempDir()
185 tr1 := &Tree{}
186 err := tr1.Init("test1", storage)
187 if err != nil {
188 b.Fatal(err)
189 }
190
191 var keys, values [][]byte
192 for i := 0; i < censusSize; i++ {
193 keys = append(keys, RandomBytes(32))
194 values = append(values, RandomBytes(32))
195 }
196
197 timer := time.Now()
198 i := 0
199 for i < censusSize-200 {
200 if fail, err := tr1.AddBatch(keys[i:i+200], values[i:i+200]); err != nil {
201 b.Fatal(err)
202 } else if len(fail) > 0 {
203 b.Fatalf("some keys failed to add on addBatch: %v", fail)
204 }
205 i += 200
206 }
207 // Add remaining claims (if size%200 != 0)
208 if i < censusSize {
209 if fail, err := tr1.AddBatch(keys[i:], values[i:]); err != nil {
210 b.Fatal(err)
211 } else if len(fail) > 0 {
212 b.Fatalf("some keys failed to add on addBatch: %v", fail)
213 }
214 }
215 b.Logf("addBatch took %d ms", time.Since(timer).Milliseconds())
216
217 timer = time.Now()
218 root1 := tr1.Root()
219 data, err := tr1.Dump(root1)
220 if err != nil {
221 b.Fatal(err)
222 }
223 b.Logf("dumped data size is: %d bytes", len(data))
224 b.Logf("dump took %d ms", time.Since(timer).Milliseconds())
225
226 // Get the size
227 s, err := tr1.Size(nil)
228 if err != nil {
229 b.Errorf("cannot get te size: %v", err)
230 }
231 if s != int64(censusSize) {
232 b.Errorf("size is wrong (have %d, expexted %d)", s, censusSize)
233 }
234
235 // Generate a proofs
236 timer = time.Now()
237 time.Sleep(5 * time.Second)
238 proofs := [][]byte{}
239 for i := 0; i < censusSize; i++ {
240 p, err := tr1.GenProof(keys[i], values[i])
241 if err != nil {
242 b.Fatal(err)
243 }
244 if len(p) == 0 {
245 b.Fatal("proof not generated")
246 }
247 proofs = append(proofs, p)
248 }
249 b.Logf("gen proofs took %d ms", time.Since(timer).Milliseconds())
250
251 // Check proofs
252 timer = time.Now()
253 for i := 0; i < censusSize; i++ {
254 valid, err := tr1.CheckProof(keys[i], values[i], root1, proofs[i])
255 if err != nil {
256 b.Fatal(err)
257 }
258 if !valid {
259 b.Errorf("proof %d is invalid", i)
260 }
261 }
262 b.Logf("check proofs took %d ms", time.Since(timer).Milliseconds())
263 b.Logf("[finished] %d ms", time.Since(totalTimer).Milliseconds())
264}
265
266func RandomBytes(n int) []byte {
267 b := make([]byte, n)
268 _, err := rand.Read(b)
269 if err != nil {
270 panic(err)
271 }
272 return b
273}