1package tree
  2
  3import (
  4	"bytes"
  5	"errors"
  6	"fmt"
  7	"os/user"
  8
  9	common3 "github.com/iden3/go-iden3/common"
 10	mkcore "github.com/iden3/go-iden3/core"
 11	db "github.com/iden3/go-iden3/db"
 12	merkletree "github.com/iden3/go-iden3/merkletree"
 13)
 14
 15type Tree struct {
 16	Storage   string
 17	Tree      *merkletree.MerkleTree
 18	DbStorage *db.LevelDbStorage
 19}
 20
 21func (t *Tree) Init(namespace string) error {
 22	if len(t.Storage) < 1 {
 23		if len(namespace) < 1 {
 24			return errors.New("namespace not valid")
 25		}
 26		usr, err := user.Current()
 27		if err == nil {
 28			t.Storage = usr.HomeDir + "/.dvote/census/" + namespace
 29		} else {
 30			t.Storage = "./dvoteTree/" + namespace
 31		}
 32	}
 33	mtdb, err := db.NewLevelDbStorage(t.Storage, false)
 34	if err != nil {
 35		return err
 36	}
 37	mt, err := merkletree.NewMerkleTree(mtdb, 140)
 38	if err != nil {
 39		return err
 40	}
 41	t.DbStorage = mtdb
 42	t.Tree = mt
 43	return nil
 44}
 45
 46func (t *Tree) Close() {
 47	defer t.Tree.Storage().Close()
 48}
 49
 50func (t *Tree) GetClaim(data []byte) (*mkcore.ClaimBasic, error) {
 51	if len(data) > 496/8 {
 52		return nil, errors.New("claim data too large")
 53	}
 54	for i := len(data); i <= 496/8; i++ {
 55		data = append(data, '\x00')
 56	}
 57	var indexSlot [400 / 8]byte
 58	var dataSlot [496 / 8]byte
 59	copy(indexSlot[:], data[:400/8])
 60	copy(dataSlot[:], data[:496/8])
 61	e := mkcore.NewClaimBasic(indexSlot, dataSlot)
 62	return e, nil
 63}
 64
 65func (t *Tree) AddClaim(data []byte) error {
 66	e, err := t.GetClaim(data)
 67	if err != nil {
 68		return err
 69	}
 70	return t.Tree.Add(e.Entry())
 71}
 72
 73func (t *Tree) GenProof(data []byte) (string, error) {
 74	e, err := t.GetClaim(data)
 75	if err != nil {
 76		return "", err
 77	}
 78	mp, err := t.Tree.GenerateProof(e.Entry().HIndex())
 79	if err != nil {
 80		return "", err
 81	}
 82	mpHex := common3.HexEncode(mp.Bytes())
 83	return mpHex, nil
 84}
 85
 86func (t *Tree) CheckProof(data []byte, mpHex string) (bool, error) {
 87	mpBytes, err := common3.HexDecode(mpHex)
 88	if err != nil {
 89		return false, err
 90	}
 91	mp, err := merkletree.NewProofFromBytes(mpBytes)
 92	if err != nil {
 93		return false, err
 94	}
 95	e, err := t.GetClaim(data)
 96	if err != nil {
 97		return false, err
 98	}
 99	return merkletree.VerifyProof(t.Tree.RootKey(), mp,
100		e.Entry().HIndex(), e.Entry().HValue()), nil
101}
102
103func (t *Tree) GetRoot() string {
104	return common3.HexEncode(t.Tree.RootKey().Bytes())
105}
106
107func (t *Tree) GetIndex(data []byte) (string, error) {
108	e, err := t.GetClaim(data)
109	if err != nil {
110		return "", err
111	}
112	index, err := t.Tree.GetDataByIndex(e.Entry().HIndex())
113	return index.String(), err
114}
115
116func (t *Tree) Dump() ([]string, error) {
117	var response []string
118	err := t.Tree.Walk(nil, func(n *merkletree.Node) {
119		if n.Type == merkletree.NodeTypeLeaf {
120			data := bytes.Trim(n.Value()[65:], "\x00")
121			response = append(response, fmt.Sprintf("%s", data))
122		}
123	})
124	return response, err
125}
126
127func (t *Tree) Snapshot(root string) (*Tree, error) {
128	var rootHash merkletree.Hash
129	snapshotTree := new(Tree)
130	rootBytes, err := common3.HexDecode(root)
131	if err != nil {
132		return snapshotTree, err
133	}
134	copy(rootHash[:32], rootBytes)
135	mt, err := t.Tree.Snapshot(&rootHash)
136	snapshotTree.Tree = mt
137	return snapshotTree, err
138}