arbo/tree.go
tree.goRaw
1/*
2Package arbo implements a Merkle Tree compatible with the circomlib
3implementation of the MerkleTree (when using the Poseidon hash function),
4following the specification from
5https://docs.iden3.io/publications/pdfs/Merkle-Tree.pdf and
6https://eprint.iacr.org/2018/955.
7
8Also allows to define which hash function to use. So for example, when working
9with zkSnarks the Poseidon hash function can be used, but when not, it can be
10used the Blake3 hash function, which improves the computation time.
11*/
12package arbo
13
14import (
15 "bytes"
16 "fmt"
17 "sync/atomic"
18 "time"
19
20 "github.com/iden3/go-merkletree/db"
21)
22
23const (
24 // PrefixValueLen defines the bytes-prefix length used for the Value
25 // bytes representation stored in the db
26 PrefixValueLen = 2
27
28 // PrefixValueEmpty is used for the first byte of a Value to indicate
29 // that is an Empty value
30 PrefixValueEmpty = 0
31 // PrefixValueLeaf is used for the first byte of a Value to indicate
32 // that is a Leaf value
33 PrefixValueLeaf = 1
34 // PrefixValueIntermediate is used for the first byte of a Value to
35 // indicate that is a Intermediate value
36 PrefixValueIntermediate = 2
37)
38
39var (
40 dbKeyRoot = []byte("root")
41 emptyValue = []byte{0}
42)
43
44// Tree defines the struct that implements the MerkleTree functionalities
45type Tree struct {
46 db db.Storage
47 lastAccess int64 // in unix time
48 maxLevels int
49 root []byte
50
51 hashFunction HashFunction
52}
53
54// NewTree returns a new Tree, if there is a Tree still in the given storage, it
55// will load it.
56func NewTree(storage db.Storage, maxLevels int, hash HashFunction) (*Tree, error) {
57 t := Tree{db: storage, maxLevels: maxLevels, hashFunction: hash}
58
59 t.updateAccessTime()
60 root, err := t.db.Get(dbKeyRoot)
61 if err == db.ErrNotFound {
62 // store new root 0
63 tx, err := t.db.NewTx()
64 if err != nil {
65 return nil, err
66 }
67 t.root = make([]byte, t.hashFunction.Len()) // empty
68 err = tx.Put(dbKeyRoot, t.root)
69 if err != nil {
70 return nil, err
71 }
72 err = tx.Commit()
73 if err != nil {
74 return nil, err
75 }
76 return &t, err
77 } else if err != nil {
78 return nil, err
79 }
80 t.root = root
81 return &t, nil
82}
83
84func (t *Tree) updateAccessTime() {
85 atomic.StoreInt64(&t.lastAccess, time.Now().Unix())
86}
87
88// LastAccess returns the last access timestamp in Unixtime
89func (t *Tree) LastAccess() int64 {
90 return atomic.LoadInt64(&t.lastAccess)
91}
92
93// Root returns the root of the Tree
94func (t *Tree) Root() []byte {
95 return t.root
96}
97
98// AddBatch adds a batch of key-values to the Tree. This method is optimized to
99// do some internal parallelization. Returns an array containing the indexes of
100// the keys failed to add.
101func (t *Tree) AddBatch(keys, values [][]byte) ([]int, error) {
102 return nil, fmt.Errorf("unimplemented")
103}
104
105// Add inserts the key-value into the Tree.
106// If the inputs come from a *big.Int, is expected that are represented by a
107// Little-Endian byte array (for circom compatibility).
108func (t *Tree) Add(k, v []byte) error {
109 // TODO check validity of key & value (for the Tree.HashFunction type)
110
111 keyPath := make([]byte, t.hashFunction.Len())
112 copy(keyPath[:], k)
113
114 path := getPath(t.maxLevels, keyPath)
115 // go down to the leaf
116 var siblings [][]byte
117 _, _, siblings, err := t.down(k, t.root, siblings, path, 0)
118 if err != nil {
119 return err
120 }
121
122 leafKey, leafValue, err := t.newLeafValue(k, v)
123 if err != nil {
124 return err
125 }
126
127 tx, err := t.db.NewTx()
128 if err != nil {
129 return err
130 }
131 if err := tx.Put(leafKey, leafValue); err != nil {
132 return err
133 }
134
135 // go up to the root
136 if len(siblings) == 0 {
137 t.root = leafKey
138 return tx.Commit()
139 }
140 root, err := t.up(tx, leafKey, siblings, path, len(siblings)-1)
141 if err != nil {
142 return err
143 }
144
145 t.root = root
146 // store root to db
147 return tx.Commit()
148}
149
150// down goes down to the leaf recursively
151func (t *Tree) down(newKey, currKey []byte, siblings [][]byte, path []bool, l int) (
152 []byte, []byte, [][]byte, error) {
153 if l > t.maxLevels-1 {
154 return nil, nil, nil, fmt.Errorf("max level")
155 }
156 var err error
157 var currValue []byte
158 emptyKey := make([]byte, t.hashFunction.Len())
159 if bytes.Equal(currKey, emptyKey) {
160 // empty value
161 return currKey, emptyValue, siblings, nil
162 }
163 currValue, err = t.db.Get(currKey)
164 if err != nil {
165 return nil, nil, nil, err
166 }
167
168 switch currValue[0] {
169 case PrefixValueEmpty: // empty
170 // TODO WIP WARNING should not be reached, as the 'if' above should avoid
171 // reaching this point
172 // return currKey, empty, siblings, nil
173 panic("should not be reached, as the 'if' above should avoid reaching this point") // TMP
174 case PrefixValueLeaf: // leaf
175 if bytes.Equal(newKey, currKey) {
176 return nil, nil, nil, fmt.Errorf("key already exists")
177 }
178
179 if !bytes.Equal(currValue, emptyValue) {
180 oldLeafKey, _ := readLeafValue(currValue)
181 oldLeafKeyFull := make([]byte, t.hashFunction.Len())
182 copy(oldLeafKeyFull[:], oldLeafKey)
183
184 // if currKey is already used, go down until paths diverge
185 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
186 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, l)
187 if err != nil {
188 return nil, nil, nil, err
189 }
190 }
191 return currKey, currValue, siblings, nil
192 case PrefixValueIntermediate: // intermediate
193 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
194 return nil, nil, nil,
195 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
196 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
197 }
198 // collect siblings while going down
199 if path[l] {
200 // right
201 lChild, rChild := readIntermediateChilds(currValue)
202 siblings = append(siblings, lChild)
203 return t.down(newKey, rChild, siblings, path, l+1)
204 }
205 // left
206 lChild, rChild := readIntermediateChilds(currValue)
207 siblings = append(siblings, rChild)
208 return t.down(newKey, lChild, siblings, path, l+1)
209 default:
210 return nil, nil, nil, fmt.Errorf("invalid value")
211 }
212}
213
214// downVirtually is used when in a leaf already exists, and a new leaf which
215// shares the path until the existing leaf is being added
216func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
217 newPath []bool, l int) ([][]byte, error) {
218 var err error
219 if l > t.maxLevels-1 {
220 return nil, fmt.Errorf("max virtual level %d", l)
221 }
222
223 if oldPath[l] == newPath[l] {
224 emptyKey := make([]byte, t.hashFunction.Len()) // empty
225 siblings = append(siblings, emptyKey)
226
227 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, l+1)
228 if err != nil {
229 return nil, err
230 }
231 return siblings, nil
232 }
233 // reached the divergence
234 siblings = append(siblings, oldKey)
235
236 return siblings, nil
237}
238
239// up goes up recursively updating the intermediate nodes
240func (t *Tree) up(tx db.Tx, key []byte, siblings [][]byte, path []bool, l int) ([]byte, error) {
241 var k, v []byte
242 var err error
243 if path[l] {
244 k, v, err = t.newIntermediate(siblings[l], key)
245 if err != nil {
246 return nil, err
247 }
248 } else {
249 k, v, err = t.newIntermediate(key, siblings[l])
250 if err != nil {
251 return nil, err
252 }
253 }
254 // store k-v to db
255 err = tx.Put(k, v)
256 if err != nil {
257 return nil, err
258 }
259
260 if l == 0 {
261 // reached the root
262 return k, nil
263 }
264
265 return t.up(tx, k, siblings, path, l-1)
266}
267
268func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
269 leafKey, err := t.hashFunction.Hash(k, v, []byte{1})
270 if err != nil {
271 return nil, nil, err
272 }
273 var leafValue []byte
274 leafValue = append(leafValue, byte(1))
275 leafValue = append(leafValue, byte(len(k)))
276 leafValue = append(leafValue, k...)
277 leafValue = append(leafValue, v...)
278 return leafKey, leafValue, nil
279}
280
281func readLeafValue(b []byte) ([]byte, []byte) {
282 if len(b) < PrefixValueLen {
283 return []byte{}, []byte{}
284 }
285
286 kLen := b[1]
287 if len(b) < PrefixValueLen+int(kLen) {
288 return []byte{}, []byte{}
289 }
290 k := b[PrefixValueLen : PrefixValueLen+kLen]
291 v := b[PrefixValueLen+kLen:]
292 return k, v
293}
294
295func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
296 b := make([]byte, PrefixValueLen+t.hashFunction.Len()*2)
297 b[0] = 2
298 b[1] = byte(len(l))
299 copy(b[PrefixValueLen:PrefixValueLen+t.hashFunction.Len()], l)
300 copy(b[PrefixValueLen+t.hashFunction.Len():], r)
301
302 key, err := t.hashFunction.Hash(l, r)
303 if err != nil {
304 return nil, nil, err
305 }
306
307 return key, b, nil
308}
309
310func readIntermediateChilds(b []byte) ([]byte, []byte) {
311 if len(b) < PrefixValueLen {
312 return []byte{}, []byte{}
313 }
314
315 lLen := b[1]
316 if len(b) < PrefixValueLen+int(lLen) {
317 return []byte{}, []byte{}
318 }
319 l := b[PrefixValueLen : PrefixValueLen+lLen]
320 r := b[PrefixValueLen+lLen:]
321 return l, r
322}
323
324func getPath(numLevels int, k []byte) []bool {
325 path := make([]bool, numLevels)
326 for n := 0; n < numLevels; n++ {
327 path[n] = k[n/8]&(1<<(n%8)) != 0
328 }
329 return path
330}
331
332// GenProof generates a MerkleTree proof for the given key. If the key exists in
333// the Tree, the proof will be of existence, if the key does not exist in the
334// tree, the proof will be of non-existence.
335func (t *Tree) GenProof(k, v []byte) ([]byte, error) {
336 // unimplemented
337 return nil, fmt.Errorf("unimplemented")
338}
339
340// Get returns the value for a given key
341func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
342 // unimplemented
343 return nil, nil, fmt.Errorf("unimplemented")
344}
345
346// CheckProof verifies the given proof
347func CheckProof(k, v, root, mproof []byte) (bool, error) {
348 // unimplemented
349 return false, fmt.Errorf("unimplemented")
350}
351
352// TODO method to export & import the full Tree without values