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}