diff --git a/.changeset/brave-baboons-move.md b/.changeset/brave-baboons-move.md new file mode 100644 index 00000000..5d87b23e --- /dev/null +++ b/.changeset/brave-baboons-move.md @@ -0,0 +1,5 @@ +--- +'chrome-extension': patch +--- + +Update how connection approvals update origin records diff --git a/apps/extension/src/approve-origin.test.ts b/apps/extension/src/approve-origin.test.ts new file mode 100644 index 00000000..2d749f33 --- /dev/null +++ b/apps/extension/src/approve-origin.test.ts @@ -0,0 +1,225 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { approveOrigin } from './approve-origin'; +import { UserChoice } from '@penumbra-zone/types/user-choice'; +import { OriginRecord } from './storage/types'; +import { PopupType } from './message/popup'; + +const mockLocalStorage = vi.hoisted(() => ({ + get: vi.fn(), + set: vi.fn(), +})); +vi.mock('./storage/local', () => ({ localExtStorage: mockLocalStorage })); + +const mockPopup = vi.hoisted(() => vi.fn()); +vi.mock('./popup', () => ({ + popup: mockPopup, +})); + +const mockTab = { + index: 2, + pinned: false, + highlighted: true, + windowId: 1, + active: true, + id: 123456, + incognito: false, + selected: false, + discarded: false, + autoDiscardable: true, + groupId: -1, + favIconUrl: 'https://image.com/favicon.ico', + title: 'Penumbra | xyz', +} satisfies chrome.tabs.Tab; + +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('throws an error if the tab ID is missing', async () => { + const messageSender = { origin: 'https://example.com' }; + await expect(approveOrigin(messageSender)).rejects.toThrow('Unsupported sender'); + }); + + it('throws an error if a frame ID is present', async () => { + const messageSender = { origin: 'https://example.com', tab: mockTab, frameId: 123 }; + await expect(approveOrigin(messageSender)).rejects.toThrow('Unsupported sender'); + }); + + it('prompts user for approval if the origin is not known', async () => { + mockLocalStorage.get.mockReturnValue(Promise.resolve([])); + const messageSender = { origin: 'https://newsite.com', tab: mockTab }; + mockPopup.mockResolvedValue({ + choice: UserChoice.Approved, + date: 123, + origin: 'https://newsite.com', + } satisfies OriginRecord); + + const choice = await approveOrigin(messageSender); + expect(choice).toBe(UserChoice.Approved); + }); + + it('returns the stored choice if the origin is already known and approved', async () => { + mockLocalStorage.get.mockReturnValue( + Promise.resolve([{ origin: 'https://example.com', choice: UserChoice.Approved }]), + ); + + const messageSender = { origin: 'https://example.com', tab: mockTab }; + const choice = await approveOrigin(messageSender); + expect(choice).toBe(UserChoice.Approved); + }); + + it('throws an error if multiple records exist for the same origin', async () => { + mockLocalStorage.get.mockReturnValue( + Promise.resolve([ + { origin: 'https://example.com', choice: UserChoice.Approved }, + { origin: 'https://example.com', choice: UserChoice.Denied }, + ]), + ); + + const messageSender = { origin: 'https://example.com', tab: mockTab }; + await expect(approveOrigin(messageSender)).rejects.toThrow( + 'There are multiple records for origin: https://example.com', + ); + }); + + it('returns Denied if the user denies the request', async () => { + mockLocalStorage.get.mockReturnValue(Promise.resolve([])); + const messageSender = { origin: 'https://newsite.com', tab: mockTab }; + mockPopup.mockResolvedValue({ + choice: UserChoice.Denied, + date: 123, + origin: 'https://newsite.com', + } satisfies OriginRecord); + + const choice = await approveOrigin(messageSender); + expect(choice).toBe(UserChoice.Denied); + }); + + it('updates an existing record if the user changes their choice from denied to approve', async () => { + mockLocalStorage.get.mockReturnValue( + Promise.resolve([{ origin: 'https://example.com', choice: UserChoice.Denied }]), + ); + const messageSender = { origin: 'https://example.com', tab: mockTab }; + mockPopup.mockResolvedValue({ + choice: UserChoice.Approved, + date: 123, + origin: 'https://example.com', + } satisfies OriginRecord); + + const choice = await approveOrigin(messageSender); + expect(choice).toBe(UserChoice.Approved); + }); + + it('returns the previously approved choice if one exists', async () => { + mockLocalStorage.get.mockReturnValue( + Promise.resolve([{ origin: 'https://example.com', choice: UserChoice.Approved }]), + ); + + const messageSender = { + origin: 'https://example.com', + tab: mockTab, + }; + const choice = await approveOrigin(messageSender); + expect(choice).toBe(UserChoice.Approved); + }); + + it('correctly updates the persisted known sites after user interaction', async () => { + mockLocalStorage.get.mockReturnValue( + Promise.resolve([{ origin: 'https://example.com', choice: UserChoice.Denied }]), + ); + const messageSender = { origin: 'https://example.com', tab: mockTab }; + const newOriginRecord = { + choice: UserChoice.Approved, + date: 123, + origin: 'https://example.com', + } satisfies OriginRecord; + mockPopup.mockResolvedValue(newOriginRecord); + + await approveOrigin(messageSender); + + expect(mockLocalStorage.set).toHaveBeenCalledWith('knownSites', [newOriginRecord]); + }); + + it('returns the stored choice if the origin is already known and ignored', async () => { + mockLocalStorage.get.mockReturnValue( + Promise.resolve([{ origin: 'https://example.com', choice: UserChoice.Ignored }]), + ); + + const messageSender = { origin: 'https://example.com', tab: mockTab }; + const choice = await approveOrigin(messageSender); + expect(choice).toBe(UserChoice.Ignored); + }); + + it('returns UserChoice.Denied if the user closes the popup without making a choice', async () => { + mockLocalStorage.get.mockReturnValue(Promise.resolve([])); + const messageSender = { origin: 'https://newsite.com', tab: mockTab }; + mockPopup.mockResolvedValue(undefined); + + const choice = await approveOrigin(messageSender); + expect(choice).toBe(UserChoice.Denied); + }); + + it('correctly handles trailing slashes in the origin', async () => { + mockLocalStorage.get.mockReturnValue( + Promise.resolve([{ origin: 'https://example.com', choice: UserChoice.Approved }]), + ); + + const messageSender = { origin: 'https://example.com/', tab: mockTab }; + const choice = await approveOrigin(messageSender); + expect(choice).toBe(UserChoice.Approved); + }); + + it('shows the popup with the correct parameters', async () => { + mockLocalStorage.get.mockReturnValue(Promise.resolve([])); + const messageSender = { origin: 'https://newsite.com', tab: mockTab }; + mockPopup.mockResolvedValue({ + choice: UserChoice.Approved, + date: 123, + origin: 'https://newsite.com', + } satisfies OriginRecord); + + await approveOrigin(messageSender); + + expect(mockPopup).toHaveBeenCalledWith({ + type: PopupType.OriginApproval, + request: { + origin: 'https://newsite.com', + favIconUrl: mockTab.favIconUrl, + title: mockTab.title, + lastRequest: undefined, + }, + }); + }); + + it('correctly persists the known sites when a new site is added', async () => { + const existingOriginRecord = { + choice: UserChoice.Approved, + date: 456, + origin: 'https://existingsite.com', + } satisfies OriginRecord; + mockLocalStorage.get.mockReturnValue(Promise.resolve([existingOriginRecord])); + + const messageSender = { origin: 'https://newsite.com', tab: mockTab }; + const newOriginRecord = { + choice: UserChoice.Approved, + date: 123, + origin: 'https://newsite.com', + } satisfies OriginRecord; + mockPopup.mockResolvedValue(newOriginRecord); + + await approveOrigin(messageSender); + + expect(mockLocalStorage.set).toHaveBeenCalledWith('knownSites', [ + existingOriginRecord, + newOriginRecord, + ]); + }); + }); +}); diff --git a/apps/extension/src/approve-origin.ts b/apps/extension/src/approve-origin.ts index 3795d3a4..add30a32 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 => { // parses the origin and returns a consistent format @@ -11,6 +13,33 @@ export const originAlreadyApproved = async (url: string): Promise => { 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'); + 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; + } + }); + await localExtStorage.set('knownSites', newKnownSites); +}; + export const approveOrigin = async ({ origin: senderOrigin, tab, @@ -21,19 +50,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 +64,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], - ); + await addOrUpdateSiteRecord(popupResponse); } return popupResponse?.choice ?? UserChoice.Denied;