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.

573 lines
16 KiB

  1. package statedb
  2. import (
  3. "errors"
  4. "fmt"
  5. "math/big"
  6. "os"
  7. "strconv"
  8. "github.com/hermeznetwork/hermez-node/common"
  9. "github.com/hermeznetwork/hermez-node/log"
  10. "github.com/iden3/go-merkletree"
  11. "github.com/iden3/go-merkletree/db"
  12. "github.com/iden3/go-merkletree/db/pebble"
  13. )
  14. // TODO(Edu): Document here how StateDB is kept consistent
  15. var (
  16. // ErrStateDBWithoutMT is used when a method that requires a MerkleTree
  17. // is called in a StateDB that does not have a MerkleTree defined
  18. ErrStateDBWithoutMT = errors.New("Can not call method to use MerkleTree in a StateDB without MerkleTree")
  19. // ErrAccountAlreadyExists is used when CreateAccount is called and the
  20. // Account already exists
  21. ErrAccountAlreadyExists = errors.New("Can not CreateAccount because Account already exists")
  22. // ErrToIdxNotFound is used when trying to get the ToIdx from ToEthAddr
  23. // or ToEthAddr&ToBJJ
  24. ErrToIdxNotFound = errors.New("ToIdx can not be found")
  25. // KeyCurrentBatch is used as key in the db to store the current BatchNum
  26. KeyCurrentBatch = []byte("k:currentbatch")
  27. // PrefixKeyIdx is the key prefix for idx in the db
  28. PrefixKeyIdx = []byte("i:")
  29. // PrefixKeyAccHash is the key prefix for account hash in the db
  30. PrefixKeyAccHash = []byte("h:")
  31. // PrefixKeyMT is the key prefix for merkle tree in the db
  32. PrefixKeyMT = []byte("m:")
  33. // PrefixKeyAddr is the key prefix for address in the db
  34. PrefixKeyAddr = []byte("a:")
  35. // PrefixKeyAddrBJJ is the key prefix for address-babyjubjub in the db
  36. PrefixKeyAddrBJJ = []byte("ab:")
  37. )
  38. const (
  39. // PathStateDB defines the subpath of the StateDB
  40. PathStateDB = "/statedb"
  41. // PathBatchNum defines the subpath of the Batch Checkpoint in the
  42. // subpath of the StateDB
  43. PathBatchNum = "/BatchNum"
  44. // PathCurrent defines the subpath of the current Batch in the subpath
  45. // of the StateDB
  46. PathCurrent = "/current"
  47. // TypeSynchronizer defines a StateDB used by the Synchronizer, that
  48. // generates the ExitTree when processing the txs
  49. TypeSynchronizer = "synchronizer"
  50. // TypeTxSelector defines a StateDB used by the TxSelector, without
  51. // computing ExitTree neither the ZKInputs
  52. TypeTxSelector = "txselector"
  53. // TypeBatchBuilder defines a StateDB used by the BatchBuilder, that
  54. // generates the ExitTree and the ZKInput when processing the txs
  55. TypeBatchBuilder = "batchbuilder"
  56. )
  57. // TypeStateDB determines the type of StateDB
  58. type TypeStateDB string
  59. // StateDB represents the StateDB object
  60. type StateDB struct {
  61. path string
  62. currentBatch common.BatchNum
  63. db *pebble.PebbleStorage
  64. mt *merkletree.MerkleTree
  65. typ TypeStateDB
  66. // idx holds the current Idx that the BatchBuilder is using
  67. idx common.Idx
  68. zki *common.ZKInputs
  69. i int // i is the current transaction index in the ZKInputs generation (zki)
  70. }
  71. // NewStateDB creates a new StateDB, allowing to use an in-memory or in-disk
  72. // storage
  73. func NewStateDB(path string, typ TypeStateDB, nLevels int) (*StateDB, error) {
  74. var sto *pebble.PebbleStorage
  75. var err error
  76. sto, err = pebble.NewPebbleStorage(path+PathStateDB+PathCurrent, false)
  77. if err != nil {
  78. return nil, err
  79. }
  80. var mt *merkletree.MerkleTree = nil
  81. if typ == TypeSynchronizer || typ == TypeBatchBuilder {
  82. mt, err = merkletree.NewMerkleTree(sto.WithPrefix(PrefixKeyMT), nLevels)
  83. if err != nil {
  84. return nil, err
  85. }
  86. }
  87. if typ == TypeTxSelector && nLevels != 0 {
  88. return nil, fmt.Errorf("invalid StateDB parameters: StateDB type==TypeStateDB can not have nLevels!=0")
  89. }
  90. sdb := &StateDB{
  91. path: path + PathStateDB,
  92. db: sto,
  93. mt: mt,
  94. typ: typ,
  95. }
  96. // load currentBatch
  97. sdb.currentBatch, err = sdb.GetCurrentBatch()
  98. if err != nil {
  99. return nil, err
  100. }
  101. // make reset (get checkpoint) at currentBatch
  102. err = sdb.reset(sdb.currentBatch, false)
  103. if err != nil {
  104. return nil, err
  105. }
  106. return sdb, nil
  107. }
  108. // DB returns the *pebble.PebbleStorage from the StateDB
  109. func (s *StateDB) DB() *pebble.PebbleStorage {
  110. return s.db
  111. }
  112. // GetCurrentBatch returns the current BatchNum stored in the StateDB
  113. func (s *StateDB) GetCurrentBatch() (common.BatchNum, error) {
  114. cbBytes, err := s.db.Get(KeyCurrentBatch)
  115. if err == db.ErrNotFound {
  116. return 0, nil
  117. }
  118. if err != nil {
  119. return 0, err
  120. }
  121. return common.BatchNumFromBytes(cbBytes)
  122. }
  123. // setCurrentBatch stores the current BatchNum in the StateDB
  124. func (s *StateDB) setCurrentBatch() error {
  125. tx, err := s.db.NewTx()
  126. if err != nil {
  127. return err
  128. }
  129. err = tx.Put(KeyCurrentBatch, s.currentBatch.Bytes())
  130. if err != nil {
  131. return err
  132. }
  133. if err := tx.Commit(); err != nil {
  134. return err
  135. }
  136. return nil
  137. }
  138. // MakeCheckpoint does a checkpoint at the given batchNum in the defined path. Internally this advances & stores the current BatchNum, and then stores a Checkpoint of the current state of the StateDB.
  139. func (s *StateDB) MakeCheckpoint() error {
  140. // advance currentBatch
  141. s.currentBatch++
  142. checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(s.currentBatch))
  143. err := s.setCurrentBatch()
  144. if err != nil {
  145. return err
  146. }
  147. // if checkpoint BatchNum already exist in disk, delete it
  148. if _, err := os.Stat(checkpointPath); !os.IsNotExist(err) {
  149. err := os.RemoveAll(checkpointPath)
  150. if err != nil {
  151. return err
  152. }
  153. } else if err != nil && !os.IsNotExist(err) {
  154. return err
  155. }
  156. // execute Checkpoint
  157. err = s.db.Pebble().Checkpoint(checkpointPath)
  158. if err != nil {
  159. return err
  160. }
  161. return nil
  162. }
  163. // DeleteCheckpoint removes if exist the checkpoint of the given batchNum
  164. func (s *StateDB) DeleteCheckpoint(batchNum common.BatchNum) error {
  165. checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(batchNum))
  166. if _, err := os.Stat(checkpointPath); os.IsNotExist(err) {
  167. return fmt.Errorf("Checkpoint with batchNum %d does not exist in DB", batchNum)
  168. }
  169. return os.RemoveAll(checkpointPath)
  170. }
  171. func pebbleMakeCheckpoint(source, dest string) error {
  172. // Remove dest folder (if it exists) before doing the checkpoint
  173. if _, err := os.Stat(dest); !os.IsNotExist(err) {
  174. err := os.RemoveAll(dest)
  175. if err != nil {
  176. return err
  177. }
  178. } else if err != nil && !os.IsNotExist(err) {
  179. return err
  180. }
  181. sto, err := pebble.NewPebbleStorage(source, false)
  182. if err != nil {
  183. return err
  184. }
  185. defer func() {
  186. errClose := sto.Pebble().Close()
  187. if errClose != nil {
  188. log.Errorw("Pebble.Close", "err", errClose)
  189. }
  190. }()
  191. // execute Checkpoint
  192. err = sto.Pebble().Checkpoint(dest)
  193. if err != nil {
  194. return err
  195. }
  196. return nil
  197. }
  198. // Reset resets the StateDB to the checkpoint at the given batchNum. Reset
  199. // does not delete the checkpoints between old current and the new current,
  200. // those checkpoints will remain in the storage, and eventually will be
  201. // deleted when MakeCheckpoint overwrites them.
  202. func (s *StateDB) Reset(batchNum common.BatchNum) error {
  203. return s.reset(batchNum, true)
  204. }
  205. // reset resets the StateDB to the checkpoint at the given batchNum. Reset
  206. // does not delete the checkpoints between old current and the new current,
  207. // those checkpoints will remain in the storage, and eventually will be
  208. // deleted when MakeCheckpoint overwrites them. `closeCurrent` will close the
  209. // currently opened db before doing the reset.
  210. func (s *StateDB) reset(batchNum common.BatchNum, closeCurrent bool) error {
  211. checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(batchNum))
  212. currentPath := s.path + PathCurrent
  213. if closeCurrent {
  214. if err := s.db.Pebble().Close(); err != nil {
  215. return err
  216. }
  217. }
  218. // remove 'current'
  219. err := os.RemoveAll(currentPath)
  220. if err != nil {
  221. return err
  222. }
  223. if batchNum == 0 {
  224. // if batchNum == 0, open the new fresh 'current'
  225. sto, err := pebble.NewPebbleStorage(currentPath, false)
  226. if err != nil {
  227. return err
  228. }
  229. s.db = sto
  230. s.idx = 255
  231. s.currentBatch = batchNum
  232. return nil
  233. }
  234. // copy 'BatchNumX' to 'current'
  235. err = pebbleMakeCheckpoint(checkpointPath, currentPath)
  236. if err != nil {
  237. return err
  238. }
  239. // open the new 'current'
  240. sto, err := pebble.NewPebbleStorage(currentPath, false)
  241. if err != nil {
  242. return err
  243. }
  244. s.db = sto
  245. // get currentBatch num
  246. s.currentBatch, err = s.GetCurrentBatch()
  247. if err != nil {
  248. return err
  249. }
  250. // idx is obtained from the statedb reset
  251. s.idx, err = s.getIdx()
  252. if err != nil {
  253. return err
  254. }
  255. if s.mt != nil {
  256. // open the MT for the current s.db
  257. mt, err := merkletree.NewMerkleTree(s.db.WithPrefix(PrefixKeyMT), s.mt.MaxLevels())
  258. if err != nil {
  259. return err
  260. }
  261. s.mt = mt
  262. }
  263. return nil
  264. }
  265. // GetAccount returns the account for the given Idx
  266. func (s *StateDB) GetAccount(idx common.Idx) (*common.Account, error) {
  267. return getAccountInTreeDB(s.db, idx)
  268. }
  269. // GetAccounts returns all the accounts in the db. Use for debugging pruposes
  270. // only.
  271. func (s *StateDB) GetAccounts() ([]common.Account, error) {
  272. idxDB := s.db.WithPrefix(PrefixKeyIdx)
  273. idxs := []common.Idx{}
  274. // NOTE: Current implementation of Iterate in the pebble interface is
  275. // not efficient, as it iterates over all keys. Improve it following
  276. // this example: https://github.com/cockroachdb/pebble/pull/923/files
  277. if err := idxDB.Iterate(func(k []byte, v []byte) (bool, error) {
  278. idx, err := common.IdxFromBytes(k)
  279. if err != nil {
  280. return false, err
  281. }
  282. idxs = append(idxs, idx)
  283. return true, nil
  284. }); err != nil {
  285. return nil, err
  286. }
  287. accs := []common.Account{}
  288. for i := range idxs {
  289. acc, err := s.GetAccount(idxs[i])
  290. if err != nil {
  291. return nil, err
  292. }
  293. accs = append(accs, *acc)
  294. }
  295. return accs, nil
  296. }
  297. // getAccountInTreeDB is abstracted from StateDB to be used from StateDB and
  298. // from ExitTree. GetAccount returns the account for the given Idx
  299. func getAccountInTreeDB(sto db.Storage, idx common.Idx) (*common.Account, error) {
  300. idxBytes, err := idx.Bytes()
  301. if err != nil {
  302. return nil, err
  303. }
  304. vBytes, err := sto.Get(append(PrefixKeyIdx, idxBytes[:]...))
  305. if err != nil {
  306. return nil, err
  307. }
  308. accBytes, err := sto.Get(append(PrefixKeyAccHash, vBytes...))
  309. if err != nil {
  310. return nil, err
  311. }
  312. var b [32 * common.NLeafElems]byte
  313. copy(b[:], accBytes)
  314. account, err := common.AccountFromBytes(b)
  315. if err != nil {
  316. return nil, err
  317. }
  318. account.Idx = idx
  319. return account, nil
  320. }
  321. // CreateAccount creates a new Account in the StateDB for the given Idx. If
  322. // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
  323. // MerkleTree, returning a CircomProcessorProof.
  324. func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
  325. cpp, err := createAccountInTreeDB(s.db, s.mt, idx, account)
  326. if err != nil {
  327. return cpp, err
  328. }
  329. // store idx by EthAddr & BJJ
  330. err = s.setIdxByEthAddrBJJ(idx, account.EthAddr, account.PublicKey)
  331. return cpp, err
  332. }
  333. // createAccountInTreeDB is abstracted from StateDB to be used from StateDB and
  334. // from ExitTree. Creates a new Account in the StateDB for the given Idx. If
  335. // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
  336. // MerkleTree, returning a CircomProcessorProof.
  337. func createAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
  338. // store at the DB the key: v, and value: leaf.Bytes()
  339. v, err := account.HashValue()
  340. if err != nil {
  341. return nil, err
  342. }
  343. accountBytes, err := account.Bytes()
  344. if err != nil {
  345. return nil, err
  346. }
  347. // store the Leaf value
  348. tx, err := sto.NewTx()
  349. if err != nil {
  350. return nil, err
  351. }
  352. idxBytes, err := idx.Bytes()
  353. if err != nil {
  354. return nil, err
  355. }
  356. _, err = tx.Get(append(PrefixKeyIdx, idxBytes[:]...))
  357. if err != db.ErrNotFound {
  358. return nil, ErrAccountAlreadyExists
  359. }
  360. err = tx.Put(append(PrefixKeyAccHash, v.Bytes()...), accountBytes[:])
  361. if err != nil {
  362. return nil, err
  363. }
  364. err = tx.Put(append(PrefixKeyIdx, idxBytes[:]...), v.Bytes())
  365. if err != nil {
  366. return nil, err
  367. }
  368. if err := tx.Commit(); err != nil {
  369. return nil, err
  370. }
  371. if mt != nil {
  372. return mt.AddAndGetCircomProof(idx.BigInt(), v)
  373. }
  374. return nil, nil
  375. }
  376. // UpdateAccount updates the Account in the StateDB for the given Idx. If
  377. // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
  378. // MerkleTree, returning a CircomProcessorProof.
  379. func (s *StateDB) UpdateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
  380. return updateAccountInTreeDB(s.db, s.mt, idx, account)
  381. }
  382. // updateAccountInTreeDB is abstracted from StateDB to be used from StateDB and
  383. // from ExitTree. Updates the Account in the StateDB for the given Idx. If
  384. // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
  385. // MerkleTree, returning a CircomProcessorProof.
  386. func updateAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
  387. // store at the DB the key: v, and value: account.Bytes()
  388. v, err := account.HashValue()
  389. if err != nil {
  390. return nil, err
  391. }
  392. accountBytes, err := account.Bytes()
  393. if err != nil {
  394. return nil, err
  395. }
  396. tx, err := sto.NewTx()
  397. if err != nil {
  398. return nil, err
  399. }
  400. err = tx.Put(append(PrefixKeyAccHash, v.Bytes()...), accountBytes[:])
  401. if err != nil {
  402. return nil, err
  403. }
  404. idxBytes, err := idx.Bytes()
  405. if err != nil {
  406. return nil, err
  407. }
  408. err = tx.Put(append(PrefixKeyIdx, idxBytes[:]...), v.Bytes())
  409. if err != nil {
  410. return nil, err
  411. }
  412. if err := tx.Commit(); err != nil {
  413. return nil, err
  414. }
  415. if mt != nil {
  416. return mt.Update(idx.BigInt(), v)
  417. }
  418. return nil, nil
  419. }
  420. // MTGetProof returns the CircomVerifierProof for a given Idx
  421. func (s *StateDB) MTGetProof(idx common.Idx) (*merkletree.CircomVerifierProof, error) {
  422. if s.mt == nil {
  423. return nil, ErrStateDBWithoutMT
  424. }
  425. return s.mt.GenerateCircomVerifierProof(idx.BigInt(), s.mt.Root())
  426. }
  427. // MTGetRoot returns the current root of the underlying Merkle Tree
  428. func (s *StateDB) MTGetRoot() *big.Int {
  429. return s.mt.Root().BigInt()
  430. }
  431. // LocalStateDB represents the local StateDB which allows to make copies from
  432. // the synchronizer StateDB, and is used by the tx-selector and the
  433. // batch-builder. LocalStateDB is an in-memory storage.
  434. type LocalStateDB struct {
  435. *StateDB
  436. synchronizerStateDB *StateDB
  437. }
  438. // NewLocalStateDB returns a new LocalStateDB connected to the given
  439. // synchronizerDB
  440. func NewLocalStateDB(path string, synchronizerDB *StateDB, typ TypeStateDB, nLevels int) (*LocalStateDB, error) {
  441. s, err := NewStateDB(path, typ, nLevels)
  442. if err != nil {
  443. return nil, err
  444. }
  445. return &LocalStateDB{
  446. s,
  447. synchronizerDB,
  448. }, nil
  449. }
  450. // Reset performs a reset in the LocaStateDB. If fromSynchronizer is true, it
  451. // gets the state from LocalStateDB.synchronizerStateDB for the given batchNum. If fromSynchronizer is false, get the state from LocalStateDB checkpoints.
  452. func (l *LocalStateDB) Reset(batchNum common.BatchNum, fromSynchronizer bool) error {
  453. if batchNum == 0 {
  454. l.idx = 0
  455. return nil
  456. }
  457. synchronizerCheckpointPath := l.synchronizerStateDB.path + PathBatchNum + strconv.Itoa(int(batchNum))
  458. checkpointPath := l.path + PathBatchNum + strconv.Itoa(int(batchNum))
  459. currentPath := l.path + PathCurrent
  460. if fromSynchronizer {
  461. // use checkpoint from SynchronizerStateDB
  462. if _, err := os.Stat(synchronizerCheckpointPath); os.IsNotExist(err) {
  463. // if synchronizerStateDB does not have checkpoint at batchNum, return err
  464. return fmt.Errorf("Checkpoint not exist in Synchronizer")
  465. }
  466. if err := l.db.Pebble().Close(); err != nil {
  467. return err
  468. }
  469. // remove 'current'
  470. err := os.RemoveAll(currentPath)
  471. if err != nil {
  472. return err
  473. }
  474. // copy synchronizer'BatchNumX' to 'current'
  475. err = pebbleMakeCheckpoint(synchronizerCheckpointPath, currentPath)
  476. if err != nil {
  477. return err
  478. }
  479. // copy synchronizer'BatchNumX' to 'BatchNumX'
  480. err = pebbleMakeCheckpoint(synchronizerCheckpointPath, checkpointPath)
  481. if err != nil {
  482. return err
  483. }
  484. // open the new 'current'
  485. sto, err := pebble.NewPebbleStorage(currentPath, false)
  486. if err != nil {
  487. return err
  488. }
  489. l.db = sto
  490. // get currentBatch num
  491. l.currentBatch, err = l.GetCurrentBatch()
  492. if err != nil {
  493. return err
  494. }
  495. // open the MT for the current s.db
  496. mt, err := merkletree.NewMerkleTree(l.db.WithPrefix(PrefixKeyMT), l.mt.MaxLevels())
  497. if err != nil {
  498. return err
  499. }
  500. l.mt = mt
  501. return nil
  502. }
  503. // use checkpoint from LocalStateDB
  504. return l.StateDB.reset(batchNum, true)
  505. }