From 6296ecb3a9c20d6c3e5e6105fc324a31f0e0a944 Mon Sep 17 00:00:00 2001 From: david0xd Date: Wed, 12 Oct 2022 18:21:24 +0200 Subject: [PATCH] Add some refactoring and new features for the CronjobController --- ...vice.test.ts => CronjobController.test.ts} | 50 +++--- .../src/services/CronjobController.ts | 159 ++++++++++++++---- .../controllers/src/test-utils/controller.ts | 64 +++++++ 3 files changed, 217 insertions(+), 56 deletions(-) rename packages/controllers/src/services/{CronjobService.test.ts => CronjobController.test.ts} (70%) diff --git a/packages/controllers/src/services/CronjobService.test.ts b/packages/controllers/src/services/CronjobController.test.ts similarity index 70% rename from packages/controllers/src/services/CronjobService.test.ts rename to packages/controllers/src/services/CronjobController.test.ts index 29c26a0439..b728115898 100644 --- a/packages/controllers/src/services/CronjobService.test.ts +++ b/packages/controllers/src/services/CronjobController.test.ts @@ -1,7 +1,7 @@ import { SnapEndowments } from '../snaps'; -import { getSnapControllerMessenger } from '../test-utils'; +import { getRestrictedCronjobControllerMessenger } from '../test-utils'; import { Timer } from '../snaps/Timer'; -import { Cronjob, CronjobService } from './CronjobService'; +import { Cronjob, CronjobController } from './CronjobController'; jest.mock('../snaps/Timer'); const MOCK_ORIGIN = 'metamask.io'; @@ -37,7 +37,7 @@ describe('CronjobService', () => { const mockSnapId = 'snap-id'; it('registers a snap', async () => { - const controllerMessenger = getSnapControllerMessenger(); + const controllerMessenger = getRestrictedCronjobControllerMessenger(); const callActionMock = jest .spyOn(controllerMessenger, 'call') @@ -50,11 +50,11 @@ describe('CronjobService', () => { return false; }); - const cronjobService = new CronjobService({ + const cronjobController = new CronjobController({ messenger: controllerMessenger, }); - await cronjobService.register(mockSnapId); + await cronjobController.register(mockSnapId); expect(callActionMock).toHaveBeenCalledWith( 'PermissionController:hasPermission', @@ -74,23 +74,23 @@ describe('CronjobService', () => { id: 'cronjob-id', request: { method: 'call_this_method', params: [] }, }; - const controllerMessenger = getSnapControllerMessenger(); + const controllerMessenger = getRestrictedCronjobControllerMessenger(); - const cronjobService = new CronjobService({ + const cronjobController = new CronjobController({ messenger: controllerMessenger, }); - await cronjobService.schedule(mockSnapId, mockJob); + await cronjobController.schedule(mockSnapId, mockJob); expect(Timer.prototype.start).toHaveBeenCalled(); - expect(cronjobService.getJob(mockJob.id)).toStrictEqual({ + expect(cronjobController.getJob(mockJob.id)).toStrictEqual({ ...mockJob, timer: expect.anything(), }); }); it('unregisters scheduled jobs for a snap', async () => { - const controllerMessenger = getSnapControllerMessenger(); + const controllerMessenger = getRestrictedCronjobControllerMessenger(); jest.spyOn(controllerMessenger, 'call').mockImplementation((method) => { if (method === 'PermissionController:hasPermission') { @@ -101,24 +101,24 @@ describe('CronjobService', () => { return false; }); - const cronjobService = new CronjobService({ + const cronjobController = new CronjobController({ messenger: controllerMessenger, }); - await cronjobService.register(mockSnapId); - const registeredJobIds = cronjobService.getRegisteredJobs(mockSnapId); + await cronjobController.register(mockSnapId); + const registeredJobIds = cronjobController.getRegisteredJobs(mockSnapId); const jobOneId = registeredJobIds?.[0] as string; const jobTwoId = registeredJobIds?.[1] as string; expect(registeredJobIds?.length).toStrictEqual(2); - cronjobService.unregister(mockSnapId); - expect(cronjobService.getRegisteredJobs(mockSnapId)).toBeUndefined(); - expect(cronjobService.getJob(jobOneId)).toBeUndefined(); - expect(cronjobService.getJob(jobTwoId)).toBeUndefined(); + cronjobController.unregister(mockSnapId); + expect(cronjobController.getRegisteredJobs(mockSnapId)).toStrictEqual([]); + expect(cronjobController.getJob(jobOneId)).toBeUndefined(); + expect(cronjobController.getJob(jobTwoId)).toBeUndefined(); }); it('should not register jobs if snap has no cronjob permission', async () => { - const controllerMessenger = getSnapControllerMessenger(); + const controllerMessenger = getRestrictedCronjobControllerMessenger(); const callActionMock = jest .spyOn(controllerMessenger, 'call') @@ -126,11 +126,11 @@ describe('CronjobService', () => { return false; }); - const cronjobService = new CronjobService({ + const cronjobController = new CronjobController({ messenger: controllerMessenger, }); - await cronjobService.register(mockSnapId); + await cronjobController.register(mockSnapId); expect(callActionMock).toHaveBeenCalledWith( 'PermissionController:hasPermission', @@ -143,11 +143,11 @@ describe('CronjobService', () => { mockSnapId, ); - expect(cronjobService.getRegisteredJobs(mockSnapId)).toBeUndefined(); + expect(cronjobController.getRegisteredJobs(mockSnapId)).toStrictEqual([]); }); it('should not register any jobs if snap has no cronjob caveats', async () => { - const controllerMessenger = getSnapControllerMessenger(); + const controllerMessenger = getRestrictedCronjobControllerMessenger(); const callActionMock = jest .spyOn(controllerMessenger, 'call') @@ -160,11 +160,11 @@ describe('CronjobService', () => { return false; }); - const cronjobService = new CronjobService({ + const cronjobController = new CronjobController({ messenger: controllerMessenger, }); - await cronjobService.register(mockSnapId); + await cronjobController.register(mockSnapId); expect(callActionMock).toHaveBeenCalledWith( 'PermissionController:hasPermission', @@ -177,6 +177,6 @@ describe('CronjobService', () => { mockSnapId, ); - expect(cronjobService.getRegisteredJobs(mockSnapId)).toBeUndefined(); + expect(cronjobController.getRegisteredJobs(mockSnapId)).toStrictEqual([]); }); }); diff --git a/packages/controllers/src/services/CronjobController.ts b/packages/controllers/src/services/CronjobController.ts index 38e1766219..56feea5ef8 100644 --- a/packages/controllers/src/services/CronjobController.ts +++ b/packages/controllers/src/services/CronjobController.ts @@ -4,30 +4,48 @@ import { HasPermission, GetPermissions, } from '@metamask/controllers'; -import { HandlerType, SnapId } from '@metamask/snap-utils'; +import { HandlerType, SnapId, TruncatedSnap } from '@metamask/snap-utils'; import { parseExpression } from 'cron-parser'; import { nanoid } from 'nanoid'; -import { GetSnap, HandleSnapRequest, SnapAdded, SnapEndowments } from '..'; +import { + GetSnap, + HandleSnapRequest, + SnapAdded, + SnapBlocked, + SnapEndowments, + SnapInstalled, + SnapRemoved, + SnapTerminated, + SnapUnblocked, + SnapUpdated, +} from '..'; import { Timer } from '../snaps/Timer'; -export type CronjobServiceActions = +export type CronjobControllerActions = | GetSnap | HandleSnapRequest | HasPermission | GetPermissions; -export type CronjobServiceEvents = SnapAdded; +export type CronjobControllerEvents = + | SnapAdded + | SnapBlocked + | SnapInstalled + | SnapRemoved + | SnapUnblocked + | SnapUpdated + | SnapTerminated; -export type CronjobServiceMessenger = RestrictedControllerMessenger< - 'CronjobService', - CronjobServiceActions, - CronjobServiceEvents, - CronjobServiceActions['type'], - CronjobServiceEvents['type'] +export type CronjobControllerMessenger = RestrictedControllerMessenger< + 'CronjobController', + CronjobControllerActions, + CronjobControllerEvents, + CronjobControllerActions['type'], + CronjobControllerEvents['type'] >; -export type CronjobServiceArgs = { - messenger: CronjobServiceMessenger; +export type CronjobControllerArgs = { + messenger: CronjobControllerMessenger; }; export type CronjobDefinition = { @@ -43,27 +61,31 @@ export type Cronjob = { id: string; } & CronjobDefinition; +export type StoredJobInformation = { + lastRun: number; +}; + export type CronjobControllerState = { snaps: Record; - jobs: Record; + jobs: Record; }; -const controllerName = 'cronjobController'; +const controllerName = 'CronjobController'; /** - * Use this service to register and schedule periodically executed jobs + * Use this controller to register and schedule periodically executed jobs * using RPC method hooks. */ export class CronjobController extends BaseController< typeof controllerName, CronjobControllerState, - CronjobServiceMessenger + CronjobControllerMessenger > { - private _messenger: CronjobServiceMessenger; + private _messenger: CronjobControllerMessenger; private _jobs: Map; - constructor({ messenger }: CronjobServiceArgs) { + constructor({ messenger }: CronjobControllerArgs) { super({ messenger, metadata: { @@ -79,6 +101,26 @@ export class CronjobController extends BaseController< this._jobs = new Map(); this._messenger = messenger; this.initializeJobSchedule(); + + this._handleEventSnapInstalled = this._handleEventSnapInstalled.bind(this); + this._handleEventSnapRemoved = this._handleEventSnapRemoved.bind(this); + this._handleEventSnapUpdated = this._handleEventSnapUpdated.bind(this); + + // Subscribe to Snap events + this.messagingSystem.subscribe( + 'SnapController:snapInstalled', + this._handleEventSnapInstalled, + ); + + this.messagingSystem.subscribe( + 'SnapController:snapRemoved', + this._handleEventSnapRemoved, + ); + + this.messagingSystem.subscribe( + 'SnapController:snapInstalled', + this._handleEventSnapUpdated, + ); } /** @@ -138,6 +180,7 @@ export class CronjobController extends BaseController< const ms = next.getTime() - now.getTime(); const timer = new Timer(ms); timer.start(() => { + this.updateJobLastRunState(job.id, next.getTime()); this._messenger.call('SnapController:handleRequest', { snapId, // @todo Decide on origin for requests like this @@ -148,12 +191,7 @@ export class CronjobController extends BaseController< this.schedule(snapId, job); }); - this.update((state) => { - state.jobs[job.id] = { - expression: job.expression, - request: job.request, - }; - }); + this.updateJobLastRunState(job.id, 0); // 0 for init, never ran actually this._jobs.set(job.id, { ...job, timer }); } @@ -175,6 +213,20 @@ export class CronjobController extends BaseController< }); } + /** + * Update time of a last run for the Cronjob specified by ID. + * + * @param jobId - ID of a cron job. + * @param lastRun - Unix timestamp when the job was last ran. + */ + updateJobLastRunState(jobId: string, lastRun: number) { + this.update((state) => { + state.jobs[jobId] = { + lastRun, + }; + }); + } + /** * Retrieve job specification by given job ID. * @@ -191,7 +243,7 @@ export class CronjobController extends BaseController< * @param jobId - ID of a job. * @returns Plain job specification. */ - getJobFromState(jobId: string): CronjobDefinition { + getJobInfoFromState(jobId: string): StoredJobInformation { return this.state.jobs[jobId]; } @@ -209,11 +261,56 @@ export class CronjobController extends BaseController< * Initialize (schedule) jobs when controller starts. */ initializeJobSchedule() { - Object.keys(this.state.snaps).forEach((snapId) => { - this.state.snaps[snapId].forEach((jobId) => { - const jobDefinition = this.getJobFromState(jobId); - this.schedule(snapId, { id: jobId, ...jobDefinition }); - }); - }); + // TODO: Do some init actions after restart? + } + + /** + * Run controller teardown process and unsubscribe from everything. + */ + protected destroy() { + super.destroy(); + + this.messagingSystem.unsubscribe( + 'SnapController:snapInstalled', + this._handleEventSnapInstalled, + ); + + this.messagingSystem.unsubscribe( + 'SnapController:snapRemoved', + this._handleEventSnapRemoved, + ); + + this.messagingSystem.unsubscribe( + 'SnapController:snapUpdated', + this._handleEventSnapUpdated, + ); + } + + /** + * Handle cron jobs on 'snapInstalled' event. + * + * @param snap - Basic Snap information. + */ + async _handleEventSnapInstalled(snap: TruncatedSnap) { + await this.register(snap.id); + } + + /** + * Handle cron jobs on 'snapRemoved' event. + * + * @param snap - Basic Snap information. + */ + _handleEventSnapRemoved(snap: TruncatedSnap) { + this.unregister(snap.id); + } + + /** + * Handle cron jobs on 'snapUpdated' event. + * + * @param snap - Basic Snap information. + */ + async _handleEventSnapUpdated(snap: TruncatedSnap) { + this.unregister(snap.id); + await this.register(snap.id); } } diff --git a/packages/controllers/src/test-utils/controller.ts b/packages/controllers/src/test-utils/controller.ts index a202c090c3..9a3dd41dad 100644 --- a/packages/controllers/src/test-utils/controller.ts +++ b/packages/controllers/src/test-utils/controller.ts @@ -8,6 +8,10 @@ import { SnapControllerEvents, SnapEndowments, } from '../snaps'; +import { + CronjobControllerActions, + CronjobControllerEvents, +} from '../services/CronjobController'; import { getNodeEES, getNodeEESMessenger } from './execution-environment'; export const getControllerMessenger = () => @@ -184,3 +188,63 @@ export const getSnapControllerWithEES = ( const controller = new SnapController(options); return [controller, _service] as const; }; + +// ------------------------------------------------ +// Mock controller messenger for Cronjob Controller +export const getCronjobControllerMessenger = () => + new ControllerMessenger< + CronjobControllerActions | AllowedActions, + CronjobControllerEvents | AllowedEvents + >(); + +export const getRestrictedCronjobControllerMessenger = ( + messenger: ReturnType< + typeof getCronjobControllerMessenger + > = getCronjobControllerMessenger(), + mocked = true, +) => { + const m = messenger.getRestricted< + 'CronjobController', + CronjobControllerActions['type'] | AllowedActions['type'], + CronjobControllerEvents['type'] | AllowedEvents['type'] + >({ + name: 'CronjobController', + allowedEvents: [ + 'ExecutionService:unhandledError', + 'ExecutionService:outboundRequest', + 'ExecutionService:outboundResponse', + 'SnapController:snapAdded', + 'SnapController:snapBlocked', + 'SnapController:snapInstalled', + 'SnapController:snapUnblocked', + 'SnapController:snapUpdated', + 'SnapController:snapRemoved', + ], + allowedActions: [ + 'ApprovalController:addRequest', + 'ExecutionService:executeSnap', + 'ExecutionService:terminateAllSnaps', + 'ExecutionService:terminateSnap', + 'ExecutionService:handleRpcRequest', + 'PermissionController:getEndowments', + 'PermissionController:hasPermission', + 'PermissionController:hasPermissions', + 'PermissionController:getPermissions', + 'SnapController:handleRequest', + ], + }); + + if (mocked) { + jest.spyOn(m, 'call').mockImplementation((method, ...args) => { + // Return false for long-running by default, and true for everything else. + if ( + method === 'PermissionController:hasPermission' && + args[1] === SnapEndowments.LongRunning + ) { + return false; + } + return true; + }); + } + return m; +};