From 086a7d04832398b228c4df9f6f29757ca37ec64c Mon Sep 17 00:00:00 2001 From: Olusegun Akintayo Date: Tue, 22 Nov 2022 20:56:26 +0400 Subject: [PATCH] Keep memstore contents after service worker restarts (#15913) * Add all controllers in memstore to store Add methods to controller to reset memstore Reset memstore when popup or tab is closed. Signed-off-by: Akintayo A. Olusegun * When profile is loaded, set isFirstTime to true.. After resetting the controllers, set the flag to false. Signed-off-by: Akintayo A. Olusegun * Remove console.logs Signed-off-by: Akintayo A. Olusegun * For some reason programmatically computing the store is not working. Signed-off-by: Akintayo A. Olusegun * Proper check for browser.storage Signed-off-by: Akintayo A. Olusegun * do a list of rest methods instead of reset controllers. Signed-off-by: Akintayo A. Olusegun * Mock controller resetStates and localstore get/set Signed-off-by: Akintayo A. Olusegun * Comments about TLC Signed-off-by: Akintayo A. Olusegun * bind this. Signed-off-by: Akintayo A. Olusegun * use globalThis instead of locastore to store first time state. Signed-off-by: Akintayo A. Olusegun * Test to check that resetStates is not called a second time Signed-off-by: Akintayo A. Olusegun * Set init state in GasFeeController and other controllers so that their state is persisted accross SW restarts Signed-off-by: Akintayo A. Olusegun * Revert localstore changes Signed-off-by: Akintayo A. Olusegun * wrap the reset states changes in MV3 flag Signed-off-by: Akintayo A. Olusegun * Remove localstore from metamask-controller Signed-off-by: Akintayo A. Olusegun * Always reset state on MMController start in MV2. Signed-off-by: Akintayo A. Olusegun * Use relative path for import of isManifestV3 Signed-off-by: Akintayo A. Olusegun * Fix unit test Signed-off-by: Akintayo A. Olusegun Signed-off-by: Akintayo A. Olusegun --- app/scripts/app-init.js | 4 + app/scripts/controllers/ens/index.js | 5 + app/scripts/controllers/swaps.js | 4 + app/scripts/controllers/transactions/index.js | 5 + app/scripts/lib/account-tracker.js | 4 + app/scripts/lib/decrypt-message-manager.js | 8 ++ .../lib/encryption-public-key-manager.js | 8 ++ app/scripts/lib/message-manager.js | 8 ++ app/scripts/lib/personal-message-manager.js | 8 ++ app/scripts/lib/typed-message-manager.js | 8 ++ .../metamask-controller.actions.test.js | 6 + app/scripts/metamask-controller.js | 114 +++++++++++++----- app/scripts/metamask-controller.test.js | 15 +++ 13 files changed, 166 insertions(+), 31 deletions(-) diff --git a/app/scripts/app-init.js b/app/scripts/app-init.js index 450c5959737e..0f61d9888ef1 100644 --- a/app/scripts/app-init.js +++ b/app/scripts/app-init.js @@ -127,6 +127,10 @@ chrome.runtime.onMessage.addListener(() => { return false; }); +chrome.runtime.onStartup.addListener(() => { + globalThis.isFirstTimeProfileLoaded = true; +}); + /* * This content script is injected programmatically because * MAIN world injection does not work properly via manifest diff --git a/app/scripts/controllers/ens/index.js b/app/scripts/controllers/ens/index.js index 85477ef7e5cf..47363a1562e6 100644 --- a/app/scripts/controllers/ens/index.js +++ b/app/scripts/controllers/ens/index.js @@ -27,6 +27,11 @@ export default class EnsController { } this.store = new ObservableStore(initState); + + this.resetState = () => { + this.store.updateState(initState); + }; + onNetworkDidChange(() => { this.store.putState(initState); const chainId = getCurrentChainId(); diff --git a/app/scripts/controllers/swaps.js b/app/scripts/controllers/swaps.js index 8c7e2c2148d6..476d020ee1ac 100644 --- a/app/scripts/controllers/swaps.js +++ b/app/scripts/controllers/swaps.js @@ -115,6 +115,10 @@ export default class SwapsController { swapsState: { ...initialState.swapsState }, }); + this.resetState = () => { + this.store.updateState({ swapsState: { ...initialState.swapsState } }); + }; + this._fetchTradesInfo = fetchTradesInfo; this._getCurrentChainId = getCurrentChainId; this._getEIP1559GasFeeEstimates = getEIP1559GasFeeEstimates; diff --git a/app/scripts/controllers/transactions/index.js b/app/scripts/controllers/transactions/index.js index 5becce48b40c..2b42ed3ff10f 100644 --- a/app/scripts/controllers/transactions/index.js +++ b/app/scripts/controllers/transactions/index.js @@ -151,6 +151,11 @@ export default class TransactionController extends EventEmitter { this.getTokenStandardAndDetails = opts.getTokenStandardAndDetails; this.memStore = new ObservableStore({}); + + this.resetState = () => { + this._updateMemstore(); + }; + this.query = new EthQuery(this.provider); this.txGasUtil = new TxGasUtil(this.provider); diff --git a/app/scripts/lib/account-tracker.js b/app/scripts/lib/account-tracker.js index 7e57a8ee9732..9f0600d65914 100644 --- a/app/scripts/lib/account-tracker.js +++ b/app/scripts/lib/account-tracker.js @@ -62,6 +62,10 @@ export default class AccountTracker { }; this.store = new ObservableStore(initState); + this.resetState = () => { + this.store.updateState(initState); + }; + this._provider = opts.provider; this._query = pify(new EthQuery(this._provider)); this._blockTracker = opts.blockTracker; diff --git a/app/scripts/lib/decrypt-message-manager.js b/app/scripts/lib/decrypt-message-manager.js index f5d221485cbb..4805ec398ca1 100644 --- a/app/scripts/lib/decrypt-message-manager.js +++ b/app/scripts/lib/decrypt-message-manager.js @@ -41,6 +41,14 @@ export default class DecryptMessageManager extends EventEmitter { unapprovedDecryptMsgs: {}, unapprovedDecryptMsgCount: 0, }); + + this.resetState = () => { + this.memStore.updateState({ + unapprovedDecryptMsgs: {}, + unapprovedDecryptMsgCount: 0, + }); + }; + this.messages = []; this.metricsEvent = opts.metricsEvent; } diff --git a/app/scripts/lib/encryption-public-key-manager.js b/app/scripts/lib/encryption-public-key-manager.js index 5f2b74993655..cf96c320a203 100644 --- a/app/scripts/lib/encryption-public-key-manager.js +++ b/app/scripts/lib/encryption-public-key-manager.js @@ -36,6 +36,14 @@ export default class EncryptionPublicKeyManager extends EventEmitter { unapprovedEncryptionPublicKeyMsgs: {}, unapprovedEncryptionPublicKeyMsgCount: 0, }); + + this.resetState = () => { + this.memStore.updateState({ + unapprovedEncryptionPublicKeyMsgs: {}, + unapprovedEncryptionPublicKeyMsgCount: 0, + }); + }; + this.messages = []; this.metricsEvent = opts.metricsEvent; } diff --git a/app/scripts/lib/message-manager.js b/app/scripts/lib/message-manager.js index 6f80bb75b7b5..dfc4e3ec3aaf 100644 --- a/app/scripts/lib/message-manager.js +++ b/app/scripts/lib/message-manager.js @@ -36,6 +36,14 @@ export default class MessageManager extends EventEmitter { unapprovedMsgs: {}, unapprovedMsgCount: 0, }); + + this.resetState = () => { + this.memStore.updateState({ + unapprovedMsgs: {}, + unapprovedMsgCount: 0, + }); + }; + this.messages = []; this.metricsEvent = metricsEvent; } diff --git a/app/scripts/lib/personal-message-manager.js b/app/scripts/lib/personal-message-manager.js index 493288907a90..338458a66db0 100644 --- a/app/scripts/lib/personal-message-manager.js +++ b/app/scripts/lib/personal-message-manager.js @@ -43,6 +43,14 @@ export default class PersonalMessageManager extends EventEmitter { unapprovedPersonalMsgs: {}, unapprovedPersonalMsgCount: 0, }); + + this.resetState = () => { + this.memStore.updateState({ + unapprovedPersonalMsgs: {}, + unapprovedPersonalMsgCount: 0, + }); + }; + this.messages = []; this.metricsEvent = metricsEvent; } diff --git a/app/scripts/lib/typed-message-manager.js b/app/scripts/lib/typed-message-manager.js index 60356ba8a522..397e27b0c3c7 100644 --- a/app/scripts/lib/typed-message-manager.js +++ b/app/scripts/lib/typed-message-manager.js @@ -43,6 +43,14 @@ export default class TypedMessageManager extends EventEmitter { unapprovedTypedMessages: {}, unapprovedTypedMessagesCount: 0, }); + + this.resetState = () => { + this.memStore.updateState({ + unapprovedTypedMessages: {}, + unapprovedTypedMessagesCount: 0, + }); + }; + this.messages = []; this.metricsEvent = metricsEvent; } diff --git a/app/scripts/metamask-controller.actions.test.js b/app/scripts/metamask-controller.actions.test.js index b3eff23d7c6d..c8b335d9435e 100644 --- a/app/scripts/metamask-controller.actions.test.js +++ b/app/scripts/metamask-controller.actions.test.js @@ -23,6 +23,12 @@ const browserPolyfillMock = { }, getPlatformInfo: async () => 'mac', }, + storage: { + local: { + get: sinon.stub().resolves({}), + set: sinon.stub().resolves(), + }, + }, }; let loggerMiddlewareMock; diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js index 6cb3c82e7bf2..47c958b84df5 100644 --- a/app/scripts/metamask-controller.js +++ b/app/scripts/metamask-controller.js @@ -100,6 +100,7 @@ import { getTokenValueParam, hexToDecimal, } from '../../shared/lib/metamask-controller-utils'; +import { isManifestV3 } from '../../shared/modules/mv3.utils'; import { onMessageReceived, checkForMultipleVersionsRunning, @@ -417,6 +418,7 @@ export default class MetamaskController extends EventEmitter { : GAS_API_BASE_URL; this.gasFeeController = new GasFeeController({ + state: initState.GasFeeController, interval: 10000, messenger: gasFeeMessenger, clientId: SWAPS_CLIENT_ID, @@ -481,26 +483,30 @@ export default class MetamaskController extends EventEmitter { ); // token exchange rate tracker - this.tokenRatesController = new TokenRatesController({ - onTokensStateChange: (listener) => - this.tokensController.subscribe(listener), - onCurrencyRateStateChange: (listener) => - this.controllerMessenger.subscribe( - `${this.currencyRateController.name}:stateChange`, - listener, - ), - onNetworkStateChange: (cb) => - this.networkController.store.subscribe((networkState) => { - const modifiedNetworkState = { - ...networkState, - provider: { - ...networkState.provider, - chainId: hexToDecimal(networkState.provider.chainId), - }, - }; - return cb(modifiedNetworkState); - }), - }); + this.tokenRatesController = new TokenRatesController( + { + onTokensStateChange: (listener) => + this.tokensController.subscribe(listener), + onCurrencyRateStateChange: (listener) => + this.controllerMessenger.subscribe( + `${this.currencyRateController.name}:stateChange`, + listener, + ), + onNetworkStateChange: (cb) => + this.networkController.store.subscribe((networkState) => { + const modifiedNetworkState = { + ...networkState, + provider: { + ...networkState.provider, + chainId: hexToDecimal(networkState.provider.chainId), + }, + }; + return cb(modifiedNetworkState); + }), + }, + undefined, + initState.TokenRatesController, + ); this.ensController = new EnsController({ provider: this.provider, @@ -719,6 +725,7 @@ export default class MetamaskController extends EventEmitter { }); this.rateLimitController = new RateLimitController({ + state: initState.RateLimitController, messenger: this.controllerMessenger.getRestricted({ name: 'RateLimitController', }), @@ -1029,6 +1036,24 @@ export default class MetamaskController extends EventEmitter { // ensure isClientOpenAndUnlocked is updated when memState updates this.on('update', (memState) => this._onStateUpdate(memState)); + /** + * All controllers in Memstore but not in store. They are not persisted. + * On chrome profile re-start, they will be re-initialized. + */ + const resetOnRestartStore = { + AccountTracker: this.accountTracker.store, + TxController: this.txController.memStore, + TokenRatesController: this.tokenRatesController, + MessageManager: this.messageManager.memStore, + PersonalMessageManager: this.personalMessageManager.memStore, + DecryptMessageManager: this.decryptMessageManager.memStore, + EncryptionPublicKeyManager: this.encryptionPublicKeyManager.memStore, + TypesMessageManager: this.typedMessageManager.memStore, + SwapsController: this.swapsController.store, + EnsController: this.ensController.store, + ApprovalController: this.approvalController, + }; + this.store.updateStructure({ AppStateController: this.appStateController.store, TransactionController: this.txController.store, @@ -1057,21 +1082,14 @@ export default class MetamaskController extends EventEmitter { CronjobController: this.cronjobController, NotificationController: this.notificationController, ///: END:ONLY_INCLUDE_IN + ...resetOnRestartStore, }); this.memStore = new ComposableObservableStore({ config: { AppStateController: this.appStateController.store, NetworkController: this.networkController.store, - AccountTracker: this.accountTracker.store, - TxController: this.txController.memStore, CachedBalancesController: this.cachedBalancesController.store, - TokenRatesController: this.tokenRatesController, - MessageManager: this.messageManager.memStore, - PersonalMessageManager: this.personalMessageManager.memStore, - DecryptMessageManager: this.decryptMessageManager.memStore, - EncryptionPublicKeyManager: this.encryptionPublicKeyManager.memStore, - TypesMessageManager: this.typedMessageManager.memStore, KeyringController: this.keyringController.memStore, PreferencesController: this.preferencesController.store, MetaMetricsController: this.metaMetricsController.store, @@ -1085,9 +1103,6 @@ export default class MetamaskController extends EventEmitter { PermissionLogController: this.permissionLogController.store, SubjectMetadataController: this.subjectMetadataController, BackupController: this.backupController, - SwapsController: this.swapsController.store, - EnsController: this.ensController.store, - ApprovalController: this.approvalController, AnnouncementController: this.announcementController, GasFeeController: this.gasFeeController, TokenListController: this.tokenListController, @@ -1099,11 +1114,36 @@ export default class MetamaskController extends EventEmitter { CronjobController: this.cronjobController, NotificationController: this.notificationController, ///: END:ONLY_INCLUDE_IN + ...resetOnRestartStore, }, controllerMessenger: this.controllerMessenger, }); this.memStore.subscribe(this.sendUpdate.bind(this)); + // if this is the first time, clear the state of by calling these methods + const resetMethods = [ + this.accountTracker.resetState, + this.txController.resetState, + this.messageManager.resetState, + this.personalMessageManager.resetState, + this.decryptMessageManager.resetState, + this.encryptionPublicKeyManager.resetState, + this.typedMessageManager.resetState, + this.swapsController.resetState, + this.ensController.resetState, + this.approvalController.clear.bind(this.approvalController), + // WE SHOULD ADD TokenListController.resetState here too. But it's not implemented yet. + ]; + + if (isManifestV3) { + if (globalThis.isFirstTimeProfileLoaded === true) { + this.resetStates(resetMethods); + } + } else { + // it's always the first time in MV2 + this.resetStates(resetMethods); + } + const password = process.env.CONF?.PASSWORD; if ( password && @@ -1137,6 +1177,18 @@ export default class MetamaskController extends EventEmitter { checkForMultipleVersionsRunning(); } + resetStates(resetMethods) { + resetMethods.forEach((resetMethod) => { + try { + resetMethod(); + } catch (err) { + console.error(err); + } + }); + + globalThis.isFirstTimeProfileLoaded = false; + } + ///: BEGIN:ONLY_INCLUDE_IN(flask) /** * Constructor helper for getting Snap permission specifications. diff --git a/app/scripts/metamask-controller.test.js b/app/scripts/metamask-controller.test.js index 5fc7232ce4e9..b2aaf665f805 100644 --- a/app/scripts/metamask-controller.test.js +++ b/app/scripts/metamask-controller.test.js @@ -102,11 +102,14 @@ const CUSTOM_RPC_CHAIN_ID = '0x539'; describe('MetaMaskController', function () { let metamaskController; + const sandbox = sinon.createSandbox(); const noop = () => undefined; before(async function () { + globalThis.isFirstTimeProfileLoaded = true; await ganacheServer.start(); + sinon.spy(MetaMaskController.prototype, 'resetStates'); }); beforeEach(function () { @@ -160,6 +163,18 @@ describe('MetaMaskController', function () { await ganacheServer.quit(); }); + describe('should reset states on first time profile load', function () { + it('should reset state', function () { + assert(metamaskController.resetStates.calledOnce); + assert.equal(globalThis.isFirstTimeProfileLoaded, false); + }); + + it('should not reset states if already set', function () { + // global.isFirstTime should also remain false + assert.equal(globalThis.isFirstTimeProfileLoaded, false); + }); + }); + describe('#getAccounts', function () { it('returns first address when dapp calls web3.eth.getAccounts', async function () { const password = 'a-fake-password';