Add Snap export handler usage tracking

Use registry for metadata

Replace controller with hook approach

Author: david0xd

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/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,7 @@ import {
   getSnapFiles,
   permissionsDiff,
   setDiff,
+  throttleTracking,
 } from './utils';
 import { SnapEndowments } from '../../snaps-rpc-methods/src/endowments';
 
@@ -221,3 +222,99 @@ describe('debouncePersistState', () => {
     expect(fn).toHaveBeenNthCalledWith(4, MOCK_LOCAL_SNAP_ID, {}, false);
   });
 });
+
+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);
+  });
+});

packages/snaps-controllers/src/utils.ts

@@ -7,6 +7,7 @@ import {
   getValidatedLocalizationFiles,
   validateAuxiliaryFiles,
   validateFetchedSnap,
+  HandlerType,
 } from '@metamask/snaps-utils';
 import type { Json } from '@metamask/utils';
 import deepEqual from 'fast-deep-equal';
@@ -375,3 +376,68 @@ export function debouncePersistState(
     );
   };
 }
+
+/**
+ * Handlers allowed for tracking.
+ */
+export const TRACKABLE_HANDLERS = [
+  HandlerType.OnHomePage,
+  HandlerType.OnInstall,
+  HandlerType.OnNameLookup,
+  HandlerType.OnRpcRequest,
+  HandlerType.OnSignature,
+  HandlerType.OnTransaction,
+  HandlerType.OnUpdate,
+] as const;
+
+/**
+ * A union type representing all possible trackable handler types.
+ */
+export type TrackableHandler = (typeof TRACKABLE_HANDLERS)[number];
+
+/**
+ * Throttles event tracking calls per unique combination of parameters.
+ *
+ * @param fn - The tracking function to throttle.
+ * @param timeout - The timeout in milliseconds. Defaults to 60000 (1 minute).
+ * @returns The throttled function.
+ */
+export function throttleTracking(
+  fn: (
+    snapId: SnapId,
+    handler: TrackableHandler,
+    success: boolean,
+    origin: string,
+  ) => Promise<void>,
+  timeout = 60000,
+) {
+  const previousCalls = new Map<string, number>();
+
+  return async (
+    snapId: SnapId,
+    handler: TrackableHandler,
+    success: boolean,
+    origin: string,
+  ): Promise<void> => {
+    const key = `${snapId}${handler}${origin}`;
+    const now = Date.now();
+    const lastCall = previousCalls.get(key) ?? 0;
+
+    if (now - lastCall >= timeout) {
+      previousCalls.set(key, now);
+      await fn(snapId, handler, success, origin);
+    }
+  };
+}
+
+/**
+ * Whether the handler type if allowed for tracking.
+ *
+ * @param handler Type of a handler.
+ * @returns True if handler is allowed for tracking, false otherwise.
+ */
+export function isTrackableHandler(
+  handler: HandlerType,
+): handler is TrackableHandler {
+  return TRACKABLE_HANDLERS.includes(handler as TrackableHandler);
+}