go-dvote/tree/tree.go

package tree

import (
	"bytes"
	"errors"
	"fmt"
	"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, '\x00')
	}
	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 common3.HexEncode(t.Tree.RootKey().Bytes())
}

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(nil, func(n *merkletree.Node) {
		if n.Type == merkletree.NodeTypeLeaf {
			data := bytes.Trim(n.Value()[65:], "\x00")
			data = bytes.Replace(data, []byte("\u0000"), nil, -1)
			response = append(response, fmt.Sprintf("%s", data))
		}
	})
	return response, err
}

func (t *Tree) Snapshot(root string) (*Tree, error) {
	var rootHash merkletree.Hash
	snapshotTree := new(Tree)
	rootBytes, err := common3.HexDecode(root)
	if err != nil {
		return snapshotTree, err
	}
	copy(rootHash[:32], rootBytes)
	mt, err := t.Tree.Snapshot(&rootHash)
	snapshotTree.Tree = mt
	return snapshotTree, err
}