asmt/asmt.go
asmt.goRaw
1package asmtree
2
3import (
4 "bytes"
5 "fmt"
6 "path"
7 "sync/atomic"
8 "time"
9
10 "go.vocdoni.io/dvote/log"
11 "go.vocdoni.io/proto/build/go/models"
12
13 "git.sr.ht/~sircmpwn/go-bare"
14 "github.com/p4u/asmt/db"
15 asmt "github.com/p4u/asmt/smt"
16)
17
18// We use go-bare for export/import the trie. In order to support
19// big census (up to 8 Million entries) we need to increase the maximums.
20const bareMaxArrayLength uint64 = 1024 * 1014 * 8 // 8 Million
21
22const bareMaxUnmarshalBytes uint64 = bareMaxArrayLength * 32 * 2 // 512 MiB
23
24type Tree struct {
25 Tree *asmt.Trie
26 db db.DB
27 public uint32
28 lastAccessUnix int64 // a unix timestamp, used via sync/atomic
29 size uint64
30 snapshotRoot []byte // if not nil, this trie is considered an inmutable snapshot
31 snapshotSize uint64
32}
33
34type Proof struct {
35 Bitmap []byte
36 Length int
37 Siblings [][]byte
38 Value []byte
39}
40
41type exportElement struct {
42 Key []byte `bare:"key"`
43 Value []byte `bare:"value"`
44}
45
46type exportData struct {
47 Elements []exportElement `bare:"elements"`
48}
49
50const (
51 MaxKeySize = 32
52 MaxValueSize = 64
53 dbRootPrefix = "this is the last root for the SMT tree"
54)
55
56// NewTree initializes a new AergoSMT tree following the Tree interface specification.
57func NewTree(name, storageDir string) (*Tree, error) {
58 tr := &Tree{}
59 err := tr.Init(name, storageDir)
60 return tr, err
61}
62
63// newTree opens or creates a merkle tree under the given storage.
64func newTree(name, storageDir string) (*asmt.Trie, db.DB, error) {
65 dir := path.Join(storageDir, name)
66 log.Debugf("creating new tree on %s", dir)
67 d := db.NewDB(db.LevelImpl, dir)
68 root := d.Get([]byte(dbRootPrefix))
69 tr := asmt.NewTrie(root, asmt.Hasher, d)
70 if root != nil {
71 if err := tr.LoadCache(root); err != nil {
72 return nil, nil, err
73 }
74 }
75 return tr, d, nil
76}
77
78// Init initializes a new asmt tree
79func (t *Tree) Init(name, storageDir string) error {
80 var err error
81 t.Tree, t.db, err = newTree(name, storageDir)
82 t.updateAccessTime()
83 t.size = 0
84 return err
85}
86
87func (t *Tree) MaxKeySize() int {
88 return MaxKeySize
89}
90
91// LastAccess returns the last time the Tree was accessed, in the form of a unix
92// timestamp.
93func (t *Tree) LastAccess() int64 {
94 return atomic.LoadInt64(&t.lastAccessUnix)
95}
96
97func (t *Tree) updateAccessTime() {
98 atomic.StoreInt64(&t.lastAccessUnix, time.Now().Unix())
99}
100
101// Publish makes a merkle tree available for queries.
102// Application layer should check IsPublish() before considering the Tree available.
103func (t *Tree) Publish() {
104 atomic.StoreUint32(&t.public, 1)
105}
106
107// UnPublish makes a merkle tree not available for queries
108func (t *Tree) UnPublish() {
109 atomic.StoreUint32(&t.public, 0)
110}
111
112// IsPublic returns true if the tree is available
113func (t *Tree) IsPublic() bool {
114 return atomic.LoadUint32(&t.public) == 1
115}
116
117// Type returns the numeric identifier for the censusTree implementation
118func (t *Tree) Type() models.Census_Type {
119 return models.Census_UNKNOWN
120}
121
122// TypeString returns the name identifying the censustree implementation
123func (t *Tree) TypeString() string {
124 return models.Census_Type_name[int32(t.Type())] // reuse the naming generated by protobuf
125}
126
127// Commit saves permanently the tree on disk
128func (t *Tree) Commit() error {
129 if t.snapshotRoot != nil {
130 return fmt.Errorf("cannot commit to a snapshot trie")
131 }
132 err := t.Tree.Commit()
133 if err != nil {
134 return err
135 }
136 t.db.Set([]byte(dbRootPrefix), t.Root())
137 return nil
138}
139
140// Add adds a new claim to the merkle tree
141// A claim is composed of two parts: index and value
142// 1.index is mandatory, the data will be used for indexing the claim into to merkle tree
143// 2.value is optional, the data will not affect the indexing
144func (t *Tree) Add(index, value []byte) error {
145 t.updateAccessTime()
146 if t.snapshotRoot != nil {
147 return fmt.Errorf("cannot add to a snapshot trie")
148 }
149 if len(index) < 4 || len(index) > MaxKeySize {
150 return fmt.Errorf("wrong key size: %d", len(index))
151 }
152 if len(value) > MaxValueSize {
153 return fmt.Errorf("index or value claim data too big")
154 }
155 _, err := t.Tree.Update([][]byte{asmt.Hasher(index)}, [][]byte{asmt.Hasher(value)})
156 if err != nil {
157 return err
158 }
159 atomic.StoreUint64(&t.size, 0) // TBD: improve this
160 return t.Commit()
161}
162
163// AddBatch adds a list of indexes and values.
164// The commit to disk is executed only once.
165// The values slince could be empty or as long as indexes.
166func (t *Tree) AddBatch(indexes, values [][]byte) ([]int, error) {
167 var wrongIndexes []int
168 t.updateAccessTime()
169 if t.snapshotRoot != nil {
170 return wrongIndexes, fmt.Errorf("cannot add to a snapshot trie")
171 }
172 if len(values) > 0 && len(indexes) != len(values) {
173 return wrongIndexes, fmt.Errorf("indexes and values have different size")
174 }
175 var value []byte
176 for i, key := range indexes {
177 if len(key) < 4 || len(key) > MaxKeySize {
178 wrongIndexes = append(wrongIndexes, i)
179 continue
180 }
181 value = nil
182 if len(values) > 0 {
183 if len(values[i]) > MaxValueSize {
184 wrongIndexes = append(wrongIndexes, i)
185 continue
186 }
187 value = values[i]
188 }
189 _, err := t.Tree.Update([][]byte{asmt.Hasher(key)}, [][]byte{asmt.Hasher(value)})
190 if err != nil {
191 return wrongIndexes, err
192 }
193 }
194 atomic.StoreUint64(&t.size, 0) // TBD: improve this
195 return wrongIndexes, t.Commit()
196}
197
198// Get returns the value of a key
199func (t *Tree) Get(key []byte) []byte { // Do something with error
200 var value []byte
201 if t.snapshotRoot != nil {
202 value, _ = t.Tree.GetWithRoot(key, t.snapshotRoot)
203 } else {
204 value, _ = t.Tree.Get(key)
205 }
206 return value
207}
208
209// GenProof generates a merkle tree proof that can be later used on CheckProof() to validate it.
210func (t *Tree) GenProof(index, value []byte) ([]byte, error) {
211 t.updateAccessTime()
212 var err error
213 var ap [][]byte
214 var pvalue, bitmap []byte
215 var length int
216 var included bool
217 key := asmt.Hasher(index)
218 if t.snapshotRoot != nil {
219 bitmap, ap, length, included, _, pvalue, err = t.Tree.MerkleProofCompressedR(key,
220 t.snapshotRoot)
221 if err != nil {
222 return nil, err
223 }
224 } else {
225 bitmap, ap, length, included, _, pvalue, err = t.Tree.MerkleProofCompressed(key)
226 if err != nil {
227 return nil, err
228 }
229 }
230 if !included {
231 return nil, nil
232 }
233 if !bytes.Equal(pvalue, asmt.Hasher(value)) {
234 return nil, fmt.Errorf("incorrect value or key on genProof")
235 }
236 return bare.Marshal(&Proof{Bitmap: bitmap, Length: length, Siblings: ap, Value: pvalue})
237}
238
239// CheckProof validates a merkle proof and its data.
240func (t *Tree) CheckProof(index, value, root, mproof []byte) (bool, error) {
241 t.updateAccessTime()
242 p := Proof{}
243 if err := bare.Unmarshal(mproof, &p); err != nil {
244 return false, err
245 }
246 if !bytes.Equal(p.Value, asmt.Hasher(value)) {
247 return false, fmt.Errorf("values mismatch %x != %x", p.Value, asmt.Hasher(value))
248 }
249 if root != nil {
250 return t.Tree.VerifyInclusionWithRootC(
251 root,
252 p.Bitmap,
253 asmt.Hasher(index),
254 p.Value,
255 p.Siblings,
256 p.Length), nil
257 }
258 if t.snapshotRoot != nil {
259 return t.Tree.VerifyInclusionWithRootC(
260 t.snapshotRoot,
261 p.Bitmap,
262 asmt.Hasher(index),
263 p.Value,
264 p.Siblings,
265 p.Length), nil
266 }
267 return t.Tree.VerifyInclusionC(
268 p.Bitmap,
269 asmt.Hasher(index),
270 p.Value,
271 p.Siblings,
272 p.Length), nil
273}
274
275// Root returns the current root hash of the merkle tree
276func (t *Tree) Root() []byte {
277 t.updateAccessTime()
278 if t.snapshotRoot != nil {
279 return t.snapshotRoot
280 }
281 return t.Tree.Root
282}
283
284// Dump returns the whole merkle tree serialized in a format that can be used on Import.
285// Byte seralization is performed using bare message protocol, it is a 40% size win over JSON.
286func (t *Tree) Dump(root []byte) ([]byte, error) {
287 t.updateAccessTime()
288 if root == nil && t.snapshotRoot != nil {
289 root = t.snapshotRoot
290 }
291 dump := exportData{}
292 t.iterateWithRoot(root, nil, func(k, v []byte) bool {
293 ee := exportElement{Key: make([]byte, len(k)), Value: make([]byte, len(v))}
294 // Copy elements since it's not safe to hold on to the []byte values from Iterate
295 copy(ee.Key, k[:])
296 copy(ee.Value, v[:])
297 dump.Elements = append(dump.Elements, ee)
298 return false
299 })
300 bare.MaxArrayLength(bareMaxArrayLength)
301 bare.MaxUnmarshalBytes(bareMaxUnmarshalBytes)
302
303 return bare.Marshal(&dump)
304}
305
306// String returns a human readable representation of the tree.
307func (t *Tree) String() string {
308 s := bytes.Buffer{}
309 t.iterate(t.snapshotRoot, func(k, v []byte) bool {
310 s.WriteString(fmt.Sprintf("%x => %x\n", k, v))
311 return false
312 })
313 return s.String()
314}
315
316// Size returns the number of leaf nodes on the merkle tree.
317// TO-DO: root is currently ignored
318func (t *Tree) Size(root []byte) (int64, error) {
319 if t.snapshotRoot != nil {
320 return int64(t.snapshotSize), nil
321 }
322 return int64(t.count()), nil
323}
324
325// DumpPlain returns the entire list of added claims for a specific root hash.
326// First return parametre are the indexes and second the values.
327// If root is not specified, the last one is used.
328func (t *Tree) DumpPlain(root []byte) ([][]byte, [][]byte, error) {
329 var indexes, values [][]byte
330 var err error
331 t.updateAccessTime()
332
333 t.iterateWithRoot(root, nil, func(k, v []byte) bool {
334 indexes = append(indexes, k)
335 values = append(values, v)
336 return false
337 })
338
339 return indexes, values, err
340}
341
342// ImportDump imports a partial or whole tree previously exported with Dump()
343func (t *Tree) ImportDump(data []byte) error {
344 t.updateAccessTime()
345 if t.snapshotRoot != nil {
346 return fmt.Errorf("cannot import to a snapshot")
347 }
348 census := new(exportData)
349 bare.MaxArrayLength(bareMaxArrayLength)
350 bare.MaxUnmarshalBytes(bareMaxUnmarshalBytes)
351 if err := bare.Unmarshal(data, census); err != nil {
352 return fmt.Errorf("importdump cannot unmarshal data: %w", err)
353 }
354 keys := [][]byte{}
355 values := [][]byte{}
356 for _, ee := range census.Elements {
357 keys = append(keys, ee.Key)
358 values = append(values, ee.Value)
359 }
360 _, err := t.Tree.Update(keys, values)
361 if err != nil {
362 return err
363 }
364 atomic.StoreUint64(&t.size, 0) // TBD: improve this
365 return t.Commit()
366}
367
368// Snapshot returns a Tree instance of a exiting merkle root.
369// A Snapshot cannot be modified.
370func (t *Tree) Snapshot(root []byte) (*Tree, error) {
371 exist, err := t.HashExists(root)
372 if err != nil {
373 return nil, err
374 }
375 if !exist {
376 return nil, fmt.Errorf("root %x does not exist, cannot build snapshot", root)
377 }
378 return &Tree{Tree: t.Tree, public: t.public, snapshotRoot: root, snapshotSize: t.count()}, nil
379}
380
381func (t *Tree) Close() error {
382 t.db.Close()
383 return nil
384}
385
386// HashExists checks if a hash exists as a node in the merkle tree
387func (t *Tree) HashExists(hash []byte) (bool, error) {
388 t.updateAccessTime()
389 return t.Tree.TrieRootExists(hash), nil
390}
391
392func (t *Tree) count() uint64 {
393 if v := atomic.LoadUint64(&t.size); v != 0 {
394 return v
395 }
396 counter := uint64(0)
397 if err := t.Tree.Walk(t.snapshotRoot, func(*asmt.WalkResult) int32 {
398 counter++
399 return 0
400 }); err != nil {
401 return 0
402 }
403 atomic.StoreUint64(&t.size, counter)
404 return counter
405}
406
407func (t *Tree) iterate(prefix []byte, callback func(key, value []byte) bool) {
408 t.Tree.Walk(t.snapshotRoot, func(v *asmt.WalkResult) int32 {
409 if callback(v.Key, v.Value) {
410 return 1
411 } else {
412 return 0
413 }
414 })
415}
416
417func (t *Tree) iterateWithRoot(root, prefix []byte, callback func(key, value []byte) bool) {
418 t.Tree.Walk(root, func(v *asmt.WalkResult) int32 {
419 if callback(v.Key, v.Value) {
420 return 1
421 } else {
422 return 0
423 }
424 })
425}