go-merkletree-iden3/merkletree.go
merkletree.goRaw
1package merkletree
2
3import (
4 "bytes"
5 "errors"
6 "fmt"
7 "math/big"
8 "sync"
9
10 "github.com/iden3/go-iden3-core/common"
11 "github.com/iden3/go-iden3-core/db"
12 cryptoUtils "github.com/iden3/go-iden3-crypto/utils"
13)
14
15const (
16 // proofFlagsLen is the byte length of the flags in the proof header (first 32
17 // bytes).
18 proofFlagsLen = 2
19 // ElemBytesLen is the length of the Hash byte array
20 ElemBytesLen = 32
21)
22
23var (
24 // ErrNodeKeyAlreadyExists is used when a node key already exists.
25 ErrNodeKeyAlreadyExists = errors.New("node already exists")
26 // ErrEntryIndexNotFound is used when no entry is found for an index.
27 ErrEntryIndexNotFound = errors.New("node index not found in the DB")
28 // ErrNodeDataBadSize is used when the data of a node has an incorrect
29 // size and can't be parsed.
30 ErrNodeDataBadSize = errors.New("node data has incorrect size in the DB")
31 // ErrReachedMaxLevel is used when a traversal of the MT reaches the
32 // maximum level.
33 ErrReachedMaxLevel = errors.New("reached maximum level of the merkle tree")
34 // ErrInvalidNodeFound is used when an invalid node is found and can't
35 // be parsed.
36 ErrInvalidNodeFound = errors.New("found an invalid node in the DB")
37 // ErrInvalidProofBytes is used when a serialized proof is invalid.
38 ErrInvalidProofBytes = errors.New("the serialized proof is invalid")
39 // ErrInvalidDBValue is used when a value in the key value DB is
40 // invalid (for example, it doen't contain a byte header and a []byte
41 // body of at least len=1.
42 ErrInvalidDBValue = errors.New("the value in the DB is invalid")
43 // ErrEntryIndexAlreadyExists is used when the entry index already
44 // exists in the tree.
45 ErrEntryIndexAlreadyExists = errors.New("the entry index already exists in the tree")
46 // ErrNotWritable is used when the MerkleTree is not writable and a write function is called
47 ErrNotWritable = errors.New("Merkle Tree not writable")
48
49 rootNodeValue = []byte("currentroot")
50 HashZero = 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}
51)
52
53type Hash [32]byte
54
55func (h Hash) String() string {
56 return new(big.Int).SetBytes(h[:]).String()
57}
58func (h *Hash) BigInt() *big.Int {
59 return new(big.Int).SetBytes(common.SwapEndianness(h[:]))
60}
61
62func NewHashFromBigInt(b *big.Int) *Hash {
63 r := &Hash{}
64 copy(r[:], common.SwapEndianness(b.Bytes()))
65 return r
66}
67
68type MerkleTree struct {
69 sync.RWMutex
70 db db.Storage
71 rootKey *Hash
72 writable bool
73 maxLevels int
74}
75
76func NewMerkleTree(storage db.Storage, maxLevels int) (*MerkleTree, error) {
77 mt := MerkleTree{db: storage, maxLevels: maxLevels, writable: true}
78
79 v, err := mt.db.Get(rootNodeValue)
80 if err != nil {
81 tx, err := mt.db.NewTx()
82 if err != nil {
83 return nil, err
84 }
85 mt.rootKey = &HashZero
86 tx.Put(rootNodeValue, mt.rootKey[:])
87 err = tx.Commit()
88 if err != nil {
89 return nil, err
90 }
91 return &mt, nil
92 }
93 mt.rootKey = &Hash{}
94 copy(mt.rootKey[:], v)
95 return &mt, nil
96}
97
98func (mt *MerkleTree) Root() *Hash {
99 return mt.rootKey
100}
101
102func (mt *MerkleTree) Add(k, v *big.Int) error {
103 // verify that the MerkleTree is writable
104 if !mt.writable {
105 return ErrNotWritable
106 }
107
108 // verfy that the ElemBytes are valid and fit inside the Finite Field.
109 if !cryptoUtils.CheckBigIntInField(k) {
110 return errors.New("Key not inside the Finite Field")
111 }
112 if !cryptoUtils.CheckBigIntInField(v) {
113 return errors.New("Value not inside the Finite Field")
114 }
115
116 tx, err := mt.db.NewTx()
117 if err != nil {
118 return err
119 }
120 mt.Lock()
121 defer mt.Unlock()
122
123 kHash := NewHashFromBigInt(k)
124 vHash := NewHashFromBigInt(v)
125 newNodeLeaf := NewNodeLeaf(kHash, vHash)
126 path := getPath(mt.maxLevels, kHash[:])
127
128 newRootKey, err := mt.addLeaf(tx, newNodeLeaf, mt.rootKey, 0, path)
129 if err != nil {
130 return err
131 }
132 mt.rootKey = newRootKey
133 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
134
135 if err := tx.Commit(); err != nil {
136 return err
137 }
138
139 return nil
140}
141
142// pushLeaf recursively pushes an existing oldLeaf down until its path diverges
143// from newLeaf, at which point both leafs are stored, all while updating the
144// path.
145func (mt *MerkleTree) pushLeaf(tx db.Tx, newLeaf *Node, oldLeaf *Node,
146 lvl int, pathNewLeaf []bool, pathOldLeaf []bool) (*Hash, error) {
147 if lvl > mt.maxLevels-2 {
148 return nil, ErrReachedMaxLevel
149 }
150 var newNodeMiddle *Node
151 if pathNewLeaf[lvl] == pathOldLeaf[lvl] { // We need to go deeper!
152 nextKey, err := mt.pushLeaf(tx, newLeaf, oldLeaf, lvl+1, pathNewLeaf, pathOldLeaf)
153 if err != nil {
154 return nil, err
155 }
156 if pathNewLeaf[lvl] {
157 newNodeMiddle = NewNodeMiddle(&HashZero, nextKey) // go right
158 } else {
159 newNodeMiddle = NewNodeMiddle(nextKey, &HashZero) // go left
160 }
161 return mt.addNode(tx, newNodeMiddle)
162 } else {
163 oldLeafKey, err := oldLeaf.Key()
164 if err != nil {
165 return nil, err
166 }
167 newLeafKey, err := newLeaf.Key()
168 if err != nil {
169 return nil, err
170 }
171
172 if pathNewLeaf[lvl] {
173 newNodeMiddle = NewNodeMiddle(oldLeafKey, newLeafKey)
174 } else {
175 newNodeMiddle = NewNodeMiddle(newLeafKey, oldLeafKey)
176 }
177 // We can add newLeaf now. We don't need to add oldLeaf because it's already in the tree.
178 _, err = mt.addNode(tx, newLeaf)
179 if err != nil {
180 return nil, err
181 }
182 return mt.addNode(tx, newNodeMiddle)
183 }
184}
185
186// addLeaf recursively adds a newLeaf in the MT while updating the path.
187func (mt *MerkleTree) addLeaf(tx db.Tx, newLeaf *Node, key *Hash,
188 lvl int, path []bool) (*Hash, error) {
189 var err error
190 var nextKey *Hash
191 if lvl > mt.maxLevels-1 {
192 return nil, ErrReachedMaxLevel
193 }
194 n, err := mt.GetNode(key)
195 if err != nil {
196 return nil, err
197 }
198 switch n.Type {
199 case NodeTypeEmpty:
200 // We can add newLeaf now
201 return mt.addNode(tx, newLeaf)
202 case NodeTypeLeaf:
203 nKey := n.Entry[0]
204 // Check if leaf node found contains the leaf node we are trying to add
205 newLeafKey := newLeaf.Entry[0]
206 if bytes.Equal(nKey[:], newLeafKey[:]) {
207 return nil, ErrEntryIndexAlreadyExists
208 }
209 pathOldLeaf := getPath(mt.maxLevels, nKey[:])
210 // We need to push newLeaf down until its path diverges from n's path
211 return mt.pushLeaf(tx, newLeaf, n, lvl, path, pathOldLeaf)
212 case NodeTypeMiddle:
213 // We need to go deeper, continue traversing the tree, left or right depending on path
214 var newNodeMiddle *Node
215 if path[lvl] {
216 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildR, lvl+1, path) // go right
217 newNodeMiddle = NewNodeMiddle(n.ChildL, nextKey)
218 } else {
219 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildL, lvl+1, path) // go left
220 newNodeMiddle = NewNodeMiddle(nextKey, n.ChildR)
221 }
222 if err != nil {
223 return nil, err
224 }
225 // Update the node to reflect the modified child
226 return mt.addNode(tx, newNodeMiddle)
227 default:
228 return nil, ErrInvalidNodeFound
229 }
230}
231
232// addNode adds a node into the MT. Empty nodes are not stored in the tree;
233// they are all the same and assumed to always exist.
234func (mt *MerkleTree) addNode(tx db.Tx, n *Node) (*Hash, error) {
235 // verify that the MerkleTree is writable
236 if !mt.writable {
237 return nil, ErrNotWritable
238 }
239 if n.Type == NodeTypeEmpty {
240 return n.Key()
241 }
242 k, err := n.Key()
243 if err != nil {
244 return nil, err
245 }
246 v := n.Value()
247 // Check that the node key doesn't already exist
248 if _, err := tx.Get(k[:]); err == nil {
249 return nil, ErrNodeKeyAlreadyExists
250 }
251 tx.Put(k[:], v)
252 return k, nil
253}
254
255// dbGet is a helper function to get the node of a key from the internal
256// storage.
257func (mt *MerkleTree) dbGet(k []byte) (NodeType, []byte, error) {
258 if bytes.Equal(k, HashZero[:]) {
259 return 0, nil, nil
260 }
261
262 value, err := mt.db.Get(k)
263 if err != nil {
264 return 0, nil, err
265 }
266
267 if len(value) < 2 {
268 return 0, nil, ErrInvalidDBValue
269 }
270 nodeType := value[0]
271 nodeBytes := value[1:]
272
273 return NodeType(nodeType), nodeBytes, nil
274}
275
276// dbInsert is a helper function to insert a node into a key in an open db
277// transaction.
278func (mt *MerkleTree) dbInsert(tx db.Tx, k []byte, t NodeType, data []byte) {
279 v := append([]byte{byte(t)}, data...)
280 tx.Put(k, v)
281}
282
283// GetNode gets a node by key from the MT. Empty nodes are not stored in the
284// tree; they are all the same and assumed to always exist.
285func (mt *MerkleTree) GetNode(key *Hash) (*Node, error) {
286 if bytes.Equal(key[:], HashZero[:]) {
287 return NewNodeEmpty(), nil
288 }
289 nBytes, err := mt.db.Get(key[:])
290 if err != nil {
291 return nil, err
292 }
293 return NewNodeFromBytes(nBytes)
294}
295
296// getPath returns the binary path, from the root to the leaf.
297func getPath(numLevels int, k []byte) []bool {
298 path := make([]bool, numLevels)
299 for n := 0; n < numLevels; n++ {
300 path[n] = common.TestBit(k[:], uint(n))
301 }
302 return path
303}
304
305// NodeAux contains the auxiliary node used in a non-existence proof.
306type NodeAux struct {
307 Key *Hash
308 Value *Hash
309}
310
311// Proof defines the required elements for a MT proof of existence or non-existence.
312type Proof struct {
313 // existence indicates wether this is a proof of existence or non-existence.
314 Existence bool
315 // depth indicates how deep in the tree the proof goes.
316 depth uint
317 // notempties is a bitmap of non-empty Siblings found in Siblings.
318 notempties [ElemBytesLen - proofFlagsLen]byte
319 // Siblings is a list of non-empty sibling keys.
320 Siblings []*Hash
321 NodeAux *NodeAux
322}
323
324// NewProofFromBytes parses a byte array into a Proof.
325func NewProofFromBytes(bs []byte) (*Proof, error) {
326 if len(bs) < ElemBytesLen {
327 return nil, ErrInvalidProofBytes
328 }
329 p := &Proof{}
330 if (bs[0] & 0x01) == 0 {
331 p.Existence = true
332 }
333 p.depth = uint(bs[1])
334 copy(p.notempties[:], bs[proofFlagsLen:ElemBytesLen])
335 siblingBytes := bs[ElemBytesLen:]
336 sibIdx := 0
337 for i := uint(0); i < p.depth; i++ {
338 if common.TestBitBigEndian(p.notempties[:], i) {
339 if len(siblingBytes) < (sibIdx+1)*ElemBytesLen {
340 return nil, ErrInvalidProofBytes
341 }
342 var sib Hash
343 copy(sib[:], siblingBytes[sibIdx*ElemBytesLen:(sibIdx+1)*ElemBytesLen])
344 p.Siblings = append(p.Siblings, &sib)
345 sibIdx++
346 }
347 }
348
349 if !p.Existence && ((bs[0] & 0x02) != 0) {
350 p.NodeAux = &NodeAux{Key: &Hash{}, Value: &Hash{}}
351 nodeAuxBytes := siblingBytes[len(p.Siblings)*ElemBytesLen:]
352 if len(nodeAuxBytes) != 2*ElemBytesLen {
353 return nil, ErrInvalidProofBytes
354 }
355 copy(p.NodeAux.Key[:], nodeAuxBytes[:ElemBytesLen])
356 copy(p.NodeAux.Value[:], nodeAuxBytes[ElemBytesLen:2*ElemBytesLen])
357 }
358 return p, nil
359}
360
361// Bytes serializes a Proof into a byte array.
362func (p *Proof) Bytes() []byte {
363 bsLen := proofFlagsLen + len(p.notempties) + ElemBytesLen*len(p.Siblings)
364 if p.NodeAux != nil {
365 bsLen += 2 * ElemBytesLen
366 }
367 bs := make([]byte, bsLen)
368
369 if !p.Existence {
370 bs[0] |= 0x01
371 }
372 bs[1] = byte(p.depth)
373 copy(bs[proofFlagsLen:len(p.notempties)+proofFlagsLen], p.notempties[:])
374 siblingsBytes := bs[len(p.notempties)+proofFlagsLen:]
375 for i, k := range p.Siblings {
376 copy(siblingsBytes[i*ElemBytesLen:(i+1)*ElemBytesLen], k[:])
377 }
378 if p.NodeAux != nil {
379 bs[0] |= 0x02
380 copy(bs[len(bs)-2*ElemBytesLen:], p.NodeAux.Key[:])
381 copy(bs[len(bs)-1*ElemBytesLen:], p.NodeAux.Value[:])
382 }
383 return bs
384}
385
386// GenerateProof generates the proof of existence (or non-existence) of an
387// Entry's hash Index for a Merkle Tree given the root.
388// If the rootKey is nil, the current merkletree root is used
389func (mt *MerkleTree) GenerateProof(k *big.Int, rootKey *Hash) (*Proof, error) {
390 p := &Proof{}
391 var siblingKey *Hash
392
393 kHash := NewHashFromBigInt(k)
394 path := getPath(mt.maxLevels, kHash[:])
395 if rootKey == nil {
396 rootKey = mt.Root()
397 }
398 nextKey := rootKey
399 for p.depth = 0; p.depth < uint(mt.maxLevels); p.depth++ {
400 n, err := mt.GetNode(nextKey)
401 if err != nil {
402 return nil, err
403 }
404 switch n.Type {
405 case NodeTypeEmpty:
406 return p, nil
407 case NodeTypeLeaf:
408 if bytes.Equal(kHash[:], n.Entry[0][:]) {
409 p.Existence = true
410 return p, nil
411 } else {
412 // We found a leaf whose entry didn't match hIndex
413 p.NodeAux = &NodeAux{Key: n.Entry[0], Value: n.Entry[1]}
414 return p, nil
415 }
416 case NodeTypeMiddle:
417 if path[p.depth] {
418 nextKey = n.ChildR
419 siblingKey = n.ChildL
420 } else {
421 nextKey = n.ChildL
422 siblingKey = n.ChildR
423 }
424 default:
425 return nil, ErrInvalidNodeFound
426 }
427 if !bytes.Equal(siblingKey[:], HashZero[:]) {
428 common.SetBitBigEndian(p.notempties[:], uint(p.depth))
429 p.Siblings = append(p.Siblings, siblingKey)
430 }
431 }
432 return nil, ErrEntryIndexNotFound
433}
434
435// VerifyProof verifies the Merkle Proof for the entry and root.
436func VerifyProof(rootKey *Hash, proof *Proof, k, v *big.Int) bool {
437 rootFromProof, err := RootFromProof(proof, k, v)
438 if err != nil {
439 return false
440 }
441 return bytes.Equal(rootKey[:], rootFromProof[:])
442}
443
444// RootFromProof calculates the root that would correspond to a tree whose
445// siblings are the ones in the proof with the claim hashing to hIndex and
446// hValue.
447func RootFromProof(proof *Proof, k, v *big.Int) (*Hash, error) {
448 kHash := NewHashFromBigInt(k)
449 vHash := NewHashFromBigInt(v)
450 sibIdx := len(proof.Siblings) - 1
451 var err error
452 var midKey *Hash
453 if proof.Existence {
454 midKey, err = LeafKey(kHash, vHash)
455 if err != nil {
456 return nil, err
457 }
458 } else {
459 if proof.NodeAux == nil {
460 midKey = &HashZero
461 } else {
462 if bytes.Equal(kHash[:], proof.NodeAux.Key[:]) {
463 return nil, fmt.Errorf("Non-existence proof being checked against hIndex equal to nodeAux")
464 }
465 midKey, err = LeafKey(proof.NodeAux.Key, proof.NodeAux.Value)
466 if err != nil {
467 return nil, err
468 }
469 }
470 }
471 path := getPath(int(proof.depth), kHash[:])
472 var siblingKey *Hash
473 for lvl := int(proof.depth) - 1; lvl >= 0; lvl-- {
474 if common.TestBitBigEndian(proof.notempties[:], uint(lvl)) {
475 siblingKey = proof.Siblings[sibIdx]
476 sibIdx--
477 } else {
478 siblingKey = &HashZero
479 }
480 if path[lvl] {
481 midKey, err = NewNodeMiddle(siblingKey, midKey).Key()
482 if err != nil {
483 return nil, err
484 }
485 } else {
486 midKey, err = NewNodeMiddle(midKey, siblingKey).Key()
487 if err != nil {
488 return nil, err
489 }
490 }
491 }
492 return midKey, nil
493}