-
Notifications
You must be signed in to change notification settings - Fork 37
/
merkletree.go
316 lines (275 loc) · 10.5 KB
/
merkletree.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
// Copyright © 2018 - 2023 Weald Technology Trading.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Package merkletree is an implementation of a Merkle tree (https://en.wikipedia.org/wiki/Merkle_tree). It provides methods to
// create a tree and generate and verify proofs. The hashing algorithm for the tree is selectable between BLAKE2b and Keccak256,
// or you can supply your own.
//
// This implementation includes advanced features salting and pollarding. Salting is the act of adding a piece of data to each
// value in the Merkle tree as it is initially hashed to form the leaves, which helps avoid rainbow table attacks on leaf hashes
// presented as part of proofs. Pollarding is the act of providing the root plus all branches to a certain height which can be
// used to reduce the size of proofs. This is useful when multiple proofs are presented against the same tree as it can reduce the
// overall size.
//
// Creating a Merkle tree requires a list of values that are each byte arrays. Once a tree has been created proofs can be generated
// using the tree's GenerateProof() function.
//
// The package includes a function VerifyProof() to verify a generated proof given only the data to prove, proof and the pollard of
// the relevant Merkle tree. This allows for efficient verification of proofs without requiring the entire Merkle tree to be stored
// or recreated.
//
// # Implementation notes
//
// The tree pads its values to the next highest power of 2; values not supplied are treated as null with a value hash of 0. This
// can be seen graphically by generating a DOT representation of the graph with DOT().
//
// If salting is enabled it appends an 4-byte value to each piece of data. The value is the binary representation of the index in
// big-endian form. Note that if there are more than 2^32 values in the tree the salt will wrap, being modulo 2^32
package merkletree
import (
"bytes"
"encoding/binary"
"encoding/hex"
"math"
"sort"
"github.com/pkg/errors"
)
// MerkleTree is the structure for the Merkle tree.
type MerkleTree struct {
// if Salt is true the Data values are salted with their index
Salt bool `json:"salt"`
// if Sorted is true, the Hash values are Sorted before hashing branch Nodes
Sorted bool `json:"sorted"`
// Hash is a pointer to the hashing struct
Hash HashType `json:"hash_type"`
// Data is the Data from which the Merkle tree is created
Data [][]byte `json:"data"`
// Nodes are the leaf and branch Nodes of the Merkle tree
Nodes [][]byte `json:"nodes"`
}
// A container which gives us the ability to sort the hashes by value
// while maintaining the relative positions of the data and it's hash.
type hashSorter struct {
// data from which the hashes are generated
data [][]byte
// Hashes of the data
hashes [][]byte
}
// Len length of the data slice.
func (s hashSorter) Len() int {
return len(s.data)
}
// Swap the given indicies in both the data and leaf slices.
func (s hashSorter) Swap(i, j int) {
s.data[i], s.data[j] = s.data[j], s.data[i]
s.hashes[i], s.hashes[j] = s.hashes[j], s.hashes[i]
}
// Compares the hash indicies, returns true if i is less than j.
func (s hashSorter) Less(i, j int) bool {
return bytes.Compare(s.hashes[i], s.hashes[j]) == -1
}
// Index of the data in the MerkleTree.
func (t *MerkleTree) indexOf(input []byte) (uint64, error) {
for i, data := range t.Data {
if bytes.Equal(data, input) {
return uint64(i), nil
}
}
return 0, errors.New("data not found")
}
// GenerateProof generates the proof for a piece of data.
// Height is the height of the pollard to verify the proof. If using the Merkle root to verify this should be 0.
// If the data is not present in the tree this will return an error.
// If the data is present in the tree this will return the hashes for each level in the tree and the index of the value in the tree.
func (t *MerkleTree) GenerateProof(data []byte, height int) (*Proof, error) {
// Find the index of the data
index, err := t.indexOf(data)
if err != nil {
return nil, err
}
return t.GenerateProofWithIndex(index, height)
}
// GenerateProofWithIndex generates the proof for the data at the given index.
// It is faster than GenerateProof() if the index is already known.
// Height is the height of the pollard to verify the proof. If using the Merkle root to verify this should be 0.
// If the index is out of range this will return an error.
// If the data is present in the tree this will return the hashes for each level in the tree and the index of the value in the tree.
func (t *MerkleTree) GenerateProofWithIndex(index uint64, height int) (*Proof, error) {
if index >= uint64(len(t.Data)) {
return nil, errors.New("index out of range")
}
proofLen := int(math.Ceil(math.Log2(float64(len(t.Data))))) - height
hashes := make([][]byte, proofLen)
cur := 0
minI := uint64(math.Pow(2, float64(height+1))) - 1
for i := index + uint64(len(t.Nodes)/2); i > minI; i /= 2 {
hashes[cur] = t.Nodes[i^1]
cur++
}
return newProof(hashes, index), nil
}
// GenerateMultiProof generates the proof for multiple pieces of data.
func (t *MerkleTree) GenerateMultiProof(data [][]byte) (*MultiProof, error) {
indices := make([]uint64, len(data))
// Step 1: generate individual proofs.
for dataIndex := range data {
index, err := t.indexOf(data[dataIndex])
if err != nil {
return nil, err
}
indices[dataIndex] = index
}
return t.GenerateMultiProofWithIndices(indices)
}
// GenerateMultiProofWithIndices generates the proof for multiple pieces of data.
func (t *MerkleTree) GenerateMultiProofWithIndices(indices []uint64) (*MultiProof, error) {
hashes := make([][][]byte, len(indices))
// Step 1: generate individual proofs.
for i, leafIndex := range indices {
tmpProof, err := t.GenerateProofWithIndex(leafIndex, 0)
if err != nil {
return nil, err
}
hashes[i] = tmpProof.Hashes
}
// Step 2: combine the hashes across all proofs and highlight all calculated indices.
proofHashes := make(map[uint64][]byte)
calculatedIndices := make([]bool, len(t.Nodes))
for i, index := range indices {
hashNum := 0
for j := index + uint64(math.Ceil(float64(len(t.Nodes))/2)); j > 1; j /= 2 {
proofHashes[j^1] = hashes[i][hashNum]
calculatedIndices[j] = true
hashNum++
}
}
// Step 3: remove any hashes that can be calculated.
for _, index := range indices {
for j := index + uint64(math.Ceil(float64(len(t.Nodes))/2)); j > 1; j /= 2 {
if calculatedIndices[j^1] {
delete(proofHashes, j^1)
}
}
}
return NewMultiProof(
WithHashes(proofHashes),
WithSalt(t.Salt),
WithSorted(t.Sorted),
WithHashType(t.Hash),
WithIndices(indices),
WithValues(uint64(len(t.Nodes)/2)),
)
}
// NewTree creates a new merkle tree using the provided information.
func NewTree(params ...Parameter) (*MerkleTree, error) {
parameters, err := parseAndCheckTreeParameters(params...)
if err != nil {
return nil, errors.Wrap(err, "problem with parameters")
}
branchesLen := int(math.Exp2(math.Ceil(math.Log2(float64(len(parameters.data))))))
// We pad our data length up to the power of 2.
nodes := make([][]byte, branchesLen*2)
// We put the leaves after the branches in the slice of nodes.
createLeaves(
parameters.data,
nodes[branchesLen:branchesLen+len(parameters.data)],
parameters.hash,
parameters.salt,
parameters.sorted,
)
// Pad the space left after the leaves.
for i := len(parameters.data) + branchesLen; i < len(nodes); i++ {
nodes[i] = make([]byte, parameters.hash.HashLength())
}
// Branches.
createBranches(
nodes,
parameters.hash,
branchesLen,
parameters.sorted,
)
tree := &MerkleTree{
Salt: parameters.salt,
Sorted: parameters.sorted,
Hash: parameters.hash,
Nodes: nodes,
Data: parameters.data,
}
return tree, nil
}
// New creates a new Merkle tree using the provided raw data and default hash type. Salting is not used.
// data must contain at least one element for it to be valid.
// Deprecated: plase use NewTree().
func New(data [][]byte) (*MerkleTree, error) {
return NewTree(WithData(data))
}
// Hashes the data slice, placing the result hashes into dest.
// salt adds a salt to the hash using the index.
// sorted sorts the leaves and data by the value of the leaf hash.
func createLeaves(data [][]byte, dest [][]byte, hash HashType, salt, sorted bool) {
indexSalt := make([]byte, 4)
for i := range data {
if salt {
binary.BigEndian.PutUint32(indexSalt, uint32(i))
dest[i] = hash.Hash(data[i], indexSalt)
} else {
dest[i] = hash.Hash(data[i])
}
}
if sorted {
sorter := hashSorter{
data: data,
hashes: dest,
}
sort.Sort(sorter)
}
}
// Create the branch nodes from the existing leaf data.
func createBranches(nodes [][]byte, hash HashType, leafOffset int, sorted bool) {
for leafIndex := leafOffset - 1; leafIndex > 0; leafIndex-- {
left := nodes[leafIndex*2]
right := nodes[leafIndex*2+1]
if sorted && bytes.Compare(left, right) == 1 {
nodes[leafIndex] = hash.Hash(right, left)
} else {
nodes[leafIndex] = hash.Hash(left, right)
}
}
}
// NewUsing creates a new Merkle tree using the provided raw data and supplied hash type.
// Salting is used, and hashes are sorted if requested.
// data must contain at least one element for it to be valid.
// Deprecated: plase use NewTree().
func NewUsing(data [][]byte, hash HashType, salt bool) (*MerkleTree, error) {
return NewTree(
WithData(data),
WithHashType(hash),
WithSalt(salt),
)
}
// Pollard returns the Merkle root plus branches to a certain height. Height 0 will return just the root, height 1 the root plus
// the two branches directly above it, height 2 the root, two branches directly above it and four branches directly above them, etc.
func (t *MerkleTree) Pollard(height int) [][]byte {
return t.Nodes[1:int(math.Exp2(float64(height+1)))]
}
// Root returns the Merkle root (hash of the root node) of the tree.
func (t *MerkleTree) Root() []byte {
return t.Nodes[1]
}
// GetSalt returns the true if the values in this Merkle tree are salted.
func (t *MerkleTree) GetSalt() bool {
return t.Salt
}
// String implements the stringer interface.
func (t *MerkleTree) String() string {
return hex.EncodeToString(t.Nodes[1])
}