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.

139 lines
3.1 KiB

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