Add subject type and method to get metadata by origin (#950)

* Add subject type and method to get metadata by origin * Add comment to enum * Rename type field to subjectType
MetaMask · Oct 11, 2023 · ac14565 · ac14565
1 parent f3b27a9
commit ac14565
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 12 deletions.
diff --git a/src/subject-metadata/SubjectMetadataController.test.ts b/src/subject-metadata/SubjectMetadataController.test.ts
@@ -6,6 +6,7 @@ import {
   SubjectMetadataControllerActions,
   SubjectMetadataControllerEvents,
   SubjectMetadataControllerMessenger,
+  SubjectType,
 } from './SubjectMetadataController';
 
 const controllerName = 'SubjectMetadataController';
@@ -48,17 +49,20 @@ function getSubjectMetadataControllerMessenger() {
  *
  * @param origin - The subject's origin
  * @param name - Optional subject name
+ * @param subjectType - Optional subject type
  * @param opts - Optional extra options for the metadata
  * @returns The created metadata object
  */
 function getSubjectMetadata(
   origin: string,
-  name?: string,
+  name: string | null = null,
+  subjectType: SubjectType | null = null,
   opts?: Record<string, Json>,
 ) {
   return {
     origin,
-    name: name ?? null,
+    name,
+    subjectType,
     iconUrl: null,
     extensionId: null,
     ...opts,
@@ -212,6 +216,34 @@ describe('SubjectMetadataController', () => {
     });
   });
 
+  describe('getSubjectMetadata', () => {
+    it('returns the subject metadata for the given origin', () => {
+      const [messenger, hasPermissionsSpy] =
+        getSubjectMetadataControllerMessenger();
+      const controller = new SubjectMetadataController({
+        messenger,
+        subjectCacheLimit: 1,
+      });
+      hasPermissionsSpy.mockImplementationOnce(() => true);
+
+      controller.addSubjectMetadata(
+        getSubjectMetadata('foo.com', 'foo', SubjectType.Snap),
+      );
+
+      controller.addSubjectMetadata(
+        getSubjectMetadata('bar.io', 'bar', SubjectType.Website),
+      );
+
+      expect(controller.getSubjectMetadata('foo.com')).toStrictEqual(
+        getSubjectMetadata('foo.com', 'foo', SubjectType.Snap),
+      );
+
+      expect(controller.getSubjectMetadata('bar.io')).toStrictEqual(
+        getSubjectMetadata('bar.io', 'bar', SubjectType.Website),
+      );
+    });
+  });
+
   describe('trimMetadataState', () => {
     it('deletes all subjects without permissions from state', () => {
       const [messenger, hasPermissionsSpy] =

diff --git a/src/subject-metadata/SubjectMetadataController.ts b/src/subject-metadata/SubjectMetadataController.ts
@@ -13,16 +13,30 @@ const controllerName = 'SubjectMetadataController';
 
 type SubjectOrigin = string;
 
+/**
+ * The different kinds of subjects that MetaMask may interact with, including
+ * third parties and itself (e.g., when the background communicated with the UI).
+ */
+export enum SubjectType {
+  Extension = 'extension',
+  Internal = 'internal',
+  Unknown = 'unknown',
+  Website = 'website',
+  Snap = 'snap',
+}
+
 export type SubjectMetadata = PermissionSubjectMetadata & {
   [key: string]: Json;
   // TODO:TS4.4 make optional
   name: string | null;
+  subjectType: SubjectType | null;
   extensionId: string | null;
   iconUrl: string | null;
 };
 
 type SubjectMetadataToAdd = PermissionSubjectMetadata & {
   name?: string | null;
+  subjectType?: SubjectType | null;
   extensionId?: string | null;
   iconUrl?: string | null;
 } & Record<string, Json>;
@@ -44,7 +58,14 @@ export type GetSubjectMetadataState = {
   handler: () => SubjectMetadataControllerState;
 };
 
-export type SubjectMetadataControllerActions = GetSubjectMetadataState;
+export type GetSubjectMetadata = {
+  type: `${typeof controllerName}:getSubjectMetadata`;
+  handler: (origin: SubjectOrigin) => SubjectMetadata | undefined;
+};
+
+export type SubjectMetadataControllerActions =
+  | GetSubjectMetadataState
+  | GetSubjectMetadata;
 
 export type SubjectMetadataStateChange = {
   type: `${typeof controllerName}:stateChange`;
@@ -80,7 +101,7 @@ export class SubjectMetadataController extends BaseController<
 > {
   private subjectCacheLimit: number;
 
-  private subjectsWithoutPermissionsEcounteredSinceStartup: Set<string>;
+  private subjectsWithoutPermissionsEncounteredSinceStartup: Set<string>;
 
   private subjectHasPermissions: GenericPermissionController['hasPermissions'];
 
@@ -110,15 +131,20 @@ export class SubjectMetadataController extends BaseController<
 
     this.subjectHasPermissions = hasPermissions;
     this.subjectCacheLimit = subjectCacheLimit;
-    this.subjectsWithoutPermissionsEcounteredSinceStartup = new Set();
+    this.subjectsWithoutPermissionsEncounteredSinceStartup = new Set();
+
+    this.messagingSystem.registerActionHandler(
+      `${this.name}:getSubjectMetadata`,
+      this.getSubjectMetadata.bind(this),
+    );
   }
 
   /**
    * Clears the state of this controller. Also resets the cache of subjects
    * encountered since startup, so as to not prematurely reach the cache limit.
    */
   clearState(): void {
-    this.subjectsWithoutPermissionsEcounteredSinceStartup.clear();
+    this.subjectsWithoutPermissionsEncounteredSinceStartup.clear();
     this.update((_draftState) => {
       return { ...defaultState };
     });
@@ -143,20 +169,22 @@ export class SubjectMetadataController extends BaseController<
       extensionId: metadata.extensionId || null,
       iconUrl: metadata.iconUrl || null,
       name: metadata.name || null,
+      subjectType: metadata.subjectType || null,
     };
 
     let originToForget: string | null = null;
     // We only delete the oldest encountered subject from the cache, again to
     // ensure that the user's experience isn't degraded by missing icons etc.
     if (
-      this.subjectsWithoutPermissionsEcounteredSinceStartup.size >=
+      this.subjectsWithoutPermissionsEncounteredSinceStartup.size >=
       this.subjectCacheLimit
     ) {
-      const cachedOrigin = this.subjectsWithoutPermissionsEcounteredSinceStartup
-        .values()
-        .next().value;
+      const cachedOrigin =
+        this.subjectsWithoutPermissionsEncounteredSinceStartup
+          .values()
+          .next().value;
 
-      this.subjectsWithoutPermissionsEcounteredSinceStartup.delete(
+      this.subjectsWithoutPermissionsEncounteredSinceStartup.delete(
         cachedOrigin,
       );
 
@@ -165,7 +193,7 @@ export class SubjectMetadataController extends BaseController<
       }
     }
 
-    this.subjectsWithoutPermissionsEcounteredSinceStartup.add(origin);
+    this.subjectsWithoutPermissionsEncounteredSinceStartup.add(origin);
 
     this.update((draftState) => {
       // Typecast: ts(2589)
@@ -176,6 +204,16 @@ export class SubjectMetadataController extends BaseController<
     });
   }
 
+  /**
+   * Gets the subject metadata for the given origin, if any.
+   *
+   * @param origin - The origin for which to get the subject metadata.
+   * @returns The subject metadata, if any, or `undefined` otherwise.
+   */
+  getSubjectMetadata(origin: SubjectOrigin): SubjectMetadata | undefined {
+    return this.state.subjectMetadata[origin];
+  }
+
   /**
    * Deletes all subjects without permissions from the controller's state.
    */