go-merkletree-old/merkletree.go
merkletree.goRaw
1package merkletree
2
3import (
4 "bytes"
5 "errors"
6
7 "github.com/syndtr/goleveldb/leveldb"
8)
9
10const (
11 // EmptyNodeType indicates the type of an EmptyNodeValue Node
12 EmptyNodeType = 00
13 // NormalNodeType indicates the type of a middle Node
14 normalNodeType = 01
15 // FinalNodeType indicates the type of middle Node that is in an optimized branch, then in the value contains the value of the final leaf node of that branch
16 finalNodeType = 02
17 // ValueNodeType indicates the type of a value Node
18 valueNodeType = 03
19 // RootNodeType indicates the type of a root Node
20 rootNodeType = 04
21)
22
23var (
24 // ErrNodeAlreadyExists is an error that indicates that a node already exists in the merkletree database
25 ErrNodeAlreadyExists = errors.New("node already exists")
26 rootNodeValue = HashBytes([]byte("root"))
27 // EmptyNodeValue is a [32]byte EmptyNodeValue array, all to zero
28 EmptyNodeValue = Hash{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
29)
30
31// Hash used in this tree, is the [32]byte keccak()
32type Hash [32]byte
33
34// Value is the interface of a generic leaf, a key value object stored in the leveldb
35type Value interface {
36 IndexLength() uint32 // returns the index length value
37 Bytes() []byte // returns the value in byte array representation
38}
39
40//MerkleTree struct with the main elements of the Merkle Tree
41type MerkleTree struct {
42 // sync.RWMutex
43 storage *leveldb.DB
44 root Hash
45 numLevels int // Height of the Merkle Tree, number of levels
46}
47
48// New generates a new Merkle Tree
49func New(storage *leveldb.DB, numLevels int) (*MerkleTree, error) {
50 var mt MerkleTree
51 mt.storage = storage
52 mt.numLevels = numLevels
53 var err error
54 _, _, rootHash, err := mt.Get(rootNodeValue)
55 if err != nil {
56 mt.root = EmptyNodeValue
57 err = mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
58 if err != nil {
59 return nil, err
60 }
61 }
62 copy(mt.root[:], rootHash)
63 return &mt, nil
64}
65
66// Root returns the merkletree.Root
67func (mt *MerkleTree) Root() Hash {
68 return mt.root
69}
70
71// NumLevels returns the merkletree.NumLevels
72func (mt *MerkleTree) NumLevels() int {
73 return mt.numLevels
74}
75
76// Add adds the leaf to the MT
77func (mt *MerkleTree) Add(v Value) error {
78 // add the leaf that we are adding
79 mt.Insert(HashBytes(v.Bytes()), valueNodeType, v.IndexLength(), v.Bytes())
80
81 hi := HashBytes(v.Bytes()[:v.IndexLength()])
82 path := getPath(mt.numLevels, hi)
83
84 nodeHash := mt.root
85 var siblings []Hash
86 for i := mt.numLevels - 2; i >= 0; i-- {
87 nodeType, indexLength, nodeBytes, err := mt.Get(nodeHash)
88 if err != nil {
89 return err
90 }
91 if nodeType == byte(finalNodeType) {
92 hiChild := HashBytes(nodeBytes[:indexLength])
93 pathChild := getPath(mt.numLevels, hiChild)
94 posDiff := comparePaths(pathChild, path)
95 if posDiff == -1 {
96 return ErrNodeAlreadyExists
97 }
98 finalNode1Hash := calcHashFromLeafAndLevel(posDiff, pathChild, HashBytes(nodeBytes))
99 mt.Insert(finalNode1Hash, finalNodeType, indexLength, nodeBytes)
100 finalNode2Hash := calcHashFromLeafAndLevel(posDiff, path, HashBytes(v.Bytes()))
101 mt.Insert(finalNode2Hash, finalNodeType, v.IndexLength(), v.Bytes())
102 // now the parent
103 var parentNode treeNode
104 if path[posDiff] {
105 parentNode = treeNode{
106 ChildL: finalNode1Hash,
107 ChildR: finalNode2Hash,
108 }
109 } else {
110 parentNode = treeNode{
111 ChildL: finalNode2Hash,
112 ChildR: finalNode1Hash,
113 }
114 }
115 siblings = append(siblings, getEmptiesBetweenIAndPosHash(mt, i, posDiff+1)...)
116
117 if mt.root, err = mt.replaceLeaf(siblings, path[posDiff+1:], parentNode.Ht(), normalNodeType, 0, parentNode.Bytes()); err != nil {
118 return err
119 }
120 mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
121 return nil
122 }
123 node := parseNodeBytes(nodeBytes)
124 var sibling Hash
125 if !path[i] {
126 nodeHash = node.ChildL
127 sibling = node.ChildR
128 } else {
129 nodeHash = node.ChildR
130 sibling = node.ChildL
131 }
132 siblings = append(siblings, sibling)
133
134 if bytes.Equal(nodeHash[:], EmptyNodeValue[:]) {
135 // if the node is EmptyNodeValue, the leaf data will go directly at that height, as a Final Node
136 if i == mt.numLevels-2 && bytes.Equal(siblings[len(siblings)-1][:], EmptyNodeValue[:]) {
137 // if the pt node is the unique in the tree, just put it into the root node
138 // this means to be in i==mt.NumLevels-2 && nodeHash==EmptyNodeValue
139 finalNodeHash := calcHashFromLeafAndLevel(i+1, path, HashBytes(v.Bytes()))
140 mt.Insert(finalNodeHash, finalNodeType, v.IndexLength(), v.Bytes())
141 mt.root = finalNodeHash
142 mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
143 return nil
144 }
145 finalNodeHash := calcHashFromLeafAndLevel(i, path, HashBytes(v.Bytes()))
146 if mt.root, err = mt.replaceLeaf(siblings, path[i:], finalNodeHash, finalNodeType, v.IndexLength(), v.Bytes()); err != nil {
147 return err
148 }
149 mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
150 return nil
151 }
152 }
153
154 var err error
155 mt.root, err = mt.replaceLeaf(siblings, path, HashBytes(v.Bytes()), valueNodeType, v.IndexLength(), v.Bytes())
156 if err != nil {
157 return err
158 }
159 mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
160 return nil
161}
162
163// GenerateProof generates the Merkle Proof from a given leafHash for the current root
164func (mt *MerkleTree) GenerateProof(hi Hash) ([]byte, error) {
165 var empties [32]byte
166
167 path := getPath(mt.numLevels, hi)
168 var siblings []Hash
169 nodeHash := mt.root
170
171 for level := 0; level < mt.numLevels-1; level++ {
172 nodeType, indexLength, nodeBytes, err := mt.Get(nodeHash)
173 if err != nil {
174 return nil, err
175 }
176 if nodeType == byte(finalNodeType) {
177 realValueInPos, err := mt.GetValueInPos(hi)
178 if err != nil {
179 return nil, err
180 }
181 if bytes.Equal(realValueInPos[:], EmptyNodeValue[:]) {
182 // go until the path is different, then get the nodes between this FinalNode and the node in the diffPath, they will be the siblings of the merkle proof
183 leafHi := HashBytes(nodeBytes[:indexLength]) // hi of element that was in the end of the branch (the finalNode)
184 pathChild := getPath(mt.numLevels, leafHi)
185
186 // get the position where the path is different
187 posDiff := comparePaths(pathChild, path)
188 if posDiff == -1 {
189 return nil, ErrNodeAlreadyExists
190 }
191
192 if posDiff != mt.NumLevels()-1-level {
193 sibling := calcHashFromLeafAndLevel(posDiff, pathChild, HashBytes(nodeBytes))
194 setbitmap(empties[:], uint(mt.NumLevels()-2-posDiff))
195 siblings = append([]Hash{sibling}, siblings...)
196 }
197
198 }
199 break
200 }
201 node := parseNodeBytes(nodeBytes)
202
203 var sibling Hash
204 if !path[mt.numLevels-level-2] {
205 nodeHash = node.ChildL
206 sibling = node.ChildR
207 } else {
208 nodeHash = node.ChildR
209 sibling = node.ChildL
210 }
211 if !bytes.Equal(sibling[:], EmptyNodeValue[:]) {
212 setbitmap(empties[:], uint(level))
213 siblings = append([]Hash{sibling}, siblings...)
214 }
215 }
216 // merge empties and siblings
217 var mp []byte
218 mp = append(mp, empties[:]...)
219 for k := range siblings {
220 mp = append(mp, siblings[k][:]...)
221 }
222 return mp, nil
223}
224
225// GetValueInPos returns the merkletree value in the position of the Hash of the Index (Hi)
226func (mt *MerkleTree) GetValueInPos(hi Hash) ([]byte, error) {
227 path := getPath(mt.numLevels, hi)
228 nodeHash := mt.root
229 for i := mt.numLevels - 2; i >= 0; i-- {
230 nodeType, indexLength, nodeBytes, err := mt.Get(nodeHash)
231 if err != nil {
232 return nodeBytes, err
233 }
234 if nodeType == byte(finalNodeType) {
235 // check if nodeBytes path is different of hi
236 index := nodeBytes[:indexLength]
237 nodeHi := HashBytes(index)
238 nodePath := getPath(mt.numLevels, nodeHi)
239 posDiff := comparePaths(path, nodePath)
240 // if is different, return an EmptyNodeValue, else return the nodeBytes
241 if posDiff != -1 {
242 return EmptyNodeValue[:], nil
243 }
244 return nodeBytes, nil
245 }
246 node := parseNodeBytes(nodeBytes)
247 if !path[i] {
248 nodeHash = node.ChildL
249 } else {
250 nodeHash = node.ChildR
251 }
252 }
253 _, _, valueBytes, err := mt.Get(nodeHash)
254 if err != nil {
255 return valueBytes, err
256 }
257 return valueBytes, nil
258}
259
260func calcHashFromLeafAndLevel(untilLevel int, path []bool, leafHash Hash) Hash {
261 nodeCurrLevel := leafHash
262 for i := 0; i < untilLevel; i++ {
263 if path[i] {
264 node := treeNode{
265 ChildL: EmptyNodeValue,
266 ChildR: nodeCurrLevel,
267 }
268 nodeCurrLevel = node.Ht()
269 } else {
270 node := treeNode{
271 ChildL: nodeCurrLevel,
272 ChildR: EmptyNodeValue,
273 }
274 nodeCurrLevel = node.Ht()
275 }
276 }
277 return nodeCurrLevel
278}
279
280func (mt *MerkleTree) replaceLeaf(siblings []Hash, path []bool, newLeafHash Hash, nodetype byte, indexLength uint32, newLeafValue []byte) (Hash, error) {
281 // add the new leaf
282 mt.Insert(newLeafHash, nodetype, indexLength, newLeafValue)
283 currNode := newLeafHash
284 // here the path is only the path[posDiff+1]
285 for i := 0; i < len(siblings); i++ {
286 if !path[i] {
287 node := treeNode{
288 ChildL: currNode,
289 ChildR: siblings[len(siblings)-1-i],
290 }
291 mt.Insert(node.Ht(), normalNodeType, 0, node.Bytes())
292 currNode = node.Ht()
293 } else {
294
295 node := treeNode{
296 ChildL: siblings[len(siblings)-1-i],
297 ChildR: currNode,
298 }
299 mt.Insert(node.Ht(), normalNodeType, 0, node.Bytes())
300 currNode = node.Ht()
301 }
302 }
303
304 return currNode, nil // currNode = root
305}
306
307// CheckProof validates the Merkle Proof for the leafHash and root
308func CheckProof(root Hash, proof []byte, hi Hash, ht Hash, numLevels int) bool {
309 var empties [32]byte
310 copy(empties[:], proof[:len(empties)])
311 hashLen := len(EmptyNodeValue)
312
313 var siblings []Hash
314 for i := len(empties); i < len(proof); i += hashLen {
315 var siblingHash Hash
316 copy(siblingHash[:], proof[i:i+hashLen])
317 siblings = append(siblings, siblingHash)
318 }
319
320 path := getPath(numLevels, hi)
321 nodeHash := ht
322 siblingUsedPos := 0
323
324 for level := numLevels - 2; level >= 0; level-- {
325 var sibling Hash
326 if testbitmap(empties[:], uint(level)) {
327 sibling = siblings[siblingUsedPos]
328 siblingUsedPos++
329 } else {
330 sibling = EmptyNodeValue
331 }
332 // calculate the nodeHash with the current nodeHash and the sibling
333 var node treeNode
334 if path[numLevels-level-2] {
335 node = treeNode{
336 ChildL: sibling,
337 ChildR: nodeHash,
338 }
339 } else {
340 node = treeNode{
341 ChildL: nodeHash,
342 ChildR: sibling,
343 }
344 }
345 // if both childs are EmptyNodeValue, the parent will be EmptyNodeValue
346 if bytes.Equal(nodeHash[:], EmptyNodeValue[:]) && bytes.Equal(sibling[:], EmptyNodeValue[:]) {
347 nodeHash = EmptyNodeValue
348 } else {
349 nodeHash = node.Ht()
350 }
351 }
352 return bytes.Equal(nodeHash[:], root[:])
353}