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