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.

456 lines
12 KiB

  1. package statedb
  2. import (
  3. "errors"
  4. "fmt"
  5. "os"
  6. "os/exec"
  7. "strconv"
  8. "github.com/hermeznetwork/hermez-node/common"
  9. "github.com/iden3/go-merkletree"
  10. "github.com/iden3/go-merkletree/db"
  11. "github.com/iden3/go-merkletree/db/pebble"
  12. )
  13. // TODO(Edu): Document here how StateDB is kept consistent
  14. // ErrStateDBWithoutMT is used when a method that requires a MerkleTree is
  15. // called in a StateDB that does not have a MerkleTree defined
  16. var ErrStateDBWithoutMT = errors.New("Can not call method to use MerkleTree in a StateDB without MerkleTree")
  17. // ErrAccountAlreadyExists is used when CreateAccount is called and the Account
  18. // already exists
  19. var ErrAccountAlreadyExists = errors.New("Can not CreateAccount because Account already exists")
  20. // ErrToIdxNotFound is used when trying to get the ToIdx from ToEthAddr or
  21. // ToEthAddr&ToBJJ
  22. var ErrToIdxNotFound = errors.New("ToIdx can not be found")
  23. // KeyCurrentBatch is used as key in the db to store the current BatchNum
  24. var KeyCurrentBatch = []byte("currentbatch")
  25. // PathStateDB defines the subpath of the StateDB
  26. const PathStateDB = "/statedb"
  27. // PathBatchNum defines the subpath of the Batch Checkpoint in the subpath of
  28. // the StateDB
  29. const PathBatchNum = "/BatchNum"
  30. // PathCurrent defines the subpath of the current Batch in the subpath of the
  31. // StateDB
  32. const PathCurrent = "/current"
  33. // StateDB represents the StateDB object
  34. type StateDB struct {
  35. path string
  36. currentBatch common.BatchNum
  37. db *pebble.PebbleStorage
  38. mt *merkletree.MerkleTree
  39. // idx holds the current Idx that the BatchBuilder is using
  40. idx common.Idx
  41. zki *common.ZKInputs
  42. i int // i is the current transaction index in the ZKInputs generation (zki)
  43. }
  44. // NewStateDB creates a new StateDB, allowing to use an in-memory or in-disk
  45. // storage
  46. func NewStateDB(path string, withMT bool, nLevels int) (*StateDB, error) {
  47. var sto *pebble.PebbleStorage
  48. var err error
  49. sto, err = pebble.NewPebbleStorage(path+PathStateDB+PathCurrent, false)
  50. if err != nil {
  51. return nil, err
  52. }
  53. var mt *merkletree.MerkleTree = nil
  54. if withMT {
  55. mt, err = merkletree.NewMerkleTree(sto, nLevels)
  56. if err != nil {
  57. return nil, err
  58. }
  59. }
  60. sdb := &StateDB{
  61. path: path + PathStateDB,
  62. db: sto,
  63. mt: mt,
  64. }
  65. // load currentBatch
  66. sdb.currentBatch, err = sdb.GetCurrentBatch()
  67. if err != nil {
  68. return nil, err
  69. }
  70. // make reset (get checkpoint) at currentBatch
  71. err = sdb.Reset(sdb.currentBatch)
  72. if err != nil {
  73. return nil, err
  74. }
  75. return sdb, nil
  76. }
  77. // DB returns the *pebble.PebbleStorage from the StateDB
  78. func (s *StateDB) DB() *pebble.PebbleStorage {
  79. return s.db
  80. }
  81. // GetCurrentBatch returns the current BatchNum stored in the StateDB
  82. func (s *StateDB) GetCurrentBatch() (common.BatchNum, error) {
  83. cbBytes, err := s.db.Get(KeyCurrentBatch)
  84. if err == db.ErrNotFound {
  85. return 0, nil
  86. }
  87. if err != nil {
  88. return 0, err
  89. }
  90. return common.BatchNumFromBytes(cbBytes)
  91. }
  92. // setCurrentBatch stores the current BatchNum in the StateDB
  93. func (s *StateDB) setCurrentBatch() error {
  94. tx, err := s.db.NewTx()
  95. if err != nil {
  96. return err
  97. }
  98. err = tx.Put(KeyCurrentBatch, s.currentBatch.Bytes())
  99. if err != nil {
  100. return err
  101. }
  102. if err := tx.Commit(); err != nil {
  103. return err
  104. }
  105. return nil
  106. }
  107. // 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.
  108. func (s *StateDB) MakeCheckpoint() error {
  109. // advance currentBatch
  110. s.currentBatch++
  111. checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(s.currentBatch))
  112. err := s.setCurrentBatch()
  113. if err != nil {
  114. return err
  115. }
  116. // if checkpoint BatchNum already exist in disk, delete it
  117. if _, err := os.Stat(checkpointPath); !os.IsNotExist(err) {
  118. err := os.RemoveAll(checkpointPath)
  119. if err != nil {
  120. return err
  121. }
  122. }
  123. // execute Checkpoint
  124. err = s.db.Pebble().Checkpoint(checkpointPath)
  125. if err != nil {
  126. return err
  127. }
  128. return nil
  129. }
  130. // DeleteCheckpoint removes if exist the checkpoint of the given batchNum
  131. func (s *StateDB) DeleteCheckpoint(batchNum common.BatchNum) error {
  132. checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(batchNum))
  133. if _, err := os.Stat(checkpointPath); os.IsNotExist(err) {
  134. return fmt.Errorf("Checkpoint with batchNum %d does not exist in DB", batchNum)
  135. }
  136. return os.RemoveAll(checkpointPath)
  137. }
  138. // Reset resets the StateDB to the checkpoint at the given batchNum. Reset
  139. // does not delete the checkpoints between old current and the new current,
  140. // those checkpoints will remain in the storage, and eventually will be
  141. // deleted when MakeCheckpoint overwrites them.
  142. func (s *StateDB) Reset(batchNum common.BatchNum) error {
  143. checkpointPath := s.path + PathBatchNum + strconv.Itoa(int(batchNum))
  144. currentPath := s.path + PathCurrent
  145. // remove 'current'
  146. err := os.RemoveAll(currentPath)
  147. if err != nil {
  148. return err
  149. }
  150. if batchNum == 0 {
  151. // if batchNum == 0, open the new fresh 'current'
  152. sto, err := pebble.NewPebbleStorage(currentPath, false)
  153. if err != nil {
  154. return err
  155. }
  156. s.db = sto
  157. s.idx = 255
  158. s.currentBatch = batchNum
  159. return nil
  160. }
  161. // copy 'BatchNumX' to 'current'
  162. cmd := exec.Command("cp", "-r", checkpointPath, currentPath) //nolint:gosec
  163. err = cmd.Run()
  164. if err != nil {
  165. return err
  166. }
  167. // open the new 'current'
  168. sto, err := pebble.NewPebbleStorage(currentPath, false)
  169. if err != nil {
  170. return err
  171. }
  172. s.db = sto
  173. // get currentBatch num
  174. s.currentBatch, err = s.GetCurrentBatch()
  175. if err != nil {
  176. return err
  177. }
  178. // idx is obtained from the statedb reset
  179. s.idx, err = s.getIdx()
  180. if err != nil {
  181. return err
  182. }
  183. if s.mt != nil {
  184. // open the MT for the current s.db
  185. mt, err := merkletree.NewMerkleTree(s.db, s.mt.MaxLevels())
  186. if err != nil {
  187. return err
  188. }
  189. s.mt = mt
  190. }
  191. return nil
  192. }
  193. // GetAccount returns the account for the given Idx
  194. func (s *StateDB) GetAccount(idx common.Idx) (*common.Account, error) {
  195. return getAccountInTreeDB(s.db, idx)
  196. }
  197. // getAccountInTreeDB is abstracted from StateDB to be used from StateDB and
  198. // from ExitTree. GetAccount returns the account for the given Idx
  199. func getAccountInTreeDB(sto db.Storage, idx common.Idx) (*common.Account, error) {
  200. idxBytes, err := idx.Bytes()
  201. if err != nil {
  202. return nil, err
  203. }
  204. vBytes, err := sto.Get(idxBytes[:])
  205. if err != nil {
  206. return nil, err
  207. }
  208. accBytes, err := sto.Get(vBytes)
  209. if err != nil {
  210. return nil, err
  211. }
  212. var b [32 * common.NLeafElems]byte
  213. copy(b[:], accBytes)
  214. return common.AccountFromBytes(b)
  215. }
  216. // CreateAccount creates a new Account in the StateDB for the given Idx. If
  217. // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
  218. // MerkleTree, returning a CircomProcessorProof.
  219. func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
  220. cpp, err := createAccountInTreeDB(s.db, s.mt, idx, account)
  221. if err != nil {
  222. return cpp, err
  223. }
  224. // store idx by EthAddr & BJJ
  225. err = s.setIdxByEthAddrBJJ(idx, account.EthAddr, account.PublicKey)
  226. return cpp, err
  227. }
  228. // createAccountInTreeDB is abstracted from StateDB to be used from StateDB and
  229. // from ExitTree. Creates a new Account in the StateDB for the given Idx. If
  230. // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
  231. // MerkleTree, returning a CircomProcessorProof.
  232. func createAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
  233. // store at the DB the key: v, and value: leaf.Bytes()
  234. v, err := account.HashValue()
  235. if err != nil {
  236. return nil, err
  237. }
  238. accountBytes, err := account.Bytes()
  239. if err != nil {
  240. return nil, err
  241. }
  242. // store the Leaf value
  243. tx, err := sto.NewTx()
  244. if err != nil {
  245. return nil, err
  246. }
  247. idxBytes, err := idx.Bytes()
  248. if err != nil {
  249. return nil, err
  250. }
  251. _, err = tx.Get(idxBytes[:])
  252. if err != db.ErrNotFound {
  253. return nil, ErrAccountAlreadyExists
  254. }
  255. err = tx.Put(v.Bytes(), accountBytes[:])
  256. if err != nil {
  257. return nil, err
  258. }
  259. err = tx.Put(idxBytes[:], v.Bytes())
  260. if err != nil {
  261. return nil, err
  262. }
  263. if err := tx.Commit(); err != nil {
  264. return nil, err
  265. }
  266. if mt != nil {
  267. return mt.AddAndGetCircomProof(idx.BigInt(), v)
  268. }
  269. return nil, nil
  270. }
  271. // UpdateAccount updates the Account in the StateDB for the given Idx. If
  272. // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
  273. // MerkleTree, returning a CircomProcessorProof.
  274. func (s *StateDB) UpdateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
  275. return updateAccountInTreeDB(s.db, s.mt, idx, account)
  276. }
  277. // updateAccountInTreeDB is abstracted from StateDB to be used from StateDB and
  278. // from ExitTree. Updates the Account in the StateDB for the given Idx. If
  279. // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
  280. // MerkleTree, returning a CircomProcessorProof.
  281. func updateAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
  282. // store at the DB the key: v, and value: account.Bytes()
  283. v, err := account.HashValue()
  284. if err != nil {
  285. return nil, err
  286. }
  287. accountBytes, err := account.Bytes()
  288. if err != nil {
  289. return nil, err
  290. }
  291. tx, err := sto.NewTx()
  292. if err != nil {
  293. return nil, err
  294. }
  295. err = tx.Put(v.Bytes(), accountBytes[:])
  296. if err != nil {
  297. return nil, err
  298. }
  299. idxBytes, err := idx.Bytes()
  300. if err != nil {
  301. return nil, err
  302. }
  303. err = tx.Put(idxBytes[:], v.Bytes())
  304. if err != nil {
  305. return nil, err
  306. }
  307. if err := tx.Commit(); err != nil {
  308. return nil, err
  309. }
  310. if mt != nil {
  311. return mt.Update(idx.BigInt(), v)
  312. }
  313. return nil, nil
  314. }
  315. // MTGetProof returns the CircomVerifierProof for a given Idx
  316. func (s *StateDB) MTGetProof(idx common.Idx) (*merkletree.CircomVerifierProof, error) {
  317. if s.mt == nil {
  318. return nil, ErrStateDBWithoutMT
  319. }
  320. return s.mt.GenerateCircomVerifierProof(idx.BigInt(), s.mt.Root())
  321. }
  322. // LocalStateDB represents the local StateDB which allows to make copies from
  323. // the synchronizer StateDB, and is used by the tx-selector and the
  324. // batch-builder. LocalStateDB is an in-memory storage.
  325. type LocalStateDB struct {
  326. *StateDB
  327. synchronizerStateDB *StateDB
  328. }
  329. // NewLocalStateDB returns a new LocalStateDB connected to the given
  330. // synchronizerDB
  331. func NewLocalStateDB(path string, synchronizerDB *StateDB, withMT bool, nLevels int) (*LocalStateDB, error) {
  332. s, err := NewStateDB(path, withMT, nLevels)
  333. if err != nil {
  334. return nil, err
  335. }
  336. return &LocalStateDB{
  337. s,
  338. synchronizerDB,
  339. }, nil
  340. }
  341. // Reset performs a reset in the LocaStateDB. If fromSynchronizer is true, it
  342. // gets the state from LocalStateDB.synchronizerStateDB for the given batchNum. If fromSynchronizer is false, get the state from LocalStateDB checkpoints.
  343. func (l *LocalStateDB) Reset(batchNum common.BatchNum, fromSynchronizer bool) error {
  344. if batchNum == 0 {
  345. l.idx = 0
  346. return nil
  347. }
  348. synchronizerCheckpointPath := l.synchronizerStateDB.path + PathBatchNum + strconv.Itoa(int(batchNum))
  349. checkpointPath := l.path + PathBatchNum + strconv.Itoa(int(batchNum))
  350. currentPath := l.path + PathCurrent
  351. if fromSynchronizer {
  352. // use checkpoint from SynchronizerStateDB
  353. if _, err := os.Stat(synchronizerCheckpointPath); os.IsNotExist(err) {
  354. // if synchronizerStateDB does not have checkpoint at batchNum, return err
  355. return fmt.Errorf("Checkpoint not exist in Synchronizer")
  356. }
  357. // remove 'current'
  358. err := os.RemoveAll(currentPath)
  359. if err != nil {
  360. return err
  361. }
  362. // copy synchronizer'BatchNumX' to 'current'
  363. cmd := exec.Command("cp", "-r", synchronizerCheckpointPath, currentPath) //nolint:gosec
  364. err = cmd.Run()
  365. if err != nil {
  366. return err
  367. }
  368. // copy synchronizer-'BatchNumX' to 'BatchNumX'
  369. cmd = exec.Command("cp", "-r", synchronizerCheckpointPath, checkpointPath) //nolint:gosec
  370. err = cmd.Run()
  371. if err != nil {
  372. return err
  373. }
  374. // open the new 'current'
  375. sto, err := pebble.NewPebbleStorage(currentPath, false)
  376. if err != nil {
  377. return err
  378. }
  379. l.db = sto
  380. // get currentBatch num
  381. l.currentBatch, err = l.GetCurrentBatch()
  382. if err != nil {
  383. return err
  384. }
  385. // open the MT for the current s.db
  386. mt, err := merkletree.NewMerkleTree(l.db, l.mt.MaxLevels())
  387. if err != nil {
  388. return err
  389. }
  390. l.mt = mt
  391. return nil
  392. }
  393. // use checkpoint from LocalStateDB
  394. return l.StateDB.Reset(batchNum)
  395. }