From ff07c7c6f1959d274d01454fb3dad3173dfe3f35 Mon Sep 17 00:00:00 2001
From: Gabe Rodriguez <grod220@gmail.com>
Date: Thu, 20 Jun 2024 13:35:38 +0200
Subject: [PATCH] Request site records late

---
 apps/extension/src/approve-origin.test.ts | 55 ++++++++++++++++++++++
 apps/extension/src/approve-origin.ts      | 56 +++++++++++++++--------
 2 files changed, 92 insertions(+), 19 deletions(-)
 create mode 100644 apps/extension/src/approve-origin.test.ts

diff --git a/apps/extension/src/approve-origin.test.ts b/apps/extension/src/approve-origin.test.ts
new file mode 100644
index 00000000..e3bbebb4
--- /dev/null
+++ b/apps/extension/src/approve-origin.test.ts
@@ -0,0 +1,55 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { approveOrigin } from './approve-origin';
+import { localExtStorage } from './storage/local';
+
+// Mocks
+vi.mock('./storage/local', () => ({
+  get: vi.fn(),
+  set: vi.fn(),
+}));
+vi.mock('./popup', () => ({
+  popup: vi.fn(),
+}));
+
+const mockTab = {
+  index: 2,
+  pinned: false,
+  highlighted: true,
+  windowId: 1,
+  active: true,
+  id: 123456,
+  incognito: false,
+  selected: false,
+  discarded: false,
+  autoDiscardable: true,
+  groupId: -1,
+};
+
+describe('originHandlers', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe('approveOrigin', () => {
+    it('throws an error if the sender origin is not supported', async () => {
+      const messageSender = { origin: 'http://insecure.com' };
+      await expect(approveOrigin(messageSender)).rejects.toThrow('Unsupported sender');
+    });
+
+    it('returns the previously approved choice if one exists', async () => {
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-call
+      localExtStorage.mockReturnValue(
+        Promise.resolve([{ origin: 'https://example.com', choice: 'Approved' }]),
+      );
+
+      const messageSender = {
+        origin: 'https://example.com',
+        tab: mockTab,
+      };
+      const choice = await approveOrigin(messageSender);
+      expect(choice).toBe('Approved');
+    });
+
+    // More tests based on different scenarios...
+  });
+});
diff --git a/apps/extension/src/approve-origin.ts b/apps/extension/src/approve-origin.ts
index 3795d3a4..49fa6ee0 100644
--- a/apps/extension/src/approve-origin.ts
+++ b/apps/extension/src/approve-origin.ts
@@ -2,6 +2,8 @@ import { localExtStorage } from './storage/local';
 import { OriginApproval, PopupType } from './message/popup';
 import { popup } from './popup';
 import { UserChoice } from '@penumbra-zone/types/user-choice';
+import { produce } from 'immer';
+import { OriginRecord } from './storage/types';
 
 export const originAlreadyApproved = async (url: string): Promise<boolean> => {
   // parses the origin and returns a consistent format
@@ -11,6 +13,35 @@ export const originAlreadyApproved = async (url: string): Promise<boolean> => {
   return existingRecord?.choice === UserChoice.Approved;
 };
 
+const getChoiceForOrigin = async (origin: string) => {
+  const knownSites = await localExtStorage.get('knownSites');
+  const existingRecords = knownSites.filter(record => record.origin === origin);
+  if (!existingRecords.length) {
+    return undefined;
+  } else if (existingRecords.length === 1) {
+    return existingRecords[0];
+  } else {
+    // TODO: It's likely that an array is not the best data structure for this in storage. Should revisit later.
+    throw new Error(`There are multiple records for origin: ${origin}`);
+  }
+};
+
+const addOrUpdateSiteRecord = async (proposal: OriginRecord) => {
+  const knownSites = await localExtStorage.get('knownSites');
+  console.log('before', knownSites);
+  const newKnownSites = produce(knownSites, allRecords => {
+    const match = allRecords.find(r => r.origin === proposal.origin);
+    if (!match) {
+      allRecords.push(proposal);
+    } else {
+      match.choice = proposal.choice;
+      match.date = proposal.date;
+    }
+  });
+  console.log('after', newKnownSites);
+  await localExtStorage.set('knownSites', newKnownSites);
+};
+
 export const approveOrigin = async ({
   origin: senderOrigin,
   tab,
@@ -21,19 +52,11 @@ export const approveOrigin = async ({
 
   // parses the origin and returns a consistent format
   const urlOrigin = new URL(senderOrigin).origin;
-  const knownSites = await localExtStorage.get('knownSites');
-
-  const siteRecords = Map.groupBy(knownSites, site => site.origin === urlOrigin);
-  const irrelevant = siteRecords.get(false) ?? []; // we need to handle these in order to write back to storage
-  const [existingRecord, ...extraRecords] = siteRecords.get(true) ?? [];
-
-  if (extraRecords.length) throw new Error('Multiple records for the same origin');
-
-  const choice = existingRecord?.choice;
+  const record = await getChoiceForOrigin(urlOrigin);
 
   // Choice already made
-  if (choice === UserChoice.Approved || choice === UserChoice.Ignored) {
-    return choice;
+  if (record && (record.choice === UserChoice.Approved || record.choice === UserChoice.Ignored)) {
+    return record.choice;
   }
 
   // It's the first or repeat ask
@@ -43,18 +66,13 @@ export const approveOrigin = async ({
       origin: urlOrigin,
       favIconUrl: tab.favIconUrl,
       title: tab.title,
-      lastRequest: existingRecord?.date,
+      lastRequest: record?.date,
     },
   });
 
+  // if user interacted with popup, update record
   if (popupResponse) {
-    void localExtStorage.set(
-      // user interacted with popup, update record
-      // TODO: is there a race condition here?  if this object has been
-      //       written after our initial read, we'll clobber them
-      'knownSites',
-      [popupResponse, ...irrelevant],
-    );
+    void addOrUpdateSiteRecord(popupResponse);
   }
 
   return popupResponse?.choice ?? UserChoice.Denied;