package asmtree
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"path"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"go.vocdoni.io/dvote/log"
|
|
"go.vocdoni.io/proto/build/go/models"
|
|
|
|
"git.sr.ht/~sircmpwn/go-bare"
|
|
"github.com/p4u/asmt/db"
|
|
asmt "github.com/p4u/asmt/smt"
|
|
)
|
|
|
|
// We use go-bare for export/import the trie. In order to support
|
|
// big census (up to 8 Million entries) we need to increase the maximums.
|
|
const bareMaxArrayLength uint64 = 1024 * 1014 * 8 // 8 Million
|
|
|
|
const bareMaxUnmarshalBytes uint64 = bareMaxArrayLength * 32 * 2 // 512 MiB
|
|
|
|
type Tree struct {
|
|
Tree *asmt.Trie
|
|
db db.DB
|
|
public uint32
|
|
lastAccessUnix int64 // a unix timestamp, used via sync/atomic
|
|
size uint64
|
|
snapshotRoot []byte // if not nil, this trie is considered an inmutable snapshot
|
|
snapshotSize uint64
|
|
}
|
|
|
|
type Proof struct {
|
|
Bitmap []byte
|
|
Length int
|
|
Siblings [][]byte
|
|
Value []byte
|
|
}
|
|
|
|
type exportElement struct {
|
|
Key []byte `bare:"key"`
|
|
Value []byte `bare:"value"`
|
|
}
|
|
|
|
type exportData struct {
|
|
Elements []exportElement `bare:"elements"`
|
|
}
|
|
|
|
const (
|
|
MaxKeySize = 32
|
|
MaxValueSize = 64
|
|
dbRootPrefix = "this is the last root for the SMT tree"
|
|
)
|
|
|
|
// NewTree initializes a new AergoSMT tree following the Tree interface specification.
|
|
func NewTree(name, storageDir string) (*Tree, error) {
|
|
tr := &Tree{}
|
|
err := tr.Init(name, storageDir)
|
|
return tr, err
|
|
}
|
|
|
|
// newTree opens or creates a merkle tree under the given storage.
|
|
func newTree(name, storageDir string) (*asmt.Trie, db.DB, error) {
|
|
dir := path.Join(storageDir, name)
|
|
log.Debugf("creating new tree on %s", dir)
|
|
d := db.NewDB(db.LevelImpl, dir)
|
|
root := d.Get([]byte(dbRootPrefix))
|
|
tr := asmt.NewTrie(root, asmt.Hasher, d)
|
|
if root != nil {
|
|
if err := tr.LoadCache(root); err != nil {
|
|
return nil, nil, err
|
|
}
|
|
}
|
|
return tr, d, nil
|
|
}
|
|
|
|
// Init initializes a new asmt tree
|
|
func (t *Tree) Init(name, storageDir string) error {
|
|
var err error
|
|
t.Tree, t.db, err = newTree(name, storageDir)
|
|
t.updateAccessTime()
|
|
t.size = 0
|
|
return err
|
|
}
|
|
|
|
func (t *Tree) MaxKeySize() int {
|
|
return MaxKeySize
|
|
}
|
|
|
|
// LastAccess returns the last time the Tree was accessed, in the form of a unix
|
|
// timestamp.
|
|
func (t *Tree) LastAccess() int64 {
|
|
return atomic.LoadInt64(&t.lastAccessUnix)
|
|
}
|
|
|
|
func (t *Tree) updateAccessTime() {
|
|
atomic.StoreInt64(&t.lastAccessUnix, time.Now().Unix())
|
|
}
|
|
|
|
// Publish makes a merkle tree available for queries.
|
|
// Application layer should check IsPublish() before considering the Tree available.
|
|
func (t *Tree) Publish() {
|
|
atomic.StoreUint32(&t.public, 1)
|
|
}
|
|
|
|
// UnPublish makes a merkle tree not available for queries
|
|
func (t *Tree) UnPublish() {
|
|
atomic.StoreUint32(&t.public, 0)
|
|
}
|
|
|
|
// IsPublic returns true if the tree is available
|
|
func (t *Tree) IsPublic() bool {
|
|
return atomic.LoadUint32(&t.public) == 1
|
|
}
|
|
|
|
// Type returns the numeric identifier for the censusTree implementation
|
|
func (t *Tree) Type() models.Census_Type {
|
|
return models.Census_UNKNOWN
|
|
}
|
|
|
|
// TypeString returns the name identifying the censustree implementation
|
|
func (t *Tree) TypeString() string {
|
|
return models.Census_Type_name[int32(t.Type())] // reuse the naming generated by protobuf
|
|
}
|
|
|
|
// Commit saves permanently the tree on disk
|
|
func (t *Tree) Commit() error {
|
|
if t.snapshotRoot != nil {
|
|
return fmt.Errorf("cannot commit to a snapshot trie")
|
|
}
|
|
err := t.Tree.Commit()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
t.db.Set([]byte(dbRootPrefix), t.Root())
|
|
return nil
|
|
}
|
|
|
|
// Add adds a new claim to the merkle tree
|
|
// A claim is composed of two parts: index and value
|
|
// 1.index is mandatory, the data will be used for indexing the claim into to merkle tree
|
|
// 2.value is optional, the data will not affect the indexing
|
|
func (t *Tree) Add(index, value []byte) error {
|
|
t.updateAccessTime()
|
|
if t.snapshotRoot != nil {
|
|
return fmt.Errorf("cannot add to a snapshot trie")
|
|
}
|
|
if len(index) < 4 || len(index) > MaxKeySize {
|
|
return fmt.Errorf("wrong key size: %d", len(index))
|
|
}
|
|
if len(value) > MaxValueSize {
|
|
return fmt.Errorf("index or value claim data too big")
|
|
}
|
|
_, err := t.Tree.Update([][]byte{asmt.Hasher(index)}, [][]byte{asmt.Hasher(value)})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
atomic.StoreUint64(&t.size, 0) // TBD: improve this
|
|
return t.Commit()
|
|
}
|
|
|
|
// AddBatch adds a list of indexes and values.
|
|
// The commit to disk is executed only once.
|
|
// The values slince could be empty or as long as indexes.
|
|
func (t *Tree) AddBatch(indexes, values [][]byte) ([]int, error) {
|
|
var wrongIndexes []int
|
|
t.updateAccessTime()
|
|
if t.snapshotRoot != nil {
|
|
return wrongIndexes, fmt.Errorf("cannot add to a snapshot trie")
|
|
}
|
|
if len(values) > 0 && len(indexes) != len(values) {
|
|
return wrongIndexes, fmt.Errorf("indexes and values have different size")
|
|
}
|
|
var value []byte
|
|
for i, key := range indexes {
|
|
if len(key) < 4 || len(key) > MaxKeySize {
|
|
wrongIndexes = append(wrongIndexes, i)
|
|
continue
|
|
}
|
|
value = nil
|
|
if len(values) > 0 {
|
|
if len(values[i]) > MaxValueSize {
|
|
wrongIndexes = append(wrongIndexes, i)
|
|
continue
|
|
}
|
|
value = values[i]
|
|
}
|
|
_, err := t.Tree.Update([][]byte{asmt.Hasher(key)}, [][]byte{asmt.Hasher(value)})
|
|
if err != nil {
|
|
return wrongIndexes, err
|
|
}
|
|
}
|
|
atomic.StoreUint64(&t.size, 0) // TBD: improve this
|
|
return wrongIndexes, t.Commit()
|
|
}
|
|
|
|
// Get returns the value of a key
|
|
func (t *Tree) Get(key []byte) []byte { // Do something with error
|
|
var value []byte
|
|
if t.snapshotRoot != nil {
|
|
value, _ = t.Tree.GetWithRoot(key, t.snapshotRoot)
|
|
} else {
|
|
value, _ = t.Tree.Get(key)
|
|
}
|
|
return value
|
|
}
|
|
|
|
// GenProof generates a merkle tree proof that can be later used on CheckProof() to validate it.
|
|
func (t *Tree) GenProof(index, value []byte) ([]byte, error) {
|
|
t.updateAccessTime()
|
|
var err error
|
|
var ap [][]byte
|
|
var pvalue, bitmap []byte
|
|
var length int
|
|
var included bool
|
|
key := asmt.Hasher(index)
|
|
if t.snapshotRoot != nil {
|
|
bitmap, ap, length, included, _, pvalue, err = t.Tree.MerkleProofCompressedR(key,
|
|
t.snapshotRoot)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
} else {
|
|
bitmap, ap, length, included, _, pvalue, err = t.Tree.MerkleProofCompressed(key)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
if !included {
|
|
return nil, nil
|
|
}
|
|
if !bytes.Equal(pvalue, asmt.Hasher(value)) {
|
|
return nil, fmt.Errorf("incorrect value or key on genProof")
|
|
}
|
|
return bare.Marshal(&Proof{Bitmap: bitmap, Length: length, Siblings: ap, Value: pvalue})
|
|
}
|
|
|
|
// CheckProof validates a merkle proof and its data.
|
|
func (t *Tree) CheckProof(index, value, root, mproof []byte) (bool, error) {
|
|
t.updateAccessTime()
|
|
p := Proof{}
|
|
if err := bare.Unmarshal(mproof, &p); err != nil {
|
|
return false, err
|
|
}
|
|
if !bytes.Equal(p.Value, asmt.Hasher(value)) {
|
|
return false, fmt.Errorf("values mismatch %x != %x", p.Value, asmt.Hasher(value))
|
|
}
|
|
if root != nil {
|
|
return t.Tree.VerifyInclusionWithRootC(
|
|
root,
|
|
p.Bitmap,
|
|
asmt.Hasher(index),
|
|
p.Value,
|
|
p.Siblings,
|
|
p.Length), nil
|
|
}
|
|
if t.snapshotRoot != nil {
|
|
return t.Tree.VerifyInclusionWithRootC(
|
|
t.snapshotRoot,
|
|
p.Bitmap,
|
|
asmt.Hasher(index),
|
|
p.Value,
|
|
p.Siblings,
|
|
p.Length), nil
|
|
}
|
|
return t.Tree.VerifyInclusionC(
|
|
p.Bitmap,
|
|
asmt.Hasher(index),
|
|
p.Value,
|
|
p.Siblings,
|
|
p.Length), nil
|
|
}
|
|
|
|
// Root returns the current root hash of the merkle tree
|
|
func (t *Tree) Root() []byte {
|
|
t.updateAccessTime()
|
|
if t.snapshotRoot != nil {
|
|
return t.snapshotRoot
|
|
}
|
|
return t.Tree.Root
|
|
}
|
|
|
|
// Dump returns the whole merkle tree serialized in a format that can be used on Import.
|
|
// Byte seralization is performed using bare message protocol, it is a 40% size win over JSON.
|
|
func (t *Tree) Dump(root []byte) ([]byte, error) {
|
|
t.updateAccessTime()
|
|
if root == nil && t.snapshotRoot != nil {
|
|
root = t.snapshotRoot
|
|
}
|
|
dump := exportData{}
|
|
t.iterateWithRoot(root, nil, func(k, v []byte) bool {
|
|
ee := exportElement{Key: make([]byte, len(k)), Value: make([]byte, len(v))}
|
|
// Copy elements since it's not safe to hold on to the []byte values from Iterate
|
|
copy(ee.Key, k[:])
|
|
copy(ee.Value, v[:])
|
|
dump.Elements = append(dump.Elements, ee)
|
|
return false
|
|
})
|
|
bare.MaxArrayLength(bareMaxArrayLength)
|
|
bare.MaxUnmarshalBytes(bareMaxUnmarshalBytes)
|
|
|
|
return bare.Marshal(&dump)
|
|
}
|
|
|
|
// String returns a human readable representation of the tree.
|
|
func (t *Tree) String() string {
|
|
s := bytes.Buffer{}
|
|
t.iterate(t.snapshotRoot, func(k, v []byte) bool {
|
|
s.WriteString(fmt.Sprintf("%x => %x\n", k, v))
|
|
return false
|
|
})
|
|
return s.String()
|
|
}
|
|
|
|
// Size returns the number of leaf nodes on the merkle tree.
|
|
// TO-DO: root is currently ignored
|
|
func (t *Tree) Size(root []byte) (int64, error) {
|
|
if t.snapshotRoot != nil {
|
|
return int64(t.snapshotSize), nil
|
|
}
|
|
return int64(t.count()), nil
|
|
}
|
|
|
|
// DumpPlain returns the entire list of added claims for a specific root hash.
|
|
// First return parametre are the indexes and second the values.
|
|
// If root is not specified, the last one is used.
|
|
func (t *Tree) DumpPlain(root []byte) ([][]byte, [][]byte, error) {
|
|
var indexes, values [][]byte
|
|
var err error
|
|
t.updateAccessTime()
|
|
|
|
t.iterateWithRoot(root, nil, func(k, v []byte) bool {
|
|
indexes = append(indexes, k)
|
|
values = append(values, v)
|
|
return false
|
|
})
|
|
|
|
return indexes, values, err
|
|
}
|
|
|
|
// ImportDump imports a partial or whole tree previously exported with Dump()
|
|
func (t *Tree) ImportDump(data []byte) error {
|
|
t.updateAccessTime()
|
|
if t.snapshotRoot != nil {
|
|
return fmt.Errorf("cannot import to a snapshot")
|
|
}
|
|
census := new(exportData)
|
|
bare.MaxArrayLength(bareMaxArrayLength)
|
|
bare.MaxUnmarshalBytes(bareMaxUnmarshalBytes)
|
|
if err := bare.Unmarshal(data, census); err != nil {
|
|
return fmt.Errorf("importdump cannot unmarshal data: %w", err)
|
|
}
|
|
keys := [][]byte{}
|
|
values := [][]byte{}
|
|
for _, ee := range census.Elements {
|
|
keys = append(keys, ee.Key)
|
|
values = append(values, ee.Value)
|
|
}
|
|
_, err := t.Tree.Update(keys, values)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
atomic.StoreUint64(&t.size, 0) // TBD: improve this
|
|
return t.Commit()
|
|
}
|
|
|
|
// Snapshot returns a Tree instance of a exiting merkle root.
|
|
// A Snapshot cannot be modified.
|
|
func (t *Tree) Snapshot(root []byte) (*Tree, error) {
|
|
exist, err := t.HashExists(root)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if !exist {
|
|
return nil, fmt.Errorf("root %x does not exist, cannot build snapshot", root)
|
|
}
|
|
return &Tree{Tree: t.Tree, public: t.public, snapshotRoot: root, snapshotSize: t.count()}, nil
|
|
}
|
|
|
|
func (t *Tree) Close() error {
|
|
t.db.Close()
|
|
return nil
|
|
}
|
|
|
|
// HashExists checks if a hash exists as a node in the merkle tree
|
|
func (t *Tree) HashExists(hash []byte) (bool, error) {
|
|
t.updateAccessTime()
|
|
return t.Tree.TrieRootExists(hash), nil
|
|
}
|
|
|
|
func (t *Tree) count() uint64 {
|
|
if v := atomic.LoadUint64(&t.size); v != 0 {
|
|
return v
|
|
}
|
|
counter := uint64(0)
|
|
if err := t.Tree.Walk(t.snapshotRoot, func(*asmt.WalkResult) int32 {
|
|
counter++
|
|
return 0
|
|
}); err != nil {
|
|
return 0
|
|
}
|
|
atomic.StoreUint64(&t.size, counter)
|
|
return counter
|
|
}
|
|
|
|
func (t *Tree) iterate(prefix []byte, callback func(key, value []byte) bool) {
|
|
t.Tree.Walk(t.snapshotRoot, func(v *asmt.WalkResult) int32 {
|
|
if callback(v.Key, v.Value) {
|
|
return 1
|
|
} else {
|
|
return 0
|
|
}
|
|
})
|
|
}
|
|
|
|
func (t *Tree) iterateWithRoot(root, prefix []byte, callback func(key, value []byte) bool) {
|
|
t.Tree.Walk(root, func(v *asmt.WalkResult) int32 {
|
|
if callback(v.Key, v.Value) {
|
|
return 1
|
|
} else {
|
|
return 0
|
|
}
|
|
})
|
|
}
|