1// Package arbotree provides the functions for creating and managing an arbo
  2// merkletree adapted to the CensusTree interface
  3package arbotree
  4
  5import (
  6	"bytes"
  7	"encoding/hex"
  8	"fmt"
  9	"path/filepath"
 10	"strings"
 11	"sync/atomic"
 12	"time"
 13
 14	"github.com/arnaucube/arbo"
 15	"github.com/iden3/go-merkletree/db/pebble"
 16	"go.vocdoni.io/dvote/censustree"
 17)
 18
 19type Tree struct {
 20	Tree           *arbo.Tree
 21	public         uint32
 22	lastAccessUnix int64 // a unix timestamp, used via sync/atomic
 23}
 24
 25func (t *Tree) Init(name, storageDir string) error {
 26	dbDir := filepath.Join(storageDir, "arbotree.db."+strings.TrimSpace(name))
 27	storage, err := pebble.NewPebbleStorage(dbDir, false) // TODO TMP
 28	if err != nil {
 29		return err
 30	}
 31
 32	mt, err := arbo.NewTree(storage, 140, arbo.HashFunctionBlake2b) // TODO here the hash function would depend on the usage
 33	if err != nil {
 34		return err
 35	}
 36	t.Tree = mt
 37	return nil
 38}
 39
 40func (t *Tree) MaxKeySize() int {
 41	return t.Tree.HashFunction().Len()
 42}
 43
 44func (t *Tree) LastAccess() int64 {
 45	return atomic.LoadInt64(&t.lastAccessUnix)
 46}
 47
 48func (t *Tree) updateAccessTime() {
 49	atomic.StoreInt64(&t.lastAccessUnix, time.Now().Unix())
 50}
 51
 52// Publish makes a merkle tree available for queries.
 53// Application layer should check IsPublish() before considering the Tree available.
 54func (t *Tree) Publish() {
 55	atomic.StoreUint32(&t.public, 1)
 56}
 57
 58// UnPublish makes a merkle tree not available for queries
 59func (t *Tree) UnPublish() {
 60	atomic.StoreUint32(&t.public, 0)
 61}
 62
 63// IsPublic returns true if the tree is available
 64func (t *Tree) IsPublic() bool {
 65	return atomic.LoadUint32(&t.public) == 1
 66}
 67
 68func (t *Tree) Add(index, value []byte) error {
 69	t.updateAccessTime()
 70	return t.Tree.Add(index, value)
 71}
 72
 73func (t *Tree) AddBatch(indexes, values [][]byte) ([]int, error) {
 74	t.updateAccessTime()
 75	return t.Tree.AddBatch(indexes, values)
 76}
 77
 78func (t *Tree) GenProof(index, value []byte) ([]byte, error) {
 79	t.updateAccessTime()
 80	v, siblings, err := t.Tree.GenProof(index)
 81	if err != nil {
 82		return nil, err
 83	}
 84	if !bytes.Equal(v, value) {
 85		return nil, fmt.Errorf("value does not match %s!=%s", hex.EncodeToString(v), hex.EncodeToString(value))
 86	}
 87	return siblings, nil
 88}
 89
 90func (t *Tree) CheckProof(index, value, root, mproof []byte) (bool, error) {
 91	t.updateAccessTime()
 92	if root == nil {
 93		root = t.Root()
 94	}
 95	return arbo.CheckProof(t.Tree.HashFunction(), index, value, root, mproof)
 96}
 97
 98func (t *Tree) Root() []byte {
 99	t.updateAccessTime()
100	return t.Tree.Root()
101}
102
103func (t *Tree) Dump(root []byte) ([]byte, error) {
104	return t.Tree.Dump() // TODO pass root once arbo is updated
105}
106
107func (t *Tree) DumpPlain(root []byte) ([][]byte, [][]byte, error) {
108	t.updateAccessTime()
109	var indexes, values [][]byte
110	// TODO pass root once arbo is updated
111	err := t.Tree.Iterate(func(k, v []byte) {
112		if v[0] != arbo.PrefixValueLeaf {
113			return
114		}
115		leafK, leafV := arbo.ReadLeafValue(v)
116		indexes = append(indexes, leafK)
117		values = append(values, leafV)
118	})
119	return indexes, values, err
120}
121
122func (t *Tree) ImportDump(data []byte) error {
123	t.updateAccessTime()
124	return t.Tree.ImportDump(data)
125}
126
127func (t *Tree) Size(root []byte) (int64, error) {
128	// TODO
129	return 0, nil
130}
131
132func (t *Tree) Snapshot(root []byte) (censustree.Tree, error) {
133	// TODO
134	return t, nil
135}
136func (t *Tree) HashExists(hash []byte) (bool, error) {
137	_, _, err := t.Tree.Get(hash)
138	if err != nil {
139		return false, err
140	}
141	return true, nil
142}