CronjobController.ts

import {
  BaseControllerV2 as BaseController,
  RestrictedControllerMessenger,
  HasPermission,
  GetPermissions,
} from '@metamask/controllers';
import { HandlerType, SnapId, TruncatedSnap } from '@metamask/snap-utils';
import { parseExpression } from 'cron-parser';
import { nanoid } from 'nanoid';
import {
  GetSnap,
  HandleSnapRequest,
  SnapAdded,
  SnapBlocked,
  SnapEndowments,
  SnapInstalled,
  SnapRemoved,
  SnapTerminated,
  SnapUnblocked,
  SnapUpdated,
} from '..';
import { Timer } from '../snaps/Timer';

export type CronjobControllerActions =
  | GetSnap
  | HandleSnapRequest
  | HasPermission
  | GetPermissions;

export type CronjobControllerEvents =
  | SnapAdded
  | SnapBlocked
  | SnapInstalled
  | SnapRemoved
  | SnapUnblocked
  | SnapUpdated
  | SnapTerminated;

export type CronjobControllerMessenger = RestrictedControllerMessenger<
  'CronjobController',
  CronjobControllerActions,
  CronjobControllerEvents,
  CronjobControllerActions['type'],
  CronjobControllerEvents['type']
>;

export type CronjobControllerArgs = {
  messenger: CronjobControllerMessenger;
};

export type CronjobDefinition = {
  expression: string;
  request: {
    method: string;
    params: string[] | number[] | boolean[];
  };
};

export type Cronjob = {
  timer?: Timer;
  id: string;
} & CronjobDefinition;

export type StoredJobInformation = {
  lastRun: number;
};

export type CronjobControllerState = {
  snaps: Record<SnapId, string[]>;
  jobs: Record<string, StoredJobInformation>;
};

const controllerName = 'CronjobController';

/**
 * Use this controller to register and schedule periodically executed jobs
 * using RPC method hooks.
 */
export class CronjobController extends BaseController<
  typeof controllerName,
  CronjobControllerState,
  CronjobControllerMessenger
> {
  private _messenger: CronjobControllerMessenger;

  private _jobs: Map<string, Cronjob>;

  constructor({ messenger }: CronjobControllerArgs) {
    super({
      messenger,
      metadata: {
        snaps: { persist: true, anonymous: false },
        jobs: { persist: true, anonymous: false },
      },
      name: controllerName,
      state: {
        snaps: {},
        jobs: {},
      },
    });
    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,
    );
  }

  /**
   * Register cron jobs for a given snap by getting specification from a permission caveats.
   * Once registered, each job will be scheduled.
   *
   * @param snapId - ID of a snap.
   */
  async register(snapId: SnapId) {
    const hasCronjob = await this._messenger.call(
      'PermissionController:hasPermission',
      snapId,
      SnapEndowments.cronjob,
    );
    if (!hasCronjob) {
      return;
    }
    const permissions = await this._messenger.call(
      'PermissionController:getPermissions',
      snapId,
    );
    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
    const cronjobPermission = permissions![SnapEndowments.cronjob]!;
    const { caveats } = cronjobPermission;

    if (!caveats) {
      return;
    }

    const definitions: CronjobDefinition[] = caveats[0].value;

    const jobs = definitions?.map((definition) => {
      const id = nanoid();
      return { ...definition, id };
    });

    this.update((state) => {
      state.snaps[snapId] = jobs?.map((job) => job.id);
    });

    jobs.forEach((job) => this.schedule(snapId, job));
  }

  /**
   * Schedule job.
   *
   * @param snapId - ID of a snap.
   * @param job - Cronjob specification.
   */
  schedule(snapId: SnapId, job: Cronjob) {
    const parsed = parseExpression(job.expression);
    if (!parsed.hasNext()) {
      return;
    }
    const next = parsed.next();
    const now = new Date();
    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
        origin: 'METAMASK',
        handler: HandlerType.OnCronjob,
        request: job.request,
      });
      this.schedule(snapId, job);
    });

    this.updateJobLastRunState(job.id, 0); // 0 for init, never ran actually
    this._jobs.set(job.id, { ...job, timer });
  }

  /**
   * Unregister all jobs related to the given snapId.
   *
   * @param snapId - ID of a snap.
   */
  unregister(snapId: SnapId) {
    const jobs = this.state.snaps[snapId];
    jobs?.forEach((id) => {
      const job = this._jobs.get(id);
      job?.timer?.cancel();
      this._jobs.delete(job?.id as string);
    });

    this.update((state) => {
      delete state.snaps[snapId];
    });
  }

  /**
   * 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.
   *
   * @param jobId - ID of a job.
   * @returns Job specification.
   */
  getJob(jobId: string): Cronjob | undefined {
    return this._jobs.get(jobId);
  }

  /**
   * Retrieve job specification from state by given job ID.
   *
   * @param jobId - ID of a job.
   * @returns Plain job specification.
   */
  getJobInfoFromState(jobId: string): StoredJobInformation {
    return this.state.jobs[jobId];
  }

  /**
   * Retrieve an array of job ID registered for a given snap identified by ID.
   *
   * @param snapId - ID of a snap.
   * @returns All registered job ID for single snap.
   */
  getRegisteredJobs(snapId: string): string[] {
    return this.state.snaps[snapId] ? this.state.snaps[snapId] : [];
  }

  /**
   * Initialize (schedule) jobs when controller starts.
   */
  initializeJobSchedule() {
    // 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);
  }
}