arnaucube
arnaucube/go-dvotePublic Self-hosted repository
go-dvote/tree/tree.go
tree/tree.goRaw
  1package tree
  2
  3import (
  4	"errors"
  5	"fmt"
  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	Namespace string
 16	Storage   string
 17	Tree      *merkletree.MerkleTree
 18	DbStorage *db.LevelDbStorage
 19}
 20
 21func (t *Tree) Init() error {
 22	if len(t.Storage) < 1 {
 23		usr, err := user.Current()
 24		if err == nil {
 25			t.Storage = usr.HomeDir + "/.dvote/Tree"
 26		} else {
 27			t.Storage = "./dvoteTree"
 28		}
 29	}
 30	mtdb, err := db.NewLevelDbStorage(t.Storage, false)
 31	if err != nil {
 32		return err
 33	}
 34	mt, err := merkletree.NewMerkleTree(mtdb, 140)
 35	if err != nil {
 36		return err
 37	}
 38	t.DbStorage = mtdb
 39	t.Tree = mt
 40	return nil
 41}
 42
 43func (t *Tree) Close() {
 44	defer t.Tree.Storage().Close()
 45}
 46
 47func (t *Tree) GetClaim(data []byte) (*mkcore.ClaimBasic, error) {
 48	if len(data) > 496/8 {
 49		return nil, errors.New("claim data too large")
 50	}
 51	for i := len(data); i <= 496/8; i++ {
 52		data = append(data, byte('0'))
 53	}
 54	var indexSlot [400 / 8]byte
 55	var dataSlot [496 / 8]byte
 56	copy(indexSlot[:], data[:400/8])
 57	copy(dataSlot[:], data[:496/8])
 58	e := mkcore.NewClaimBasic(indexSlot, dataSlot)
 59	return e, nil
 60}
 61
 62func (t *Tree) AddClaim(data []byte) error {
 63	e, err := t.GetClaim(data)
 64	if err != nil {
 65		return err
 66	}
 67	return t.Tree.Add(e.Entry())
 68}
 69
 70func (t *Tree) GenProof(data []byte) (string, error) {
 71	e, err := t.GetClaim(data)
 72	if err != nil {
 73		return "", err
 74	}
 75	mp, err := t.Tree.GenerateProof(e.Entry().HIndex())
 76	if err != nil {
 77		return "", err
 78	}
 79	mpHex := common3.HexEncode(mp.Bytes())
 80	return mpHex, nil
 81}
 82
 83func (t *Tree) CheckProof(data []byte, mpHex string) (bool, error) {
 84	mpBytes, err := common3.HexDecode(mpHex)
 85	if err != nil {
 86		return false, err
 87	}
 88	mp, err := merkletree.NewProofFromBytes(mpBytes)
 89	if err != nil {
 90		return false, err
 91	}
 92	e, err := t.GetClaim(data)
 93	if err != nil {
 94		return false, err
 95	}
 96	return merkletree.VerifyProof(t.Tree.RootKey(), mp,
 97		e.Entry().HIndex(), e.Entry().HValue()), nil
 98}
 99
100func (t *Tree) GetRoot() string {
101	return t.Tree.RootKey().String()
102}
103
104func (t *Tree) Dump() ([]string, error) {
105	var response []string
106
107	err := t.Tree.Walk(t.Tree.RootKey(), func(n *merkletree.Node) {
108		if n.Type == merkletree.NodeTypeLeaf {
109			response = append(response, fmt.Sprintf("|%s", n.Entry.Data))
110		}
111	})
112	return response, err
113}
114
115func (t *Tree) Snapshot() (string, error) {
116	snapshotNamespace := t.Tree.RootKey().String()
117	return snapshotNamespace, nil
118}