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 if mt.root, err = mt.replaceLeaf(siblings, path[posDiff+1:], parentNode.Ht(), normalNodeType, 0, parentNode.Bytes()); err != nil {
117 return err
118 }
119 mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
120 return nil
121 }
122 node := parseNodeBytes(nodeBytes)
123 var sibling Hash
124 if !path[i] {
125 nodeHash = node.ChildL
126 sibling = node.ChildR
127 } else {
128 nodeHash = node.ChildR
129 sibling = node.ChildL
130 }
131 siblings = append(siblings, sibling)
132
133 if bytes.Equal(nodeHash[:], EmptyNodeValue[:]) {
134 // if the node is EmptyNodeValue, the leaf data will go directly at that height, as a Final Node
135 if i == mt.numLevels-2 && bytes.Equal(siblings[len(siblings)-1][:], EmptyNodeValue[:]) {
136 // if the pt node is the unique in the tree, just put it into the root node
137 // this means to be in i==mt.NumLevels-2 && nodeHash==EmptyNodeValue
138 finalNodeHash := calcHashFromLeafAndLevel(i+1, path, HashBytes(v.Bytes()))
139 mt.Insert(finalNodeHash, finalNodeType, v.IndexLength(), v.Bytes())
140 mt.root = finalNodeHash
141 mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
142 return nil
143 }
144 finalNodeHash := calcHashFromLeafAndLevel(i, path, HashBytes(v.Bytes()))
145 if mt.root, err = mt.replaceLeaf(siblings, path[i:], finalNodeHash, finalNodeType, v.IndexLength(), v.Bytes()); err != nil {
146 return err
147 }
148 mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
149 return nil
150 }
151 }
152
153 var err error
154 mt.root, err = mt.replaceLeaf(siblings, path, HashBytes(v.Bytes()), valueNodeType, v.IndexLength(), v.Bytes())
155 if err != nil {
156 return err
157 }
158 mt.Insert(rootNodeValue, rootNodeType, 0, mt.root[:])
159 return nil
160}
161
162// GenerateProof generates the Merkle Proof from a given leafHash for the current root
163func (mt *MerkleTree) GenerateProof(hi Hash) ([]byte, error) {
164 var empties [32]byte
165
166 path := getPath(mt.numLevels, hi)
167 var siblings []Hash
168 nodeHash := mt.root
169
170 for level := 0; level < mt.numLevels-1; level++ {
171 nodeType, indexLength, nodeBytes, err := mt.Get(nodeHash)
172 if err != nil {
173 return nil, err
174 }
175 if nodeType == byte(finalNodeType) {
176 realValueInPos, err := mt.GetValueInPos(hi)
177 if err != nil {
178 return nil, err
179 }
180 if bytes.Equal(realValueInPos[:], EmptyNodeValue[:]) {
181 // 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
182 leafHi := HashBytes(nodeBytes[:indexLength]) // hi of element that was in the end of the branch (the finalNode)
183 pathChild := getPath(mt.numLevels, leafHi)
184
185 // get the position where the path is different
186 posDiff := comparePaths(pathChild, path)
187 if posDiff == -1 {
188 return nil, ErrNodeAlreadyExists
189 }
190
191 if posDiff != mt.NumLevels()-1-level {
192 sibling := calcHashFromLeafAndLevel(posDiff, pathChild, HashBytes(nodeBytes))
193 setbitmap(empties[:], uint(mt.NumLevels()-2-posDiff))
194 siblings = append([]Hash{sibling}, siblings...)
195 }
196
197 }
198 break
199 }
200 node := parseNodeBytes(nodeBytes)
201
202 var sibling Hash
203 if !path[mt.numLevels-level-2] {
204 nodeHash = node.ChildL
205 sibling = node.ChildR
206 } else {
207 nodeHash = node.ChildR
208 sibling = node.ChildL
209 }
210 if !bytes.Equal(sibling[:], EmptyNodeValue[:]) {
211 setbitmap(empties[:], uint(level))
212 siblings = append([]Hash{sibling}, siblings...)
213 }
214 }
215 // merge empties and siblings
216 var mp []byte
217 mp = append(mp, empties[:]...)
218 for k := range siblings {
219 mp = append(mp, siblings[k][:]...)
220 }
221 return mp, nil
222}
223
224// GetValueInPos returns the merkletree value in the position of the Hash of the Index (Hi)
225func (mt *MerkleTree) GetValueInPos(hi Hash) ([]byte, error) {
226 path := getPath(mt.numLevels, hi)
227 nodeHash := mt.root
228 for i := mt.numLevels - 2; i >= 0; i-- {
229 nodeType, indexLength, nodeBytes, err := mt.Get(nodeHash)
230 if err != nil {
231 return nodeBytes, err
232 }
233 if nodeType == byte(finalNodeType) {
234 // check if nodeBytes path is different of hi
235 index := nodeBytes[:indexLength]
236 hi := HashBytes(index)
237 nodePath := getPath(mt.numLevels, hi)
238 posDiff := comparePaths(path, nodePath)
239 // if is different, return an EmptyNodeValue, else return the nodeBytes
240 if posDiff != -1 {
241 return EmptyNodeValue[:], nil
242 }
243 return nodeBytes, nil
244 }
245 node := parseNodeBytes(nodeBytes)
246 if !path[i] {
247 nodeHash = node.ChildL
248 } else {
249 nodeHash = node.ChildR
250 }
251 }
252 _, _, valueBytes, err := mt.Get(nodeHash)
253 if err != nil {
254 return valueBytes, err
255 }
256 return valueBytes, nil
257}
258
259func calcHashFromLeafAndLevel(untilLevel int, path []bool, leafHash Hash) Hash {
260 nodeCurrLevel := leafHash
261 for i := 0; i < untilLevel; i++ {
262 if path[i] {
263 node := treeNode{
264 ChildL: EmptyNodeValue,
265 ChildR: nodeCurrLevel,
266 }
267 nodeCurrLevel = node.Ht()
268 } else {
269 node := treeNode{
270 ChildL: nodeCurrLevel,
271 ChildR: EmptyNodeValue,
272 }
273 nodeCurrLevel = node.Ht()
274 }
275 }
276 return nodeCurrLevel
277}
278
279func (mt *MerkleTree) replaceLeaf(siblings []Hash, path []bool, newLeafHash Hash, nodetype byte, indexLength uint32, newLeafValue []byte) (Hash, error) {
280 // add the new leaf
281 mt.Insert(newLeafHash, nodetype, indexLength, newLeafValue)
282 currNode := newLeafHash
283 // here the path is only the path[posDiff+1]
284 for i := 0; i < len(siblings); i++ {
285 if !path[i] {
286 node := treeNode{
287 ChildL: currNode,
288 ChildR: siblings[len(siblings)-1-i],
289 }
290 mt.Insert(node.Ht(), normalNodeType, 0, node.Bytes())
291 currNode = node.Ht()
292 } else {
293
294 node := treeNode{
295 ChildL: siblings[len(siblings)-1-i],
296 ChildR: currNode,
297 }
298 mt.Insert(node.Ht(), normalNodeType, 0, node.Bytes())
299 currNode = node.Ht()
300 }
301 }
302
303 return currNode, nil // currNode = root
304}
305
306// CheckProof validates the Merkle Proof for the leafHash and root
307func CheckProof(root Hash, proof []byte, hi Hash, ht Hash, numLevels int) bool {
308 var empties [32]byte
309 copy(empties[:], proof[:len(empties)])
310 hashLen := len(EmptyNodeValue)
311
312 var siblings []Hash
313 for i := len(empties); i < len(proof); i += hashLen {
314 var siblingHash Hash
315 copy(siblingHash[:], proof[i:i+hashLen])
316 siblings = append(siblings, siblingHash)
317 }
318
319 path := getPath(numLevels, hi)
320 nodeHash := ht
321 siblingUsedPos := 0
322
323 for level := numLevels - 2; level >= 0; level-- {
324 var sibling Hash
325 if testbitmap(empties[:], uint(level)) {
326 sibling = siblings[siblingUsedPos]
327 siblingUsedPos++
328 } else {
329 sibling = EmptyNodeValue
330 }
331 // calculate the nodeHash with the current nodeHash and the sibling
332 var node treeNode
333 if path[numLevels-level-2] {
334 node = treeNode{
335 ChildL: sibling,
336 ChildR: nodeHash,
337 }
338 } else {
339 node = treeNode{
340 ChildL: nodeHash,
341 ChildR: sibling,
342 }
343 }
344 // if both childs are EmptyNodeValue, the parent will be EmptyNodeValue
345 if bytes.Equal(nodeHash[:], EmptyNodeValue[:]) && bytes.Equal(sibling[:], EmptyNodeValue[:]) {
346 nodeHash = EmptyNodeValue
347 } else {
348 nodeHash = node.Ht()
349 }
350 }
351 return bytes.Equal(nodeHash[:], root[:])
352}