You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

1581 lines
43 KiB

3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
  1. /*
  2. Package arbo implements a Merkle Tree compatible with the circomlib
  3. implementation of the MerkleTree, following the specification from
  4. https://docs.iden3.io/publications/pdfs/Merkle-Tree.pdf and
  5. https://eprint.iacr.org/2018/955.
  6. Allows to define which hash function to use. So for example, when working with
  7. zkSnarks the Poseidon hash function can be used, but when not, it can be used
  8. the Blake2b hash function, which has much faster computation time.
  9. */
  10. package arbo
  11. import (
  12. "bufio"
  13. "bytes"
  14. "encoding/binary"
  15. "encoding/hex"
  16. "fmt"
  17. "io"
  18. "math"
  19. "runtime"
  20. "sync"
  21. "go.vocdoni.io/dvote/db"
  22. )
  23. const (
  24. // PrefixValueLen defines the bytes-prefix length used for the Value
  25. // bytes representation stored in the db
  26. PrefixValueLen = 2
  27. // PrefixValueEmpty is used for the first byte of a Value to indicate
  28. // that is an Empty value
  29. PrefixValueEmpty = 0
  30. // PrefixValueLeaf is used for the first byte of a Value to indicate
  31. // that is a Leaf value
  32. PrefixValueLeaf = 1
  33. // PrefixValueIntermediate is used for the first byte of a Value to
  34. // indicate that is a Intermediate value
  35. PrefixValueIntermediate = 2
  36. // nChars is used to crop the Graphviz nodes labels
  37. nChars = 4
  38. // maxUint8 is the max size of key length
  39. maxUint8 = int(^uint8(0)) // 2**8 -1
  40. // maxUint16 is the max size of value length
  41. maxUint16 = int(^uint16(0)) // 2**16 -1
  42. )
  43. var (
  44. // DefaultThresholdNLeafs defines the threshold number of leafs in the
  45. // tree that determines if AddBatch will work in memory or in disk. It
  46. // is defined when calling NewTree, and if set to 0 it will work always
  47. // in disk.
  48. DefaultThresholdNLeafs = 65536
  49. dbKeyRoot = []byte("root")
  50. dbKeyNLeafs = []byte("nleafs")
  51. emptyValue = []byte{0}
  52. // ErrKeyNotFound is used when a key is not found in the db neither in
  53. // the current db Batch.
  54. ErrKeyNotFound = fmt.Errorf("key not found")
  55. // ErrKeyAlreadyExists is used when trying to add a key as leaf to the
  56. // tree that already exists.
  57. ErrKeyAlreadyExists = fmt.Errorf("key already exists")
  58. // ErrInvalidValuePrefix is used when going down into the tree, a value
  59. // is read from the db and has an unrecognized prefix.
  60. ErrInvalidValuePrefix = fmt.Errorf("invalid value prefix")
  61. // ErrDBNoTx is used when trying to use Tree.dbPut but Tree.dbBatch==nil
  62. ErrDBNoTx = fmt.Errorf("dbPut error: no db Batch")
  63. // ErrMaxLevel indicates when going down into the tree, the max level is
  64. // reached
  65. ErrMaxLevel = fmt.Errorf("max level reached")
  66. // ErrMaxVirtualLevel indicates when going down into the tree, the max
  67. // virtual level is reached
  68. ErrMaxVirtualLevel = fmt.Errorf("max virtual level reached")
  69. // ErrSnapshotNotEditable indicates when the tree is a non writable
  70. // snapshot, thus can not be modified
  71. ErrSnapshotNotEditable = fmt.Errorf("snapshot tree can not be edited")
  72. // ErrTreeNotEmpty indicates when the tree was expected to be empty and
  73. // it is not
  74. ErrTreeNotEmpty = fmt.Errorf("tree is not empty")
  75. )
  76. // Tree defines the struct that implements the MerkleTree functionalities
  77. type Tree struct {
  78. sync.Mutex
  79. db db.Database
  80. maxLevels int
  81. // thresholdNLeafs defines the threshold number of leafs in the tree
  82. // that determines if AddBatch will work in memory or in disk. It is
  83. // defined when calling NewTree, and if set to 0 it will work always in
  84. // disk.
  85. thresholdNLeafs int
  86. snapshotRoot []byte
  87. hashFunction HashFunction
  88. // TODO in the methods that use it, check if emptyHash param is len>0
  89. // (check if it has been initialized)
  90. emptyHash []byte
  91. dbg *dbgStats
  92. }
  93. // Config defines the configuration for calling NewTree & NewTreeWithTx methods
  94. type Config struct {
  95. Database db.Database
  96. MaxLevels int
  97. ThresholdNLeafs int
  98. HashFunction HashFunction
  99. }
  100. // NewTree returns a new Tree, if there is a Tree still in the given database, it
  101. // will load it.
  102. func NewTree(cfg Config) (*Tree, error) {
  103. wTx := cfg.Database.WriteTx()
  104. defer wTx.Discard()
  105. t, err := NewTreeWithTx(wTx, cfg)
  106. if err != nil {
  107. return nil, err
  108. }
  109. if err = wTx.Commit(); err != nil {
  110. return nil, err
  111. }
  112. return t, nil
  113. }
  114. // NewTreeWithTx returns a new Tree using the given db.WriteTx, which will not
  115. // be ccommited inside this method, if there is a Tree still in the given
  116. // database, it will load it.
  117. func NewTreeWithTx(wTx db.WriteTx, cfg Config) (*Tree, error) {
  118. // if thresholdNLeafs is set to 0, use the DefaultThresholdNLeafs
  119. if cfg.ThresholdNLeafs == 0 {
  120. cfg.ThresholdNLeafs = DefaultThresholdNLeafs
  121. }
  122. t := Tree{db: cfg.Database, maxLevels: cfg.MaxLevels,
  123. thresholdNLeafs: cfg.ThresholdNLeafs, hashFunction: cfg.HashFunction}
  124. t.emptyHash = make([]byte, t.hashFunction.Len()) // empty
  125. _, err := wTx.Get(dbKeyRoot)
  126. if err == db.ErrKeyNotFound {
  127. // store new root 0 (empty)
  128. if err = wTx.Set(dbKeyRoot, t.emptyHash); err != nil {
  129. return nil, err
  130. }
  131. if err = t.setNLeafs(wTx, 0); err != nil {
  132. return nil, err
  133. }
  134. return &t, nil
  135. } else if err != nil {
  136. return nil, err
  137. }
  138. return &t, nil
  139. }
  140. // Root returns the root of the Tree
  141. func (t *Tree) Root() ([]byte, error) {
  142. rTx := t.db.ReadTx()
  143. defer rTx.Discard()
  144. return t.RootWithTx(rTx)
  145. }
  146. // RootWithTx returns the root of the Tree using the given db.ReadTx
  147. func (t *Tree) RootWithTx(rTx db.ReadTx) ([]byte, error) {
  148. // if snapshotRoot is defined, means that the tree is a snapshot, and
  149. // the root is not obtained from the db, but from the snapshotRoot
  150. // parameter
  151. if t.snapshotRoot != nil {
  152. return t.snapshotRoot, nil
  153. }
  154. // get db root
  155. return rTx.Get(dbKeyRoot)
  156. }
  157. func (t *Tree) setRoot(wTx db.WriteTx, root []byte) error {
  158. return wTx.Set(dbKeyRoot, root)
  159. }
  160. // HashFunction returns Tree.hashFunction
  161. func (t *Tree) HashFunction() HashFunction {
  162. return t.hashFunction
  163. }
  164. // editable returns true if the tree is editable, and false when is not
  165. // editable (because is a snapshot tree)
  166. func (t *Tree) editable() bool {
  167. return t.snapshotRoot == nil
  168. }
  169. // Invalid is used when a key-value can not be added trough AddBatch, and
  170. // contains the index of the key-value and the error.
  171. type Invalid struct {
  172. Index int
  173. Error error
  174. }
  175. // AddBatch adds a batch of key-values to the Tree. Returns an array containing
  176. // the indexes of the keys failed to add. Supports empty values as input
  177. // parameters, which is equivalent to 0 valued byte array.
  178. func (t *Tree) AddBatch(keys, values [][]byte) ([]Invalid, error) {
  179. wTx := t.db.WriteTx()
  180. defer wTx.Discard()
  181. invalids, err := t.AddBatchWithTx(wTx, keys, values)
  182. if err != nil {
  183. return invalids, err
  184. }
  185. return invalids, wTx.Commit()
  186. }
  187. // AddBatchWithTx does the same than the AddBatch method, but allowing to pass
  188. // the db.WriteTx that is used. The db.WriteTx will not be committed inside
  189. // this method.
  190. func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
  191. t.Lock()
  192. defer t.Unlock()
  193. if !t.editable() {
  194. return nil, ErrSnapshotNotEditable
  195. }
  196. e := []byte{}
  197. // equal the number of keys & values
  198. if len(keys) > len(values) {
  199. // add missing values
  200. for i := len(values); i < len(keys); i++ {
  201. values = append(values, e)
  202. }
  203. } else if len(keys) < len(values) {
  204. // crop extra values
  205. values = values[:len(keys)]
  206. }
  207. nLeafs, err := t.GetNLeafsWithTx(wTx)
  208. if err != nil {
  209. return nil, err
  210. }
  211. if nLeafs > t.thresholdNLeafs {
  212. return t.addBatchInDisk(wTx, keys, values)
  213. }
  214. return t.addBatchInMemory(wTx, keys, values)
  215. }
  216. func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
  217. nCPU := flp2(runtime.NumCPU())
  218. if nCPU == 1 || len(keys) < nCPU {
  219. var invalids []Invalid
  220. for i := 0; i < len(keys); i++ {
  221. if err := t.addWithTx(wTx, keys[i], values[i]); err != nil {
  222. invalids = append(invalids, Invalid{i, err})
  223. }
  224. }
  225. return invalids, nil
  226. }
  227. kvs, invalids, err := keysValuesToKvs(t.maxLevels, keys, values)
  228. if err != nil {
  229. return nil, err
  230. }
  231. buckets := splitInBuckets(kvs, nCPU)
  232. root, err := t.RootWithTx(wTx)
  233. if err != nil {
  234. return nil, err
  235. }
  236. l := int(math.Log2(float64(nCPU)))
  237. subRoots, err := t.getSubRootsAtLevel(wTx, root, l+1)
  238. if err != nil {
  239. return nil, err
  240. }
  241. if len(subRoots) != nCPU {
  242. // Already populated Tree but Unbalanced.
  243. // add one key at each bucket, and then continue with the flow
  244. for i := 0; i < len(buckets); i++ {
  245. // add one leaf of the bucket, if there is an error when
  246. // adding the k-v, try to add the next one of the bucket
  247. // (until one is added)
  248. inserted := -1
  249. for j := 0; j < len(buckets[i]); j++ {
  250. if newRoot, err := t.add(wTx, root, 0,
  251. buckets[i][j].k, buckets[i][j].v); err == nil {
  252. inserted = j
  253. root = newRoot
  254. break
  255. }
  256. }
  257. // remove the inserted element from buckets[i]
  258. if inserted != -1 {
  259. buckets[i] = append(buckets[i][:inserted], buckets[i][inserted+1:]...)
  260. }
  261. }
  262. subRoots, err = t.getSubRootsAtLevel(wTx, root, l+1)
  263. if err != nil {
  264. return nil, err
  265. }
  266. }
  267. if len(subRoots) != nCPU {
  268. return nil, fmt.Errorf("This error should not be reached."+
  269. " len(subRoots) != nCPU, len(subRoots)=%d, nCPU=%d."+
  270. " Please report it in a new issue:"+
  271. " https://github.com/vocdoni/arbo/issues/new", len(subRoots), nCPU)
  272. }
  273. invalidsInBucket := make([][]Invalid, nCPU)
  274. txs := make([]db.WriteTx, nCPU)
  275. for i := 0; i < nCPU; i++ {
  276. txs[i] = t.db.WriteTx()
  277. err := txs[i].Apply(wTx)
  278. if err != nil {
  279. return nil, err
  280. }
  281. }
  282. var wg sync.WaitGroup
  283. wg.Add(nCPU)
  284. for i := 0; i < nCPU; i++ {
  285. go func(cpu int) {
  286. // use different wTx for each cpu, after once all
  287. // are done, iter over the cpuWTxs and copy their
  288. // content into the main wTx
  289. for j := 0; j < len(buckets[cpu]); j++ {
  290. newSubRoot, err := t.add(txs[cpu], subRoots[cpu],
  291. l, buckets[cpu][j].k, buckets[cpu][j].v)
  292. if err != nil {
  293. invalidsInBucket[cpu] = append(invalidsInBucket[cpu],
  294. Invalid{buckets[cpu][j].pos, err})
  295. continue
  296. }
  297. // if there has not been errors, set the new subRoots[cpu]
  298. subRoots[cpu] = newSubRoot
  299. }
  300. wg.Done()
  301. }(i)
  302. }
  303. wg.Wait()
  304. for i := 0; i < nCPU; i++ {
  305. if err := wTx.Apply(txs[i]); err != nil {
  306. return nil, err
  307. }
  308. txs[i].Discard()
  309. }
  310. for i := 0; i < len(invalidsInBucket); i++ {
  311. invalids = append(invalids, invalidsInBucket[i]...)
  312. }
  313. newRoot, err := t.upFromSubRoots(wTx, subRoots)
  314. if err != nil {
  315. return nil, err
  316. }
  317. // update dbKeyNLeafs
  318. if err := t.SetRootWithTx(wTx, newRoot); err != nil {
  319. return nil, err
  320. }
  321. // update nLeafs
  322. if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
  323. return nil, err
  324. }
  325. return invalids, nil
  326. }
  327. func (t *Tree) upFromSubRoots(wTx db.WriteTx, subRoots [][]byte) ([]byte, error) {
  328. // is a method of Tree just to get access to t.hashFunction and
  329. // t.emptyHash.
  330. // go up from subRoots to up, storing nodes in the given WriteTx
  331. // once up at the root, store it in the WriteTx using the dbKeyRoot
  332. if len(subRoots) == 1 {
  333. return subRoots[0], nil
  334. }
  335. // get the subRoots values to know the node types of each subRoot
  336. nodeTypes := make([]byte, len(subRoots))
  337. for i := 0; i < len(subRoots); i++ {
  338. if bytes.Equal(subRoots[i], t.emptyHash) {
  339. nodeTypes[i] = PrefixValueEmpty
  340. continue
  341. }
  342. v, err := wTx.Get(subRoots[i])
  343. if err != nil {
  344. return nil, err
  345. }
  346. nodeTypes[i] = v[0]
  347. }
  348. var newSubRoots [][]byte
  349. for i := 0; i < len(subRoots); i += 2 {
  350. if (bytes.Equal(subRoots[i], t.emptyHash) && bytes.Equal(subRoots[i+1], t.emptyHash)) ||
  351. (nodeTypes[i] == PrefixValueLeaf && bytes.Equal(subRoots[i+1], t.emptyHash)) {
  352. // when both sub nodes are empty, the parent is also empty
  353. // or
  354. // when 1st sub node is a leaf but the 2nd is empty, the
  355. // leaf is used as 'parent'
  356. newSubRoots = append(newSubRoots, subRoots[i])
  357. continue
  358. }
  359. if bytes.Equal(subRoots[i], t.emptyHash) && nodeTypes[i+1] == PrefixValueLeaf {
  360. // when 2nd sub node is a leaf but the 1st is empty,
  361. // the leaf is used as 'parent'
  362. newSubRoots = append(newSubRoots, subRoots[i+1])
  363. continue
  364. }
  365. k, v, err := t.newIntermediate(subRoots[i], subRoots[i+1])
  366. if err != nil {
  367. return nil, err
  368. }
  369. // store k-v to db
  370. if err = wTx.Set(k, v); err != nil {
  371. return nil, err
  372. }
  373. newSubRoots = append(newSubRoots, k)
  374. }
  375. return t.upFromSubRoots(wTx, newSubRoots)
  376. }
  377. func (t *Tree) getSubRootsAtLevel(rTx db.ReadTx, root []byte, l int) ([][]byte, error) {
  378. // go at level l and return each node key, where each node key is the
  379. // subRoot of the subTree that starts there
  380. var subRoots [][]byte
  381. err := t.iterWithStop(rTx, root, 0, func(currLvl int, k, v []byte) bool {
  382. if currLvl == l && !bytes.Equal(k, t.emptyHash) {
  383. subRoots = append(subRoots, k)
  384. }
  385. if currLvl >= l {
  386. return true // to stop the iter from going down
  387. }
  388. return false
  389. })
  390. return subRoots, err
  391. }
  392. func (t *Tree) addBatchInMemory(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
  393. vt, err := t.loadVT(wTx)
  394. if err != nil {
  395. return nil, err
  396. }
  397. invalids, err := vt.addBatch(keys, values)
  398. if err != nil {
  399. return nil, err
  400. }
  401. // once the VirtualTree is build, compute the hashes
  402. pairs, err := vt.computeHashes()
  403. if err != nil {
  404. // currently invalids in computeHashes are not counted,
  405. // but should not be needed, as if there is an error there is
  406. // nothing stored in the db and the error is returned
  407. return nil, err
  408. }
  409. // store pairs in db
  410. for i := 0; i < len(pairs); i++ {
  411. if err := wTx.Set(pairs[i][0], pairs[i][1]); err != nil {
  412. return nil, err
  413. }
  414. }
  415. // store root (from the vt) to db
  416. if vt.root != nil {
  417. if err := wTx.Set(dbKeyRoot, vt.root.h); err != nil {
  418. return nil, err
  419. }
  420. }
  421. // update nLeafs
  422. if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
  423. return nil, err
  424. }
  425. return invalids, nil
  426. }
  427. // loadVT loads a new virtual tree (vt) from the current Tree, which contains
  428. // the same leafs.
  429. func (t *Tree) loadVT(rTx db.ReadTx) (vt, error) {
  430. vt := newVT(t.maxLevels, t.hashFunction)
  431. vt.params.dbg = t.dbg
  432. var callbackErr error
  433. err := t.IterateWithStopWithTx(rTx, nil, func(_ int, k, v []byte) bool {
  434. if v[0] != PrefixValueLeaf {
  435. return false
  436. }
  437. leafK, leafV := ReadLeafValue(v)
  438. if err := vt.add(0, leafK, leafV); err != nil {
  439. callbackErr = err
  440. return true
  441. }
  442. return false
  443. })
  444. if callbackErr != nil {
  445. return vt, callbackErr
  446. }
  447. return vt, err
  448. }
  449. // Add inserts the key-value into the Tree. If the inputs come from a
  450. // *big.Int, is expected that are represented by a Little-Endian byte array
  451. // (for circom compatibility).
  452. func (t *Tree) Add(k, v []byte) error {
  453. wTx := t.db.WriteTx()
  454. defer wTx.Discard()
  455. if err := t.AddWithTx(wTx, k, v); err != nil {
  456. return err
  457. }
  458. return wTx.Commit()
  459. }
  460. // AddWithTx does the same than the Add method, but allowing to pass the
  461. // db.WriteTx that is used. The db.WriteTx will not be committed inside this
  462. // method.
  463. func (t *Tree) AddWithTx(wTx db.WriteTx, k, v []byte) error {
  464. t.Lock()
  465. defer t.Unlock()
  466. if !t.editable() {
  467. return ErrSnapshotNotEditable
  468. }
  469. return t.addWithTx(wTx, k, v)
  470. }
  471. // warning: addWithTx does not use the Tree mutex, the mutex is responsibility
  472. // of the methods calling this method, and same with t.editable().
  473. func (t *Tree) addWithTx(wTx db.WriteTx, k, v []byte) error {
  474. root, err := t.RootWithTx(wTx)
  475. if err != nil {
  476. return err
  477. }
  478. root, err = t.add(wTx, root, 0, k, v) // add from level 0
  479. if err != nil {
  480. return err
  481. }
  482. // store root to db
  483. if err := t.setRoot(wTx, root); err != nil {
  484. return err
  485. }
  486. // update nLeafs
  487. if err = t.incNLeafs(wTx, 1); err != nil {
  488. return err
  489. }
  490. return nil
  491. }
  492. // keyPathFromKey returns the keyPath and checks that the key is not bigger
  493. // than maximum key length for the tree maxLevels size.
  494. // This is because if the key bits length is bigger than the maxLevels of the
  495. // tree, two different keys that their difference is at the end, will collision
  496. // in the same leaf of the tree (at the max depth).
  497. func keyPathFromKey(maxLevels int, k []byte) ([]byte, error) {
  498. maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
  499. if len(k) > maxKeyLen {
  500. return nil, fmt.Errorf("len(k) can not be bigger than ceil(maxLevels/8), where"+
  501. " len(k): %d, maxLevels: %d, max key len=ceil(maxLevels/8): %d. Might need"+
  502. " a bigger tree depth (maxLevels>=%d) in order to input keys of length %d",
  503. len(k), maxLevels, maxKeyLen, len(k)*8, len(k)) //nolint:gomnd
  504. }
  505. keyPath := make([]byte, maxKeyLen) //nolint:gomnd
  506. copy(keyPath[:], k)
  507. return keyPath, nil
  508. }
  509. // checkKeyValueLen checks the key length and value length. This method is used
  510. // when adding single leafs and also when adding a batch. The limits of lengths
  511. // used are derived from the encoding of tree dumps: 1 byte to define the
  512. // length of the keys (2^8-1 bytes length)), and 2 bytes to define the length
  513. // of the values (2^16-1 bytes length).
  514. func checkKeyValueLen(k, v []byte) error {
  515. if len(k) > maxUint8 {
  516. return fmt.Errorf("len(k)=%v, can not be bigger than %v",
  517. len(k), maxUint8)
  518. }
  519. if len(v) > maxUint16 {
  520. return fmt.Errorf("len(v)=%v, can not be bigger than %v",
  521. len(v), maxUint16)
  522. }
  523. return nil
  524. }
  525. func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) {
  526. if err := checkKeyValueLen(k, v); err != nil {
  527. return nil, err
  528. }
  529. keyPath, err := keyPathFromKey(t.maxLevels, k)
  530. if err != nil {
  531. return nil, err
  532. }
  533. path := getPath(t.maxLevels, keyPath)
  534. // go down to the leaf
  535. var siblings [][]byte
  536. _, _, siblings, err = t.down(wTx, k, root, siblings, path, fromLvl, false)
  537. if err != nil {
  538. return nil, err
  539. }
  540. leafKey, leafValue, err := t.newLeafValue(k, v)
  541. if err != nil {
  542. return nil, err
  543. }
  544. if err := wTx.Set(leafKey, leafValue); err != nil {
  545. return nil, err
  546. }
  547. // go up to the root
  548. if len(siblings) == 0 {
  549. // return the leafKey as root
  550. return leafKey, nil
  551. }
  552. root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, fromLvl)
  553. if err != nil {
  554. return nil, err
  555. }
  556. return root, nil
  557. }
  558. // down goes down to the leaf recursively
  559. func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte,
  560. path []bool, currLvl int, getLeaf bool) (
  561. []byte, []byte, [][]byte, error) {
  562. if currLvl > t.maxLevels {
  563. return nil, nil, nil, ErrMaxLevel
  564. }
  565. var err error
  566. var currValue []byte
  567. if bytes.Equal(currKey, t.emptyHash) {
  568. // empty value
  569. return currKey, emptyValue, siblings, nil
  570. }
  571. currValue, err = rTx.Get(currKey)
  572. if err != nil {
  573. return nil, nil, nil, err
  574. }
  575. switch currValue[0] {
  576. case PrefixValueEmpty: // empty
  577. fmt.Printf("newKey: %s, currKey: %s, currLvl: %d, currValue: %s\n",
  578. hex.EncodeToString(newKey), hex.EncodeToString(currKey),
  579. currLvl, hex.EncodeToString(currValue))
  580. panic("This point should not be reached, as the 'if currKey==t.emptyHash'" +
  581. " above should avoid reaching this point. This panic is temporary" +
  582. " for reporting purposes, will be deleted in future versions." +
  583. " Please paste this log (including the previous log lines) in a" +
  584. " new issue: https://github.com/vocdoni/arbo/issues/new") // TMP
  585. case PrefixValueLeaf: // leaf
  586. if !bytes.Equal(currValue, emptyValue) {
  587. if getLeaf {
  588. return currKey, currValue, siblings, nil
  589. }
  590. oldLeafKey, _ := ReadLeafValue(currValue)
  591. if bytes.Equal(newKey, oldLeafKey) {
  592. return nil, nil, nil, ErrKeyAlreadyExists
  593. }
  594. oldLeafKeyFull, err := keyPathFromKey(t.maxLevels, oldLeafKey)
  595. if err != nil {
  596. return nil, nil, nil, err
  597. }
  598. // if currKey is already used, go down until paths diverge
  599. oldPath := getPath(t.maxLevels, oldLeafKeyFull)
  600. siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
  601. if err != nil {
  602. return nil, nil, nil, err
  603. }
  604. }
  605. return currKey, currValue, siblings, nil
  606. case PrefixValueIntermediate: // intermediate
  607. if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
  608. return nil, nil, nil,
  609. fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
  610. PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
  611. }
  612. // collect siblings while going down
  613. if path[currLvl] {
  614. // right
  615. lChild, rChild := ReadIntermediateChilds(currValue)
  616. siblings = append(siblings, lChild)
  617. return t.down(rTx, newKey, rChild, siblings, path, currLvl+1, getLeaf)
  618. }
  619. // left
  620. lChild, rChild := ReadIntermediateChilds(currValue)
  621. siblings = append(siblings, rChild)
  622. return t.down(rTx, newKey, lChild, siblings, path, currLvl+1, getLeaf)
  623. default:
  624. return nil, nil, nil, ErrInvalidValuePrefix
  625. }
  626. }
  627. // downVirtually is used when in a leaf already exists, and a new leaf which
  628. // shares the path until the existing leaf is being added
  629. func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
  630. newPath []bool, currLvl int) ([][]byte, error) {
  631. var err error
  632. if currLvl > t.maxLevels-1 {
  633. return nil, ErrMaxVirtualLevel
  634. }
  635. if oldPath[currLvl] == newPath[currLvl] {
  636. siblings = append(siblings, t.emptyHash)
  637. siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
  638. if err != nil {
  639. return nil, err
  640. }
  641. return siblings, nil
  642. }
  643. // reached the divergence
  644. siblings = append(siblings, oldKey)
  645. return siblings, nil
  646. }
  647. // up goes up recursively updating the intermediate nodes
  648. func (t *Tree) up(wTx db.WriteTx, key []byte, siblings [][]byte, path []bool,
  649. currLvl, toLvl int) ([]byte, error) {
  650. var k, v []byte
  651. var err error
  652. if path[currLvl+toLvl] {
  653. k, v, err = t.newIntermediate(siblings[currLvl], key)
  654. if err != nil {
  655. return nil, err
  656. }
  657. } else {
  658. k, v, err = t.newIntermediate(key, siblings[currLvl])
  659. if err != nil {
  660. return nil, err
  661. }
  662. }
  663. // store k-v to db
  664. if err = wTx.Set(k, v); err != nil {
  665. return nil, err
  666. }
  667. if currLvl == 0 {
  668. // reached the root
  669. return k, nil
  670. }
  671. return t.up(wTx, k, siblings, path, currLvl-1, toLvl)
  672. }
  673. func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
  674. t.dbg.incHash()
  675. return newLeafValue(t.hashFunction, k, v)
  676. }
  677. // newLeafValue takes a key & value from a leaf, and computes the leaf hash,
  678. // which is used as the leaf key. And the value is the concatenation of the
  679. // inputed key & value. The output of this function is used as key-value to
  680. // store the leaf in the DB.
  681. // [ 1 byte | 1 byte | N bytes | M bytes ]
  682. // [ type of node | length of key | key | value ]
  683. func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
  684. if err := checkKeyValueLen(k, v); err != nil {
  685. return nil, nil, err
  686. }
  687. leafKey, err := hashFunc.Hash(k, v, []byte{1})
  688. if err != nil {
  689. return nil, nil, err
  690. }
  691. var leafValue []byte
  692. leafValue = append(leafValue, byte(PrefixValueLeaf))
  693. leafValue = append(leafValue, byte(len(k)))
  694. leafValue = append(leafValue, k...)
  695. leafValue = append(leafValue, v...)
  696. return leafKey, leafValue, nil
  697. }
  698. // ReadLeafValue reads from a byte array the leaf key & value
  699. func ReadLeafValue(b []byte) ([]byte, []byte) {
  700. if len(b) < PrefixValueLen {
  701. return []byte{}, []byte{}
  702. }
  703. kLen := b[1]
  704. if len(b) < PrefixValueLen+int(kLen) {
  705. return []byte{}, []byte{}
  706. }
  707. k := b[PrefixValueLen : PrefixValueLen+kLen]
  708. v := b[PrefixValueLen+kLen:]
  709. return k, v
  710. }
  711. func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
  712. t.dbg.incHash()
  713. return newIntermediate(t.hashFunction, l, r)
  714. }
  715. // newIntermediate takes the left & right keys of a intermediate node, and
  716. // computes its hash. Returns the hash of the node, which is the node key, and a
  717. // byte array that contains the value (which contains the left & right child
  718. // keys) to store in the DB.
  719. // [ 1 byte | 1 byte | N bytes | N bytes ]
  720. // [ type of node | length of left key | left key | right key ]
  721. func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
  722. b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
  723. b[0] = PrefixValueIntermediate
  724. if len(l) > maxUint8 {
  725. return nil, nil, fmt.Errorf("newIntermediate: len(l) > %v", maxUint8)
  726. }
  727. b[1] = byte(len(l))
  728. copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
  729. copy(b[PrefixValueLen+hashFunc.Len():], r)
  730. key, err := hashFunc.Hash(l, r)
  731. if err != nil {
  732. return nil, nil, err
  733. }
  734. return key, b, nil
  735. }
  736. // ReadIntermediateChilds reads from a byte array the two childs keys
  737. func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
  738. if len(b) < PrefixValueLen {
  739. return []byte{}, []byte{}
  740. }
  741. lLen := b[1]
  742. if len(b) < PrefixValueLen+int(lLen) {
  743. return []byte{}, []byte{}
  744. }
  745. l := b[PrefixValueLen : PrefixValueLen+lLen]
  746. r := b[PrefixValueLen+lLen:]
  747. return l, r
  748. }
  749. func getPath(numLevels int, k []byte) []bool {
  750. path := make([]bool, numLevels)
  751. for n := 0; n < numLevels; n++ {
  752. path[n] = k[n/8]&(1<<(n%8)) != 0
  753. }
  754. return path
  755. }
  756. // Update updates the value for a given existing key. If the given key does not
  757. // exist, returns an error.
  758. func (t *Tree) Update(k, v []byte) error {
  759. wTx := t.db.WriteTx()
  760. defer wTx.Discard()
  761. if err := t.UpdateWithTx(wTx, k, v); err != nil {
  762. return err
  763. }
  764. return wTx.Commit()
  765. }
  766. // UpdateWithTx does the same than the Update method, but allowing to pass the
  767. // db.WriteTx that is used. The db.WriteTx will not be committed inside this
  768. // method.
  769. func (t *Tree) UpdateWithTx(wTx db.WriteTx, k, v []byte) error {
  770. t.Lock()
  771. defer t.Unlock()
  772. if !t.editable() {
  773. return ErrSnapshotNotEditable
  774. }
  775. keyPath, err := keyPathFromKey(t.maxLevels, k)
  776. if err != nil {
  777. return err
  778. }
  779. path := getPath(t.maxLevels, keyPath)
  780. root, err := t.RootWithTx(wTx)
  781. if err != nil {
  782. return err
  783. }
  784. var siblings [][]byte
  785. _, valueAtBottom, siblings, err := t.down(wTx, k, root, siblings, path, 0, true)
  786. if err != nil {
  787. return err
  788. }
  789. oldKey, _ := ReadLeafValue(valueAtBottom)
  790. if !bytes.Equal(oldKey, k) {
  791. return ErrKeyNotFound
  792. }
  793. leafKey, leafValue, err := t.newLeafValue(k, v)
  794. if err != nil {
  795. return err
  796. }
  797. if err := wTx.Set(leafKey, leafValue); err != nil {
  798. return err
  799. }
  800. // go up to the root
  801. if len(siblings) == 0 {
  802. return t.setRoot(wTx, leafKey)
  803. }
  804. root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, 0)
  805. if err != nil {
  806. return err
  807. }
  808. // store root to db
  809. if err := t.setRoot(wTx, root); err != nil {
  810. return err
  811. }
  812. return nil
  813. }
  814. // GenProof generates a MerkleTree proof for the given key. The leaf value is
  815. // returned, together with the packed siblings of the proof, and a boolean
  816. // parameter that indicates if the proof is of existence (true) or not (false).
  817. func (t *Tree) GenProof(k []byte) ([]byte, []byte, []byte, bool, error) {
  818. rTx := t.db.ReadTx()
  819. defer rTx.Discard()
  820. return t.GenProofWithTx(rTx, k)
  821. }
  822. // GenProofWithTx does the same than the GenProof method, but allowing to pass
  823. // the db.ReadTx that is used.
  824. func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) {
  825. keyPath, err := keyPathFromKey(t.maxLevels, k)
  826. if err != nil {
  827. return nil, nil, nil, false, err
  828. }
  829. path := getPath(t.maxLevels, keyPath)
  830. root, err := t.RootWithTx(rTx)
  831. if err != nil {
  832. return nil, nil, nil, false, err
  833. }
  834. // go down to the leaf
  835. var siblings [][]byte
  836. _, value, siblings, err := t.down(rTx, k, root, siblings, path, 0, true)
  837. if err != nil {
  838. return nil, nil, nil, false, err
  839. }
  840. s, err := PackSiblings(t.hashFunction, siblings)
  841. if err != nil {
  842. return nil, nil, nil, false, err
  843. }
  844. leafK, leafV := ReadLeafValue(value)
  845. if !bytes.Equal(k, leafK) {
  846. // key not in tree, proof of non-existence
  847. return leafK, leafV, s, false, nil
  848. }
  849. return leafK, leafV, s, true, nil
  850. }
  851. // PackSiblings packs the siblings into a byte array.
  852. // [ 2 byte | 2 byte | L bytes | S * N bytes ]
  853. // [ full length | bitmap length (L) | bitmap | N non-zero siblings ]
  854. // Where the bitmap indicates if the sibling is 0 or a value from the siblings
  855. // array. And S is the size of the output of the hash function used for the
  856. // Tree. The 2 2-byte that define the full length and bitmap length, are
  857. // encoded in little-endian.
  858. func PackSiblings(hashFunc HashFunction, siblings [][]byte) ([]byte, error) {
  859. var b []byte
  860. var bitmap []bool
  861. emptySibling := make([]byte, hashFunc.Len())
  862. for i := 0; i < len(siblings); i++ {
  863. if bytes.Equal(siblings[i], emptySibling) {
  864. bitmap = append(bitmap, false)
  865. } else {
  866. bitmap = append(bitmap, true)
  867. b = append(b, siblings[i]...)
  868. }
  869. }
  870. bitmapBytes := bitmapToBytes(bitmap)
  871. l := len(bitmapBytes)
  872. if l > maxUint16 {
  873. return nil, fmt.Errorf("PackSiblings: bitmapBytes length > %v", maxUint16)
  874. }
  875. fullLen := 4 + l + len(b) //nolint:gomnd
  876. if fullLen > maxUint16 {
  877. return nil, fmt.Errorf("PackSiblings: fullLen > %v", maxUint16)
  878. }
  879. res := make([]byte, fullLen)
  880. binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length
  881. binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length
  882. copy(res[4:4+l], bitmapBytes)
  883. copy(res[4+l:], b)
  884. return res, nil
  885. }
  886. // UnpackSiblings unpacks the siblings from a byte array.
  887. func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
  888. fullLen := binary.LittleEndian.Uint16(b[0:2])
  889. l := binary.LittleEndian.Uint16(b[2:4]) // bitmap bytes length
  890. if len(b) != int(fullLen) {
  891. return nil,
  892. fmt.Errorf("expected len: %d, current len: %d",
  893. fullLen, len(b))
  894. }
  895. bitmapBytes := b[4 : 4+l]
  896. bitmap := bytesToBitmap(bitmapBytes)
  897. siblingsBytes := b[4+l:]
  898. iSibl := 0
  899. emptySibl := make([]byte, hashFunc.Len())
  900. var siblings [][]byte
  901. for i := 0; i < len(bitmap); i++ {
  902. if iSibl >= len(siblingsBytes) {
  903. break
  904. }
  905. if bitmap[i] {
  906. siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
  907. iSibl += hashFunc.Len()
  908. } else {
  909. siblings = append(siblings, emptySibl)
  910. }
  911. }
  912. return siblings, nil
  913. }
  914. func bitmapToBytes(bitmap []bool) []byte {
  915. bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
  916. b := make([]byte, bitmapBytesLen)
  917. for i := 0; i < len(bitmap); i++ {
  918. if bitmap[i] {
  919. b[i/8] |= 1 << (i % 8)
  920. }
  921. }
  922. return b
  923. }
  924. func bytesToBitmap(b []byte) []bool {
  925. var bitmap []bool
  926. for i := 0; i < len(b); i++ {
  927. for j := 0; j < 8; j++ {
  928. bitmap = append(bitmap, b[i]&(1<<j) > 0)
  929. }
  930. }
  931. return bitmap
  932. }
  933. // Get returns the value in the Tree for a given key. If the key is not found,
  934. // will return the error ErrKeyNotFound, and in the leafK & leafV parameters
  935. // will be placed the data found in the tree in the leaf that was on the path
  936. // going to the input key.
  937. func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
  938. rTx := t.db.ReadTx()
  939. defer rTx.Discard()
  940. return t.GetWithTx(rTx, k)
  941. }
  942. // GetWithTx does the same than the Get method, but allowing to pass the
  943. // db.ReadTx that is used. If the key is not found, will return the error
  944. // ErrKeyNotFound, and in the leafK & leafV parameters will be placed the data
  945. // found in the tree in the leaf that was on the path going to the input key.
  946. func (t *Tree) GetWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, error) {
  947. keyPath, err := keyPathFromKey(t.maxLevels, k)
  948. if err != nil {
  949. return nil, nil, err
  950. }
  951. path := getPath(t.maxLevels, keyPath)
  952. root, err := t.RootWithTx(rTx)
  953. if err != nil {
  954. return nil, nil, err
  955. }
  956. // go down to the leaf
  957. var siblings [][]byte
  958. _, value, _, err := t.down(rTx, k, root, siblings, path, 0, true)
  959. if err != nil {
  960. return nil, nil, err
  961. }
  962. leafK, leafV := ReadLeafValue(value)
  963. if !bytes.Equal(k, leafK) {
  964. return leafK, leafV, ErrKeyNotFound
  965. }
  966. return leafK, leafV, nil
  967. }
  968. // CheckProof verifies the given proof. The proof verification depends on the
  969. // HashFunction passed as parameter.
  970. func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
  971. siblings, err := UnpackSiblings(hashFunc, packedSiblings)
  972. if err != nil {
  973. return false, err
  974. }
  975. keyPath := make([]byte, int(math.Ceil(float64(len(siblings))/float64(8)))) //nolint:gomnd
  976. copy(keyPath[:], k)
  977. key, _, err := newLeafValue(hashFunc, k, v)
  978. if err != nil {
  979. return false, err
  980. }
  981. path := getPath(len(siblings), keyPath)
  982. for i := len(siblings) - 1; i >= 0; i-- {
  983. if path[i] {
  984. key, _, err = newIntermediate(hashFunc, siblings[i], key)
  985. if err != nil {
  986. return false, err
  987. }
  988. } else {
  989. key, _, err = newIntermediate(hashFunc, key, siblings[i])
  990. if err != nil {
  991. return false, err
  992. }
  993. }
  994. }
  995. if bytes.Equal(key[:], root) {
  996. return true, nil
  997. }
  998. return false, nil
  999. }
  1000. func (t *Tree) incNLeafs(wTx db.WriteTx, nLeafs int) error {
  1001. oldNLeafs, err := t.GetNLeafsWithTx(wTx)
  1002. if err != nil {
  1003. return err
  1004. }
  1005. newNLeafs := oldNLeafs + nLeafs
  1006. return t.setNLeafs(wTx, newNLeafs)
  1007. }
  1008. func (t *Tree) setNLeafs(wTx db.WriteTx, nLeafs int) error {
  1009. b := make([]byte, 8)
  1010. binary.LittleEndian.PutUint64(b, uint64(nLeafs))
  1011. if err := wTx.Set(dbKeyNLeafs, b); err != nil {
  1012. return err
  1013. }
  1014. return nil
  1015. }
  1016. // GetNLeafs returns the number of Leafs of the Tree.
  1017. func (t *Tree) GetNLeafs() (int, error) {
  1018. rTx := t.db.ReadTx()
  1019. defer rTx.Discard()
  1020. return t.GetNLeafsWithTx(rTx)
  1021. }
  1022. // GetNLeafsWithTx does the same than the GetNLeafs method, but allowing to
  1023. // pass the db.ReadTx that is used.
  1024. func (t *Tree) GetNLeafsWithTx(rTx db.ReadTx) (int, error) {
  1025. b, err := rTx.Get(dbKeyNLeafs)
  1026. if err != nil {
  1027. return 0, err
  1028. }
  1029. nLeafs := binary.LittleEndian.Uint64(b)
  1030. return int(nLeafs), nil
  1031. }
  1032. // SetRoot sets the root to the given root
  1033. func (t *Tree) SetRoot(root []byte) error {
  1034. wTx := t.db.WriteTx()
  1035. defer wTx.Discard()
  1036. if err := t.SetRootWithTx(wTx, root); err != nil {
  1037. return err
  1038. }
  1039. return wTx.Commit()
  1040. }
  1041. // SetRootWithTx sets the root to the given root using the given db.WriteTx
  1042. func (t *Tree) SetRootWithTx(wTx db.WriteTx, root []byte) error {
  1043. if !t.editable() {
  1044. return ErrSnapshotNotEditable
  1045. }
  1046. if root == nil {
  1047. return fmt.Errorf("can not SetRoot with nil root")
  1048. }
  1049. // check that the root exists in the db
  1050. if !bytes.Equal(root, t.emptyHash) {
  1051. if _, err := wTx.Get(root); err == ErrKeyNotFound {
  1052. return fmt.Errorf("can not SetRoot with root %x, as it"+
  1053. " does not exist in the db", root)
  1054. } else if err != nil {
  1055. return err
  1056. }
  1057. }
  1058. return wTx.Set(dbKeyRoot, root)
  1059. }
  1060. // Snapshot returns a read-only copy of the Tree from the given root
  1061. func (t *Tree) Snapshot(fromRoot []byte) (*Tree, error) {
  1062. // allow to define which root to use
  1063. if fromRoot == nil {
  1064. var err error
  1065. fromRoot, err = t.Root()
  1066. if err != nil {
  1067. return nil, err
  1068. }
  1069. }
  1070. rTx := t.db.ReadTx()
  1071. defer rTx.Discard()
  1072. // check that the root exists in the db
  1073. if !bytes.Equal(fromRoot, t.emptyHash) {
  1074. if _, err := rTx.Get(fromRoot); err == ErrKeyNotFound {
  1075. return nil,
  1076. fmt.Errorf("can not do a Snapshot with root %x,"+
  1077. " as it does not exist in the db", fromRoot)
  1078. } else if err != nil {
  1079. return nil, err
  1080. }
  1081. }
  1082. return &Tree{
  1083. db: t.db,
  1084. maxLevels: t.maxLevels,
  1085. snapshotRoot: fromRoot,
  1086. emptyHash: t.emptyHash,
  1087. hashFunction: t.hashFunction,
  1088. dbg: t.dbg,
  1089. }, nil
  1090. }
  1091. // Iterate iterates through the full Tree, executing the given function on each
  1092. // node of the Tree.
  1093. func (t *Tree) Iterate(fromRoot []byte, f func([]byte, []byte)) error {
  1094. rTx := t.db.ReadTx()
  1095. defer rTx.Discard()
  1096. return t.IterateWithTx(rTx, fromRoot, f)
  1097. }
  1098. // IterateWithTx does the same than the Iterate method, but allowing to pass
  1099. // the db.ReadTx that is used.
  1100. func (t *Tree) IterateWithTx(rTx db.ReadTx, fromRoot []byte, f func([]byte, []byte)) error {
  1101. // allow to define which root to use
  1102. if fromRoot == nil {
  1103. var err error
  1104. fromRoot, err = t.RootWithTx(rTx)
  1105. if err != nil {
  1106. return err
  1107. }
  1108. }
  1109. return t.iter(rTx, fromRoot, f)
  1110. }
  1111. // IterateWithStop does the same than Iterate, but with int for the current
  1112. // level, and a boolean parameter used by the passed function, is to indicate to
  1113. // stop iterating on the branch when the method returns 'true'.
  1114. func (t *Tree) IterateWithStop(fromRoot []byte, f func(int, []byte, []byte) bool) error {
  1115. rTx := t.db.ReadTx()
  1116. defer rTx.Discard()
  1117. // allow to define which root to use
  1118. if fromRoot == nil {
  1119. var err error
  1120. fromRoot, err = t.RootWithTx(rTx)
  1121. if err != nil {
  1122. return err
  1123. }
  1124. }
  1125. return t.iterWithStop(rTx, fromRoot, 0, f)
  1126. }
  1127. // IterateWithStopWithTx does the same than the IterateWithStop method, but
  1128. // allowing to pass the db.ReadTx that is used.
  1129. func (t *Tree) IterateWithStopWithTx(rTx db.ReadTx, fromRoot []byte,
  1130. f func(int, []byte, []byte) bool) error {
  1131. // allow to define which root to use
  1132. if fromRoot == nil {
  1133. var err error
  1134. fromRoot, err = t.RootWithTx(rTx)
  1135. if err != nil {
  1136. return err
  1137. }
  1138. }
  1139. return t.iterWithStop(rTx, fromRoot, 0, f)
  1140. }
  1141. func (t *Tree) iterWithStop(rTx db.ReadTx, k []byte, currLevel int,
  1142. f func(int, []byte, []byte) bool) error {
  1143. var v []byte
  1144. var err error
  1145. if bytes.Equal(k, t.emptyHash) {
  1146. v = t.emptyHash
  1147. } else {
  1148. v, err = rTx.Get(k)
  1149. if err != nil {
  1150. return err
  1151. }
  1152. }
  1153. currLevel++
  1154. switch v[0] {
  1155. case PrefixValueEmpty:
  1156. f(currLevel, k, v)
  1157. case PrefixValueLeaf:
  1158. f(currLevel, k, v)
  1159. case PrefixValueIntermediate:
  1160. stop := f(currLevel, k, v)
  1161. if stop {
  1162. return nil
  1163. }
  1164. l, r := ReadIntermediateChilds(v)
  1165. if err = t.iterWithStop(rTx, l, currLevel, f); err != nil {
  1166. return err
  1167. }
  1168. if err = t.iterWithStop(rTx, r, currLevel, f); err != nil {
  1169. return err
  1170. }
  1171. default:
  1172. return ErrInvalidValuePrefix
  1173. }
  1174. return nil
  1175. }
  1176. func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
  1177. f2 := func(currLvl int, k, v []byte) bool {
  1178. f(k, v)
  1179. return false
  1180. }
  1181. return t.iterWithStop(rTx, k, 0, f2)
  1182. }
  1183. // Dump exports all the Tree leafs in a byte array
  1184. func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
  1185. return t.dump(fromRoot, nil)
  1186. }
  1187. // DumpWriter exports all the Tree leafs writing the bytes in the given Writer
  1188. func (t *Tree) DumpWriter(fromRoot []byte, w io.Writer) error {
  1189. _, err := t.dump(fromRoot, w)
  1190. return err
  1191. }
  1192. // dump exports all the Tree leafs. If the given w is nil, it will return a
  1193. // byte array with the dump, if w contains a *bufio.Writer, it will write the
  1194. // dump in w.
  1195. // The format of the dump is the following:
  1196. // Dump length: [ N * (3+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
  1197. // [ 1 byte | 2 byte | S bytes | len(v) bytes ]
  1198. // [ len(k) | len(v) | key | value ]
  1199. // Where S is the size of the output of the hash function used for the Tree.
  1200. func (t *Tree) dump(fromRoot []byte, w io.Writer) ([]byte, error) {
  1201. // allow to define which root to use
  1202. if fromRoot == nil {
  1203. var err error
  1204. fromRoot, err = t.Root()
  1205. if err != nil {
  1206. return nil, err
  1207. }
  1208. }
  1209. // WARNING current encoding only supports keys of 255 bytes and values
  1210. // of 65535 bytes each (due using only 1 and 2 bytes for the length
  1211. // headers). These lengths are checked on leaf addition by the function
  1212. // checkKeyValueLen.
  1213. var b []byte
  1214. var callbackErr error
  1215. err := t.IterateWithStop(fromRoot, func(_ int, k, v []byte) bool {
  1216. if v[0] != PrefixValueLeaf {
  1217. return false
  1218. }
  1219. leafK, leafV := ReadLeafValue(v)
  1220. kv := make([]byte, 3+len(leafK)+len(leafV))
  1221. if len(leafK) > maxUint8 {
  1222. callbackErr = fmt.Errorf("len(leafK) > %v", maxUint8)
  1223. return true
  1224. }
  1225. kv[0] = byte(len(leafK))
  1226. if len(leafV) > maxUint16 {
  1227. callbackErr = fmt.Errorf("len(leafV) > %v", maxUint16)
  1228. return true
  1229. }
  1230. binary.LittleEndian.PutUint16(kv[1:3], uint16(len(leafV)))
  1231. copy(kv[3:3+len(leafK)], leafK)
  1232. copy(kv[3+len(leafK):], leafV)
  1233. if w == nil {
  1234. b = append(b, kv...)
  1235. } else {
  1236. n, err := w.Write(kv)
  1237. if err != nil {
  1238. callbackErr = fmt.Errorf("dump: w.Write, %s", err)
  1239. return true
  1240. }
  1241. if n != len(kv) {
  1242. callbackErr = fmt.Errorf("dump: w.Write n!=len(kv), %s", err)
  1243. return true
  1244. }
  1245. }
  1246. return false
  1247. })
  1248. if callbackErr != nil {
  1249. return nil, callbackErr
  1250. }
  1251. return b, err
  1252. }
  1253. // ImportDump imports the leafs (that have been exported with the Dump method)
  1254. // in the Tree, reading them from the given byte array.
  1255. func (t *Tree) ImportDump(b []byte) error {
  1256. bytesReader := bytes.NewReader(b)
  1257. r := bufio.NewReader(bytesReader)
  1258. return t.ImportDumpReader(r)
  1259. }
  1260. // ImportDumpReader imports the leafs (that have been exported with the Dump
  1261. // method) in the Tree, reading them from the given reader.
  1262. func (t *Tree) ImportDumpReader(r io.Reader) error {
  1263. if !t.editable() {
  1264. return ErrSnapshotNotEditable
  1265. }
  1266. root, err := t.Root()
  1267. if err != nil {
  1268. return err
  1269. }
  1270. if !bytes.Equal(root, t.emptyHash) {
  1271. return ErrTreeNotEmpty
  1272. }
  1273. var keys, values [][]byte
  1274. for {
  1275. l := make([]byte, 3)
  1276. _, err = io.ReadFull(r, l)
  1277. if err == io.EOF {
  1278. break
  1279. } else if err != nil {
  1280. return err
  1281. }
  1282. lenK := int(l[0])
  1283. k := make([]byte, lenK)
  1284. _, err = io.ReadFull(r, k)
  1285. if err != nil {
  1286. return err
  1287. }
  1288. lenV := binary.LittleEndian.Uint16(l[1:3])
  1289. v := make([]byte, lenV)
  1290. _, err = io.ReadFull(r, v)
  1291. if err != nil {
  1292. return err
  1293. }
  1294. keys = append(keys, k)
  1295. values = append(values, v)
  1296. }
  1297. if _, err = t.AddBatch(keys, values); err != nil {
  1298. return err
  1299. }
  1300. return nil
  1301. }
  1302. // Graphviz iterates across the full tree to generate a string Graphviz
  1303. // representation of the tree and writes it to w
  1304. func (t *Tree) Graphviz(w io.Writer, fromRoot []byte) error {
  1305. return t.GraphvizFirstNLevels(w, fromRoot, t.maxLevels)
  1306. }
  1307. // GraphvizFirstNLevels iterates across the first NLevels of the tree to
  1308. // generate a string Graphviz representation of the first NLevels of the tree
  1309. // and writes it to w
  1310. func (t *Tree) GraphvizFirstNLevels(w io.Writer, fromRoot []byte, untilLvl int) error {
  1311. fmt.Fprintf(w, `digraph hierarchy {
  1312. node [fontname=Monospace,fontsize=10,shape=box]
  1313. `)
  1314. rTx := t.db.ReadTx()
  1315. defer rTx.Discard()
  1316. if fromRoot == nil {
  1317. var err error
  1318. fromRoot, err = t.RootWithTx(rTx)
  1319. if err != nil {
  1320. return err
  1321. }
  1322. }
  1323. nEmpties := 0
  1324. err := t.iterWithStop(rTx, fromRoot, 0, func(currLvl int, k, v []byte) bool {
  1325. if currLvl == untilLvl {
  1326. return true // to stop the iter from going down
  1327. }
  1328. switch v[0] {
  1329. case PrefixValueEmpty:
  1330. case PrefixValueLeaf:
  1331. fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
  1332. // key & value from the leaf
  1333. kB, vB := ReadLeafValue(v)
  1334. fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
  1335. hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
  1336. hex.EncodeToString(vB[:nChars]))
  1337. fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
  1338. hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
  1339. case PrefixValueIntermediate:
  1340. l, r := ReadIntermediateChilds(v)
  1341. lStr := hex.EncodeToString(l[:nChars])
  1342. rStr := hex.EncodeToString(r[:nChars])
  1343. eStr := ""
  1344. if bytes.Equal(l, t.emptyHash) {
  1345. lStr = fmt.Sprintf("empty%v", nEmpties)
  1346. eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
  1347. lStr)
  1348. nEmpties++
  1349. }
  1350. if bytes.Equal(r, t.emptyHash) {
  1351. rStr = fmt.Sprintf("empty%v", nEmpties)
  1352. eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
  1353. rStr)
  1354. nEmpties++
  1355. }
  1356. fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
  1357. lStr, rStr)
  1358. fmt.Fprint(w, eStr)
  1359. default:
  1360. }
  1361. return false
  1362. })
  1363. fmt.Fprintf(w, "}\n")
  1364. return err
  1365. }
  1366. // PrintGraphviz prints the output of Tree.Graphviz
  1367. func (t *Tree) PrintGraphviz(fromRoot []byte) error {
  1368. if fromRoot == nil {
  1369. var err error
  1370. fromRoot, err = t.Root()
  1371. if err != nil {
  1372. return err
  1373. }
  1374. }
  1375. return t.PrintGraphvizFirstNLevels(fromRoot, t.maxLevels)
  1376. }
  1377. // PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
  1378. func (t *Tree) PrintGraphvizFirstNLevels(fromRoot []byte, untilLvl int) error {
  1379. if fromRoot == nil {
  1380. var err error
  1381. fromRoot, err = t.Root()
  1382. if err != nil {
  1383. return err
  1384. }
  1385. }
  1386. w := bytes.NewBufferString("")
  1387. fmt.Fprintf(w,
  1388. "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+":\n")
  1389. err := t.GraphvizFirstNLevels(w, fromRoot, untilLvl)
  1390. if err != nil {
  1391. fmt.Println(w)
  1392. return err
  1393. }
  1394. fmt.Fprintf(w,
  1395. "End of Graphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+"\n--------\n")
  1396. fmt.Println(w)
  1397. return nil
  1398. }
  1399. // TODO circom proofs
  1400. // TODO data structure for proofs (including root, key, value, siblings,
  1401. // hashFunction) + method to verify that data structure