1package merkle
2
3import (
4 "bytes"
5 "crypto/sha256"
6 "encoding/hex"
7 "errors"
8)
9
10type Hashable interface {
11 Bytes() []byte
12}
13
14type nodes []Node
15
16type Node struct {
17 hash []byte
18
19 position uint8
20}
21
22func NewNode(hash []byte, position uint8) Node {
23 return Node{
24 hash: hash,
25 position: position,
26 }
27}
28
29func (n Node) Position() uint8 {
30 return n.position
31}
32
33func (n Node) Hash() string {
34 return hex.EncodeToString(n.hash[:])
35}
36
37type Tree struct {
38 layers []nodes
39}
40
41// Root return the merkle root of the tree
42func (t *Tree) Root() string {
43 for _, l := range t.layers {
44 if len(l) == 1 {
45 return l[0].Hash()
46 }
47 }
48 return ""
49}
50
51// NewTree create a new Merkle Tree
52func NewTree(data []Hashable) *Tree {
53 tree := &Tree{}
54
55 leaves := make([]Node, len(data))
56
57 for i, d := range data {
58 hash := sha256.Sum256(d.Bytes())
59 leaves[i] = Node{hash: hash[:]}
60 }
61
62 tree.layers = []nodes{nodes(leaves)}
63
64 var buff bytes.Buffer
65 for len(leaves) > 1 {
66 level := make([]Node, 0, len(leaves)/2+1)
67 for i := 0; i < len(leaves); i += 2 {
68 buff.Reset()
69
70 if i < len(leaves)-1 {
71 buff.Write(leaves[i].hash)
72 buff.Write(leaves[i+1].hash)
73 hash := sha256.Sum256(buff.Bytes())
74 level = append(level, Node{
75 hash: hash[:],
76 })
77 } else {
78 level = append(level, leaves[i])
79 }
80 }
81 leaves = level
82 tree.layers = append(tree.layers, level)
83 }
84 return tree
85}
86
87// Proof return a MerkleProof
88func (t *Tree) Proof(data Hashable) ([]Node, error) {
89 targetHash := sha256.Sum256(data.Bytes())
90 targetIndex := -1
91
92 for i, layer := range t.layers[0] {
93 if bytes.Equal(targetHash[:], layer.hash) {
94 targetIndex = i
95 break
96 }
97 }
98
99 if targetIndex == -1 {
100 return nil, errors.New("target not found")
101 }
102
103 proofs := make([]Node, 0, len(t.layers))
104
105 for _, layer := range t.layers {
106 var pairIndex int
107
108 if targetIndex%2 == 0 {
109 pairIndex = targetIndex + 1
110 } else {
111 pairIndex = targetIndex - 1
112 }
113 if pairIndex < len(layer) {
114 proofs = append(proofs, Node{
115 hash: layer[pairIndex].hash,
116 position: uint8(targetIndex) % 2,
117 })
118 }
119 targetIndex /= 2
120 }
121 return proofs, nil
122}
123
124// Verify if a merkle proof is valid
125func (t *Tree) Verify(leaf Hashable, proofs []Node) bool {
126 return Verify(t.Root(), leaf, proofs)
127}
128
129// Verify if a merkle proof is valid
130func Verify(root string, leaf Hashable, proofs []Node) bool {
131 hash := sha256.Sum256(leaf.Bytes())
132
133 for i := 0; i < len(proofs); i += 1 {
134 var h []byte
135 if proofs[i].position == 0 {
136 h = append(hash[:], proofs[i].hash...)
137 } else {
138 h = append(proofs[i].hash, hash[:]...)
139 }
140 hash = sha256.Sum256(h)
141 }
142 return hex.EncodeToString(hash[:]) == root
143}
merkle.gno
2.58 Kb ยท 143 lines