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