Add Snap export handler usage tracking

Use registry for metadata

Replace controller with hook approach

Fix broken unit tests

Add unit tests

Author: david0xd

packages/snaps-controllers/src/snaps/SnapController.test.tsx

@@ -4702,7 +4702,7 @@ describe('SnapController', () => {
         },
       });
 
-      expect(rootMessenger.call).toHaveBeenCalledTimes(4);
+      expect(rootMessenger.call).toHaveBeenCalledTimes(5);
       expect(rootMessenger.call).toHaveBeenCalledWith(
         'ExecutionService:handleRpcRequest',
         MOCK_SNAP_ID,
@@ -11061,4 +11061,108 @@ describe('SnapController', () => {
       snapController.destroy();
     });
   });
+
+  describe('SnapController:trackEvent', () => {
+    it('should track event for allowed handler', async () => {
+      const mockTrackEvent = jest.fn();
+      const rootMessenger = getControllerMessenger();
+      const executionEnvironmentStub = new ExecutionEnvironmentStub(
+        getNodeEESMessenger(rootMessenger),
+      ) as unknown as NodeThreadExecutionService;
+
+      const [snapController] = getSnapControllerWithEES(
+        getSnapControllerWithEESOptions({
+          rootMessenger,
+          trackEvent: mockTrackEvent,
+          state: {
+            snaps: getPersistedSnapsState(),
+          },
+        }),
+        executionEnvironmentStub,
+      );
+
+      const snap = snapController.getExpect(MOCK_SNAP_ID);
+      await snapController.startSnap(snap.id);
+
+      await snapController.handleRequest({
+        snapId: snap.id,
+        origin: MOCK_ORIGIN,
+        handler: HandlerType.OnRpcRequest,
+        request: {
+          jsonrpc: '2.0',
+          method: 'test',
+          params: {},
+          id: 1,
+        },
+      });
+
+      expect(mockTrackEvent).toHaveBeenCalledTimes(1);
+      expect(mockTrackEvent).toHaveBeenCalledWith({
+        category: 'Snaps',
+        event: 'SnapExportUsed',
+        properties: {
+          export: 'onRpcRequest',
+          origin: 'https://example.com',
+          // eslint-disable-next-line @typescript-eslint/naming-convention
+          snap_category: null,
+          // eslint-disable-next-line @typescript-eslint/naming-convention
+          snap_id: 'npm:@metamask/example-snap',
+          success: true,
+        },
+      });
+      snapController.destroy();
+    });
+
+    it('should not track event for disallowed handler', async () => {
+      const mockTrackEvent = jest.fn();
+      const rootMessenger = getControllerMessenger();
+
+      rootMessenger.registerActionHandler(
+        'PermissionController:getPermissions',
+        () => ({
+          [SnapEndowments.Cronjob]: {
+            caveats: [{ type: SnapCaveatType.SnapCronjob, value: '* * * * *' }],
+            date: 1664187844588,
+            id: 'izn0WGUO8cvq_jqvLQuQP',
+            invoker: MOCK_SNAP_ID,
+            parentCapability: SnapEndowments.Cronjob,
+          },
+        }),
+      );
+
+      const executionEnvironmentStub = new ExecutionEnvironmentStub(
+        getNodeEESMessenger(rootMessenger),
+      ) as unknown as NodeThreadExecutionService;
+
+      const [snapController] = getSnapControllerWithEES(
+        getSnapControllerWithEESOptions({
+          environmentEndowmentPermissions: ['endowment:cronjob'],
+          rootMessenger,
+          trackEvent: mockTrackEvent,
+          state: {
+            snaps: getPersistedSnapsState(),
+          },
+        }),
+        executionEnvironmentStub,
+      );
+
+      const snap = snapController.getExpect(MOCK_SNAP_ID);
+      await snapController.startSnap(snap.id);
+
+      await snapController.handleRequest({
+        snapId: snap.id,
+        origin: MOCK_ORIGIN,
+        handler: HandlerType.OnCronjob,
+        request: {
+          jsonrpc: '2.0',
+          method: 'test',
+          params: {},
+          id: 1,
+        },
+      });
+
+      expect(mockTrackEvent).not.toHaveBeenCalled();
+      snapController.destroy();
+    });
+  });
 });

packages/snaps-controllers/src/snaps/SnapController.ts

@@ -176,7 +176,9 @@ import {
   hasTimedOut,
   permissionsDiff,
   setDiff,
+  throttleTracking,
   withTimeout,
+  isTrackableHandler,
 } from '../utils';
 
 export const controllerName = 'SnapController';
@@ -775,6 +777,11 @@ type SnapControllerArgs = {
    * object to fall back to the default cryptographic functions.
    */
   clientCryptography?: CryptographicFunctions;
+
+  /**
+   * MetaMetrics event tracking hook.
+   */
+  trackEvent: TrackEventHook;
 };
 
 type AddSnapArgs = {
@@ -795,6 +802,14 @@ type SetSnapArgs = Omit<AddSnapArgs, 'location' | 'versionRange'> & {
   hideSnapBranding?: boolean;
 };
 
+type TrackingEventPayload = {
+  event: string;
+  category: string;
+  properties: Record<string, Json>;
+};
+
+type TrackEventHook = (event: TrackingEventPayload) => void;
+
 const defaultState: SnapControllerState = {
   snaps: {},
   snapStates: {},
@@ -880,6 +895,10 @@ export class SnapController extends BaseController<
 
   readonly #preinstalledSnaps: PreinstalledSnap[] | null;
 
+  readonly #trackEvent: TrackEventHook;
+
+  readonly #trackSnapExport: ReturnType<typeof throttleTracking>;
+
   constructor({
     closeAllConnections,
     messenger,
@@ -898,6 +917,7 @@ export class SnapController extends BaseController<
     getMnemonicSeed,
     getFeatureFlags = () => ({}),
     clientCryptography,
+    trackEvent,
   }: SnapControllerArgs) {
     super({
       messenger,
@@ -960,6 +980,7 @@ export class SnapController extends BaseController<
     this._onOutboundResponse = this._onOutboundResponse.bind(this);
     this.#rollbackSnapshots = new Map();
     this.#snapsRuntimeData = new Map();
+    this.#trackEvent = trackEvent;
 
     this.#pollForLastRequestStatus();
 
@@ -1025,6 +1046,30 @@ export class SnapController extends BaseController<
     Object.values(this.state?.snaps ?? {}).forEach((snap) =>
       this.#setupRuntime(snap.id),
     );
+
+    this.#trackSnapExport = throttleTracking(
+      async (
+        snapId: SnapId,
+        handler: string,
+        success: boolean,
+        origin: string,
+      ) => {
+        const snapMetadata = await this.getRegistryMetadata(snapId);
+        this.#trackEvent({
+          event: 'SnapExportUsed',
+          category: 'Snaps',
+          properties: {
+            // eslint-disable-next-line @typescript-eslint/naming-convention
+            snap_id: snapId,
+            export: handler,
+            // eslint-disable-next-line @typescript-eslint/naming-convention
+            snap_category: snapMetadata?.category ?? null,
+            success,
+            origin,
+          },
+        });
+      },
+    );
   }
 
   /**
@@ -3584,10 +3629,19 @@ export class SnapController extends BaseController<
 
         this.#recordSnapRpcRequestFinish(snapId, transformedRequest.id);
 
+        if (isTrackableHandler(handlerType)) {
+          await this.#trackSnapExport(snapId, handlerType, true, origin);
+        }
+
         return transformedResult;
       } catch (error) {
         // We flag the RPC request as finished early since termination may affect pending requests
         this.#recordSnapRpcRequestFinish(snapId, transformedRequest.id);
+
+        if (isTrackableHandler(handlerType)) {
+          await this.#trackSnapExport(snapId, handlerType, false, origin);
+        }
+
         const [jsonRpcError, handled] = unwrapError(error);
 
         if (!handled) {

packages/snaps-controllers/src/test-utils/controller.ts

@@ -593,6 +593,7 @@ export const getSnapControllerOptions = (
       Promise.resolve(TEST_SECRET_RECOVERY_PHRASE_SEED_BYTES),
     clientCryptography: {},
     encryptor: getSnapControllerEncryptor(),
+    trackEvent: jest.fn(),
     ...opts,
   } as SnapControllerConstructorParams;
 
@@ -626,6 +627,7 @@ export const getSnapControllerWithEESOptions = ({
       Promise.resolve(TEST_SECRET_RECOVERY_PHRASE_SEED_BYTES),
     encryptor: getSnapControllerEncryptor(),
     fetchFunction: jest.fn(),
+    trackEvent: jest.fn(),
     ...options,
   } as SnapControllerConstructorParams & {
     rootMessenger: ReturnType<typeof getControllerMessenger>;

packages/snaps-controllers/src/utils.test.ts

@@ -1,4 +1,4 @@
-import { VirtualFile } from '@metamask/snaps-utils';
+import { HandlerType, VirtualFile } from '@metamask/snaps-utils';
 import {
   getMockSnapFiles,
   getSnapManifest,
@@ -20,6 +20,8 @@ import {
   getSnapFiles,
   permissionsDiff,
   setDiff,
+  throttleTracking,
+  TRACKABLE_HANDLERS,
 } from './utils';
 import { SnapEndowments } from '../../snaps-rpc-methods/src/endowments';
 
@@ -221,3 +223,122 @@ describe('debouncePersistState', () => {
     expect(fn).toHaveBeenNthCalledWith(4, MOCK_LOCAL_SNAP_ID, {}, false);
   });
 });
+
+describe('TRACKABLE_HANDLERS', () => {
+  it('should contain the expected handler types', () => {
+    expect(TRACKABLE_HANDLERS).toStrictEqual([
+      HandlerType.OnHomePage,
+      HandlerType.OnInstall,
+      HandlerType.OnNameLookup,
+      HandlerType.OnRpcRequest,
+      HandlerType.OnSignature,
+      HandlerType.OnTransaction,
+      HandlerType.OnUpdate,
+    ]);
+  });
+
+  it('should be a readonly array', () => {
+    expect(Object.isFrozen(TRACKABLE_HANDLERS)).toBe(true);
+  });
+
+  it('should contain unique values', () => {
+    const uniqueValues = new Set(TRACKABLE_HANDLERS);
+    expect(uniqueValues.size).toBe(TRACKABLE_HANDLERS.length);
+  });
+});
+
+describe('throttleTracking', () => {
+  beforeAll(() => {
+    jest.useFakeTimers();
+  });
+
+  afterAll(() => {
+    jest.useRealTimers();
+  });
+
+  it('throttles tracking calls based on unique combinations of snapId, handler, and origin', async () => {
+    const fn = jest.fn();
+    const throttled = throttleTracking(fn, 1000);
+
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin1');
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin1');
+    await throttled(MOCK_SNAP_ID, HandlerType.OnRpcRequest, true, 'origin1');
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin2');
+
+    expect(fn).toHaveBeenCalledTimes(3);
+    expect(fn).toHaveBeenNthCalledWith(
+      1,
+      MOCK_SNAP_ID,
+      HandlerType.OnHomePage,
+      true,
+      'origin1',
+    );
+    expect(fn).toHaveBeenNthCalledWith(
+      2,
+      MOCK_SNAP_ID,
+      HandlerType.OnRpcRequest,
+      true,
+      'origin1',
+    );
+    expect(fn).toHaveBeenNthCalledWith(
+      3,
+      MOCK_SNAP_ID,
+      HandlerType.OnHomePage,
+      true,
+      'origin2',
+    );
+
+    jest.advanceTimersByTime(500);
+
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin1');
+    await throttled(MOCK_SNAP_ID, HandlerType.OnRpcRequest, true, 'origin1');
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin2');
+
+    expect(fn).toHaveBeenCalledTimes(3);
+
+    jest.advanceTimersByTime(600);
+
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin1');
+    await throttled(MOCK_SNAP_ID, HandlerType.OnRpcRequest, true, 'origin1');
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin2');
+
+    expect(fn).toHaveBeenCalledTimes(6);
+    expect(fn).toHaveBeenNthCalledWith(
+      4,
+      MOCK_SNAP_ID,
+      HandlerType.OnHomePage,
+      true,
+      'origin1',
+    );
+    expect(fn).toHaveBeenNthCalledWith(
+      5,
+      MOCK_SNAP_ID,
+      HandlerType.OnRpcRequest,
+      true,
+      'origin1',
+    );
+    expect(fn).toHaveBeenNthCalledWith(
+      6,
+      MOCK_SNAP_ID,
+      HandlerType.OnHomePage,
+      true,
+      'origin2',
+    );
+  });
+
+  it('uses default timeout of 60000ms when no timeout is specified', async () => {
+    const fn = jest.fn();
+    const throttled = throttleTracking(fn);
+
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin1');
+    expect(fn).toHaveBeenCalledTimes(1);
+
+    jest.advanceTimersByTime(59999);
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin1');
+    expect(fn).toHaveBeenCalledTimes(1);
+
+    jest.advanceTimersByTime(2);
+    await throttled(MOCK_SNAP_ID, HandlerType.OnHomePage, true, 'origin1');
+    expect(fn).toHaveBeenCalledTimes(2);
+  });
+});