You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

150 lines
3.3 KiB

package tree
import (
"bytes"
"errors"
"os/user"
common3 "github.com/iden3/go-iden3/common"
mkcore "github.com/iden3/go-iden3/core"
db "github.com/iden3/go-iden3/db"
merkletree "github.com/iden3/go-iden3/merkletree"
)
type Tree struct {
Storage string
Tree *merkletree.MerkleTree
DbStorage *db.LevelDbStorage
}
func (t *Tree) Init(namespace string) error {
if len(t.Storage) < 1 {
if len(namespace) < 1 {
return errors.New("namespace not valid")
}
usr, err := user.Current()
if err == nil {
t.Storage = usr.HomeDir + "/.dvote/census/" + namespace
} else {
t.Storage = "./dvoteTree/" + namespace
}
}
mtdb, err := db.NewLevelDbStorage(t.Storage, false)
if err != nil {
return err
}
mt, err := merkletree.NewMerkleTree(mtdb, 140)
if err != nil {
return err
}
t.DbStorage = mtdb
t.Tree = mt
return nil
}
func (t *Tree) Close() {
defer t.Tree.Storage().Close()
}
func (t *Tree) GetClaim(data []byte) (*mkcore.ClaimBasic, error) {
if len(data) > 496/8 {
return nil, errors.New("claim data too large")
}
for i := len(data); i <= 496/8; i++ {
data = append(data, byte('.'))
}
var indexSlot [400 / 8]byte
var dataSlot [496 / 8]byte
copy(indexSlot[:], data[:400/8])
copy(dataSlot[:], data[:496/8])
e := mkcore.NewClaimBasic(indexSlot, dataSlot)
return e, nil
}
func (t *Tree) AddClaim(data []byte) error {
e, err := t.GetClaim(data)
if err != nil {
return err
}
return t.Tree.Add(e.Entry())
}
func (t *Tree) GenProof(data []byte) (string, error) {
e, err := t.GetClaim(data)
if err != nil {
return "", err
}
mp, err := t.Tree.GenerateProof(e.Entry().HIndex())
if err != nil {
return "", err
}
mpHex := common3.HexEncode(mp.Bytes())
return mpHex, nil
}
func (t *Tree) CheckProof(data []byte, mpHex string) (bool, error) {
mpBytes, err := common3.HexDecode(mpHex)
if err != nil {
return false, err
}
mp, err := merkletree.NewProofFromBytes(mpBytes)
if err != nil {
return false, err
}
e, err := t.GetClaim(data)
if err != nil {
return false, err
}
return merkletree.VerifyProof(t.Tree.RootKey(), mp,
e.Entry().HIndex(), e.Entry().HValue()), nil
}
func (t *Tree) GetRoot() string {
return t.Tree.RootKey().String()
}
func (t *Tree) GetIndex(data []byte) (string, error) {
e, err := t.GetClaim(data)
if err != nil {
return "", err
}
index, err := t.Tree.GetDataByIndex(e.Entry().HIndex())
return index.String(), err
}
/*
func (t *Tree) Dump() ([]string, error) {
var response []string
err := t.Tree.Walk(t.Tree.RootKey(), func(n *merkletree.Node) {
if n.Type == merkletree.NodeTypeLeaf {
rawValue := n.Value()
var cleanValue []byte
for i := 0; i < len(rawValue); i++ {
if rawValue[i] == byte('.') {
break
}
cleanValue = append(cleanValue, rawValue[i])
}
response = append(response, fmt.Sprintf("%s", cleanValue))
}
})
return response, err
}
*/
func (t *Tree) Dump() (string, error) {
w := bytes.NewBufferString("")
err := t.Tree.DumpClaims(w, nil) // as rootKey we can pass a nil pointer, and it will use the current RootKey
return w.String(), err
}
func (t *Tree) Snapshot(root string) (*Tree, error) {
var rootHash merkletree.Hash
copy(rootHash[:32], root)
mt, err := t.Tree.Snapshot(&rootHash)
snapshotTree := new(Tree)
snapshotTree.Tree = mt
return snapshotTree, err
}