Skip to content

Commit

Permalink
Request site records late (#46)
Browse files Browse the repository at this point in the history
* Request site records late
* add tests
  • Loading branch information
grod220 authored Jun 20, 2024
1 parent 237ee61 commit 93cc701
Show file tree
Hide file tree
Showing 3 changed files with 265 additions and 19 deletions.
5 changes: 5 additions & 0 deletions .changeset/brave-baboons-move.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'chrome-extension': patch
---

Update how connection approvals update origin records
225 changes: 225 additions & 0 deletions apps/extension/src/approve-origin.test.ts
Original file line number Diff line number Diff line change
@@ -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,
]);
});
});
});
54 changes: 35 additions & 19 deletions apps/extension/src/approve-origin.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -11,6 +13,33 @@ 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');
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,
Expand All @@ -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
Expand All @@ -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;
Expand Down

0 comments on commit 93cc701

Please sign in to comment.