go-merkletree-iden3/merkletree.go
merkletree.goRaw
1package merkletree
2
3import (
4 "bytes"
5 "errors"
6 "math/big"
7 "sync"
8
9 "github.com/iden3/go-iden3-core/common"
10 "github.com/iden3/go-iden3-core/db"
11 cryptoUtils "github.com/iden3/go-iden3-crypto/utils"
12)
13
14const (
15 // proofFlagsLen is the byte length of the flags in the proof header (first 32
16 // bytes).
17 proofFlagsLen = 2
18 // ElemBytesLen is the length of the Hash byte array
19 ElemBytesLen = 32
20)
21
22var (
23 // ErrNodeKeyAlreadyExists is used when a node key already exists.
24 ErrNodeKeyAlreadyExists = errors.New("node already exists")
25 // ErrEntryIndexNotFound is used when no entry is found for an index.
26 ErrEntryIndexNotFound = errors.New("node index not found in the DB")
27 // ErrNodeDataBadSize is used when the data of a node has an incorrect
28 // size and can't be parsed.
29 ErrNodeDataBadSize = errors.New("node data has incorrect size in the DB")
30 // ErrReachedMaxLevel is used when a traversal of the MT reaches the
31 // maximum level.
32 ErrReachedMaxLevel = errors.New("reached maximum level of the merkle tree")
33 // ErrInvalidNodeFound is used when an invalid node is found and can't
34 // be parsed.
35 ErrInvalidNodeFound = errors.New("found an invalid node in the DB")
36 // ErrInvalidProofBytes is used when a serialized proof is invalid.
37 ErrInvalidProofBytes = errors.New("the serialized proof is invalid")
38 // ErrInvalidDBValue is used when a value in the key value DB is
39 // invalid (for example, it doen't contain a byte header and a []byte
40 // body of at least len=1.
41 ErrInvalidDBValue = errors.New("the value in the DB is invalid")
42 // ErrEntryIndexAlreadyExists is used when the entry index already
43 // exists in the tree.
44 ErrEntryIndexAlreadyExists = errors.New("the entry index already exists in the tree")
45 // ErrNotWritable is used when the MerkleTree is not writable and a write function is called
46 ErrNotWritable = errors.New("Merkle Tree not writable")
47
48 rootNodeValue = []byte("currentroot")
49 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}
50)
51
52type Hash [32]byte
53
54func (h Hash) String() string {
55 return new(big.Int).SetBytes(h[:]).String()
56}
57func (h *Hash) BigInt() *big.Int {
58 return new(big.Int).SetBytes(common.SwapEndianness(h[:]))
59}
60
61func NewHashFromBigInt(b *big.Int) *Hash {
62 r := &Hash{}
63 copy(r[:], common.SwapEndianness(b.Bytes()))
64 return r
65}
66
67type MerkleTree struct {
68 sync.RWMutex
69 db db.Storage
70 rootKey *Hash
71 writable bool
72 maxLevels int
73}
74
75func NewMerkleTree(storage db.Storage, maxLevels int) (*MerkleTree, error) {
76 mt := MerkleTree{db: storage, maxLevels: maxLevels, writable: true}
77
78 v, err := mt.db.Get(rootNodeValue)
79 if err != nil {
80 tx, err := mt.db.NewTx()
81 if err != nil {
82 return nil, err
83 }
84 mt.rootKey = &HashZero
85 tx.Put(rootNodeValue, mt.rootKey[:])
86 err = tx.Commit()
87 if err != nil {
88 return nil, err
89 }
90 return &mt, nil
91 }
92 mt.rootKey = &Hash{}
93 copy(mt.rootKey[:], v)
94 return &mt, nil
95}
96
97func (mt *MerkleTree) Root() *Hash {
98 return mt.rootKey
99}
100
101func (mt *MerkleTree) Add(k, v *big.Int) error {
102 // verify that the MerkleTree is writable
103 if !mt.writable {
104 return ErrNotWritable
105 }
106
107 // verfy that the ElemBytes are valid and fit inside the Finite Field.
108 if !cryptoUtils.CheckBigIntInField(k) {
109 return errors.New("Key not inside the Finite Field")
110 }
111 if !cryptoUtils.CheckBigIntInField(v) {
112 return errors.New("Value not inside the Finite Field")
113 }
114
115 tx, err := mt.db.NewTx()
116 if err != nil {
117 return err
118 }
119 mt.Lock()
120 defer mt.Unlock()
121
122 kHash := NewHashFromBigInt(k)
123 vHash := NewHashFromBigInt(v)
124 newNodeLeaf := NewNodeLeaf(kHash, vHash)
125 path := getPath(mt.maxLevels, kHash[:])
126
127 newRootKey, err := mt.addLeaf(tx, newNodeLeaf, mt.rootKey, 0, path)
128 if err != nil {
129 return err
130 }
131 mt.rootKey = newRootKey
132 mt.dbInsert(tx, rootNodeValue, DBEntryTypeRoot, mt.rootKey[:])
133
134 if err := tx.Commit(); err != nil {
135 return err
136 }
137
138 return nil
139}
140
141// pushLeaf recursively pushes an existing oldLeaf down until its path diverges
142// from newLeaf, at which point both leafs are stored, all while updating the
143// path.
144func (mt *MerkleTree) pushLeaf(tx db.Tx, newLeaf *Node, oldLeaf *Node,
145 lvl int, pathNewLeaf []bool, pathOldLeaf []bool) (*Hash, error) {
146 if lvl > mt.maxLevels-2 {
147 return nil, ErrReachedMaxLevel
148 }
149 var newNodeMiddle *Node
150 if pathNewLeaf[lvl] == pathOldLeaf[lvl] { // We need to go deeper!
151 nextKey, err := mt.pushLeaf(tx, newLeaf, oldLeaf, lvl+1, pathNewLeaf, pathOldLeaf)
152 if err != nil {
153 return nil, err
154 }
155 if pathNewLeaf[lvl] {
156 newNodeMiddle = NewNodeMiddle(&HashZero, nextKey) // go right
157 } else {
158 newNodeMiddle = NewNodeMiddle(nextKey, &HashZero) // go left
159 }
160 return mt.addNode(tx, newNodeMiddle)
161 } else {
162 oldLeafKey, err := oldLeaf.Key()
163 if err != nil {
164 return nil, err
165 }
166 newLeafKey, err := newLeaf.Key()
167 if err != nil {
168 return nil, err
169 }
170
171 if pathNewLeaf[lvl] {
172 newNodeMiddle = NewNodeMiddle(oldLeafKey, newLeafKey)
173 } else {
174 newNodeMiddle = NewNodeMiddle(newLeafKey, oldLeafKey)
175 }
176 // We can add newLeaf now. We don't need to add oldLeaf because it's already in the tree.
177 _, err = mt.addNode(tx, newLeaf)
178 if err != nil {
179 return nil, err
180 }
181 return mt.addNode(tx, newNodeMiddle)
182 }
183}
184
185// addLeaf recursively adds a newLeaf in the MT while updating the path.
186func (mt *MerkleTree) addLeaf(tx db.Tx, newLeaf *Node, key *Hash,
187 lvl int, path []bool) (*Hash, error) {
188 var err error
189 var nextKey *Hash
190 if lvl > mt.maxLevels-1 {
191 return nil, ErrReachedMaxLevel
192 }
193 n, err := mt.GetNode(key)
194 if err != nil {
195 return nil, err
196 }
197 switch n.Type {
198 case NodeTypeEmpty:
199 // We can add newLeaf now
200 return mt.addNode(tx, newLeaf)
201 case NodeTypeLeaf:
202 nKey := n.Entry[0]
203 // Check if leaf node found contains the leaf node we are trying to add
204 newLeafKey := newLeaf.Entry[0]
205 if bytes.Equal(nKey[:], newLeafKey[:]) {
206 return nil, ErrEntryIndexAlreadyExists
207 }
208 pathOldLeaf := getPath(mt.maxLevels, nKey[:])
209 // We need to push newLeaf down until its path diverges from n's path
210 return mt.pushLeaf(tx, newLeaf, n, lvl, path, pathOldLeaf)
211 case NodeTypeMiddle:
212 // We need to go deeper, continue traversing the tree, left or right depending on path
213 var newNodeMiddle *Node
214 if path[lvl] {
215 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildR, lvl+1, path) // go right
216 newNodeMiddle = NewNodeMiddle(n.ChildL, nextKey)
217 } else {
218 nextKey, err = mt.addLeaf(tx, newLeaf, n.ChildL, lvl+1, path) // go left
219 newNodeMiddle = NewNodeMiddle(nextKey, n.ChildR)
220 }
221 if err != nil {
222 return nil, err
223 }
224 // Update the node to reflect the modified child
225 return mt.addNode(tx, newNodeMiddle)
226 default:
227 return nil, ErrInvalidNodeFound
228 }
229}
230
231// addNode adds a node into the MT. Empty nodes are not stored in the tree;
232// they are all the same and assumed to always exist.
233func (mt *MerkleTree) addNode(tx db.Tx, n *Node) (*Hash, error) {
234 // verify that the MerkleTree is writable
235 if !mt.writable {
236 return nil, ErrNotWritable
237 }
238 if n.Type == NodeTypeEmpty {
239 return n.Key()
240 }
241 k, err := n.Key()
242 if err != nil {
243 return nil, err
244 }
245 v := n.Value()
246 // Check that the node key doesn't already exist
247 if _, err := tx.Get(k[:]); err == nil {
248 return nil, ErrNodeKeyAlreadyExists
249 }
250 tx.Put(k[:], v)
251 return k, nil
252}
253
254// dbGet is a helper function to get the node of a key from the internal
255// storage.
256func (mt *MerkleTree) dbGet(k []byte) (NodeType, []byte, error) {
257 if bytes.Equal(k, HashZero[:]) {
258 return 0, nil, nil
259 }
260
261 value, err := mt.db.Get(k)
262 if err != nil {
263 return 0, nil, err
264 }
265
266 if len(value) < 2 {
267 return 0, nil, ErrInvalidDBValue
268 }
269 nodeType := value[0]
270 nodeBytes := value[1:]
271
272 return NodeType(nodeType), nodeBytes, nil
273}
274
275// dbInsert is a helper function to insert a node into a key in an open db
276// transaction.
277func (mt *MerkleTree) dbInsert(tx db.Tx, k []byte, t NodeType, data []byte) {
278 v := append([]byte{byte(t)}, data...)
279 tx.Put(k, v)
280}
281
282// GetNode gets a node by key from the MT. Empty nodes are not stored in the
283// tree; they are all the same and assumed to always exist.
284func (mt *MerkleTree) GetNode(key *Hash) (*Node, error) {
285 if bytes.Equal(key[:], HashZero[:]) {
286 return NewNodeEmpty(), nil
287 }
288 nBytes, err := mt.db.Get(key[:])
289 if err != nil {
290 return nil, err
291 }
292 return NewNodeFromBytes(nBytes)
293}
294
295// getPath returns the binary path, from the root to the leaf.
296func getPath(numLevels int, k []byte) []bool {
297 path := make([]bool, numLevels)
298 for n := 0; n < numLevels; n++ {
299 path[n] = common.TestBit(k[:], uint(n))
300 }
301 return path
302}