Skip to content

Commit

Permalink
KEEP-397 Remove messages when tab closes (#413)
Browse files Browse the repository at this point in the history
  • Loading branch information
Sergey Slipchenko authored Sep 15, 2022
1 parent 2d9ded4 commit d185827
Show file tree
Hide file tree
Showing 8 changed files with 345 additions and 158 deletions.
90 changes: 57 additions & 33 deletions src/background.ts
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,6 @@ import {
MessageStoreItem,
} from 'messages/types';
import { CreateWalletInput } from 'wallets/types';
import pump from 'pump';
import { collectBalances } from 'balances/utils';
import { BalancesItem } from 'balances/types';

Expand Down Expand Up @@ -101,35 +100,28 @@ Sentry.init({

extension.runtime.onConnect.addListener(async remotePort => {
const bgService = await bgPromise;
const portStream = new PortStream(remotePort);

if (remotePort.name === 'contentscript') {
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
const origin = new URL(remotePort.sender!.url!).hostname;
bgService.setupPageConnection(portStream, origin);
bgService.setupPageConnection(remotePort);
} else {
bgService.setupUiConnection(portStream);
bgService.setupUiConnection(remotePort);
}
});

extension.runtime.onConnectExternal.addListener(async remotePort => {
const bgService = await bgPromise;
const portStream = new PortStream(remotePort);
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
const origin = new URL(remotePort.sender!.url!).hostname;
bgService.setupPageConnection(portStream, origin);
bgService.setupPageConnection(remotePort);
});

extension.runtime.onUpdateAvailable.addListener(async () => {
await backupStorage();
await extension.runtime.reload();
extension.runtime.reload();
});

extension.runtime.onInstalled.addListener(async details => {
const bgService = await bgPromise;

if (details.reason === extension.runtime.OnInstalledReason.UPDATE) {
bgService.messageController.clearUnusedMessages();
bgService.assetInfoController.addTickersForExistingAssets();
bgService.vaultController.migrate();
bgService.addressBookController.migrate();
Expand All @@ -153,9 +145,9 @@ async function setupBackgroundService() {
(global as any).background = backgroundService;
}

const updateBadge = async () => {
const { selectedAccount } =
backgroundService.extensionStorage.getState('selectedAccount');
const updateBadge = () => {
const selectedAccount =
backgroundService.preferencesController.getSelectedAccount();
const messages = backgroundService.messageController.getUnapproved();
const notifications =
backgroundService.notificationsController.getGroupNotificationsByAccount(
Expand All @@ -172,7 +164,7 @@ async function setupBackgroundService() {
// update badge
backgroundService.messageController.on('Update badge', updateBadge);
backgroundService.notificationsController.on('Update badge', updateBadge);
await updateBadge();
updateBadge();
// open new tab
backgroundService.messageController.on('Open new tab', url => {
extension.tabs.create({ url });
Expand Down Expand Up @@ -200,6 +192,9 @@ async function setupBackgroundService() {
return tabsManager.closeCurrentTab();
});

backgroundService.messageController.clearMessages();
windowManager.closeWindow();

return backgroundService;
}

Expand Down Expand Up @@ -759,7 +754,7 @@ class BackgroundService extends EventEmitter {
};
}

async validatePermission(origin: string) {
async validatePermission(origin: string, connectionId: string) {
const { selectedAccount } = this.getState('selectedAccount');

if (!selectedAccount) throw ERRORS.EMPTY_KEEPER();
Expand Down Expand Up @@ -791,6 +786,7 @@ class BackgroundService extends EventEmitter {
if (!messageId) {
const messageData: MessageInput = {
origin,
connectionId,
title: null,
options: {},
broadcast: false,
Expand Down Expand Up @@ -823,7 +819,7 @@ class BackgroundService extends EventEmitter {
}
}

getNewMessageFn(origin?: string): NewMessageFn {
getNewMessageFn(origin?: string, connectionId?: string): NewMessageFn {
return async (data, type, options, broadcast, title = '') => {
if (data.type === 1000) {
type = 'auth';
Expand All @@ -832,13 +828,14 @@ class BackgroundService extends EventEmitter {
data.isRequest = true;
}

if (origin) {
await this.validatePermission(origin);
if (origin != null && connectionId != null) {
await this.validatePermission(origin, connectionId);
}

const { selectedAccount } = this.getState('selectedAccount');

const { noSign, ...result } = await this.messageController.newMessage({
connectionId,
data,
type,
title,
Expand Down Expand Up @@ -871,8 +868,8 @@ class BackgroundService extends EventEmitter {
};
}

getInpageApi(origin: string) {
const newMessage = this.getNewMessageFn(origin);
getInpageApi(origin: string, connectionId: string) {
const newMessage = this.getNewMessageFn(origin, connectionId);

const newNotification = async (
data:
Expand Down Expand Up @@ -996,7 +993,7 @@ class BackgroundService extends EventEmitter {
throw !initialized ? ERRORS.INIT_KEEPER() : ERRORS.EMPTY_KEEPER();
}

await this.validatePermission(origin);
await this.validatePermission(origin, connectionId);

return await newNotification(data);
},
Expand All @@ -1011,7 +1008,7 @@ class BackgroundService extends EventEmitter {
throw !initialized ? ERRORS.INIT_KEEPER() : ERRORS.EMPTY_KEEPER();
}

await this.validatePermission(origin);
await this.validatePermission(origin, connectionId);

return this._publicState(origin);
},
Expand Down Expand Up @@ -1048,7 +1045,7 @@ class BackgroundService extends EventEmitter {
throw ERRORS.INVALID_FORMAT(undefined, 'publicKey is invalid');
}

await this.validatePermission(origin);
await this.validatePermission(origin, connectionId);

return this.walletController.getKEK(
selectedAccount.address,
Expand Down Expand Up @@ -1080,7 +1077,7 @@ class BackgroundService extends EventEmitter {
throw ERRORS.INVALID_FORMAT(undefined, 'publicKey is invalid');
}

await this.validatePermission(origin);
await this.validatePermission(origin, connectionId);

return this.walletController.encryptMessage(
selectedAccount.address,
Expand Down Expand Up @@ -1113,7 +1110,7 @@ class BackgroundService extends EventEmitter {
throw ERRORS.INVALID_FORMAT(undefined, 'publicKey is invalid');
}

await this.validatePermission(origin);
await this.validatePermission(origin, connectionId);

return this.walletController.decryptMessage(
selectedAccount.address,
Expand All @@ -1126,9 +1123,8 @@ class BackgroundService extends EventEmitter {
};
}

setupUiConnection(connectionStream: pump.Stream) {
const api = this.getApi();
const dnode = setupDnode(connectionStream, api, 'api');
setupUiConnection(remotePort: chrome.runtime.Port) {
const dnode = setupDnode(new PortStream(remotePort), this.getApi(), 'api');

// eslint-disable-next-line @typescript-eslint/no-explicit-any
const remoteHandler = (remote: any) => {
Expand All @@ -1149,9 +1145,37 @@ class BackgroundService extends EventEmitter {
dnode.on('remote', remoteHandler);
}

setupPageConnection(connectionStream: pump.Stream, origin: string) {
const inpageApi = this.getInpageApi(origin);
setupDnode(connectionStream, inpageApi, 'inpageApi');
setupPageConnection(remotePort: chrome.runtime.Port) {
const { sender } = remotePort;

if (!sender || !sender.url) {
return;
}

const origin = new URL(sender.url).hostname;
const connectionId = uuidv4();
const inpageApi = this.getInpageApi(origin, connectionId);

const dnode = setupDnode(
new PortStream(remotePort),
inpageApi,
'inpageApi'
);

dnode.once('end', () => {
this.messageController.removeMessagesFromConnection(connectionId);

const messages = this.messageController.getUnapproved();

const notifications =
this.notificationsController.getGroupNotificationsByAccount(
this.preferencesController.getSelectedAccount()
);

if (messages.length === 0 && notifications.length === 0) {
this.emit('Close notification');
}
});
}

_getCurrentNetwork(account: PreferencesAccount | undefined) {
Expand Down
41 changes: 15 additions & 26 deletions src/controllers/message.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import ObservableStore from 'obs-store';
import { extension } from 'lib/extension';
import { MsgStatus, MSG_STATUSES } from '../constants';
import { MSG_STATUSES } from '../constants';
import { v4 as uuidv4 } from 'uuid';
import log from 'loglevel';
import EventEmitter from 'events';
Expand Down Expand Up @@ -344,6 +344,20 @@ export class MessageController extends EventEmitter {
this._updateMessagesByTimeout();
}

removeMessagesFromConnection(connectionId: string) {
const { messages } = this.store.getState();

messages.forEach(message => {
if (message.connectionId === connectionId) {
this.reject(message.id);
}
});

this._updateStore(
messages.filter(message => message.connectionId !== connectionId)
);
}

clearMessages(ids?: string | string[]) {
if (typeof ids === 'string') {
this._deleteMessage(ids);
Expand All @@ -354,31 +368,6 @@ export class MessageController extends EventEmitter {
}
}

clearUnusedMessages() {
const unusedStatuses: MsgStatus[] = [
MSG_STATUSES.REJECTED,
MSG_STATUSES.REJECTED_FOREVER,
MSG_STATUSES.SIGNED,
MSG_STATUSES.PUBLISHED,
MSG_STATUSES.FAILED,
];
const unusedMessages: MessageStoreItem[] = [];
const actualMessages: MessageStoreItem[] = [];

this.messages.messages.forEach(message => {
const { status } = message;
if (unusedStatuses.indexOf(status) === -1) {
actualMessages.push(message);
} else {
unusedMessages.push(message);
}
});

if (unusedMessages.length) {
this._updateStore(actualMessages);
}
}

getUnapproved() {
return this.messages.messages.filter(
({ status }) => status === MSG_STATUSES.UNAPPROVED
Expand Down
2 changes: 2 additions & 0 deletions src/messages/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import {
} from '@waves/waves-transactions/dist/requests/custom-data';

export type MessageInput = {
connectionId?: string;
account: PreferencesAccount;
broadcast?: boolean;
options?: {
Expand Down Expand Up @@ -165,6 +166,7 @@ interface TxData {
}

export type MessageStoreItem = {
connectionId?: string;
account: PreferencesAccount;
broadcast?: boolean;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
Expand Down
2 changes: 0 additions & 2 deletions src/ui/actions/updateState.ts
Original file line number Diff line number Diff line change
Expand Up @@ -288,8 +288,6 @@ export function createUpdateState(store: UiStore) {
}

const nfts = getParam(state.nfts, null);
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
if (nfts && !equals(nfts, currentState.nfts)) {
store.dispatch({
type: ACTION.UPDATE_NFTS,
Expand Down
16 changes: 9 additions & 7 deletions src/ui/reducers/notifications.ts
Original file line number Diff line number Diff line change
Expand Up @@ -76,21 +76,23 @@ const getActiveFromState = (
): ActivePopupState | null => {
// Can activeMessage
if (state != null && (state.msg || state.notify)) {
let { msg, notify } = state;

// Update from messages
if (state.msg) {
if (msg) {
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
const msgItem = allMessages.find(item => item.id === state.msg!.id);
state.msg = msgItem || state.msg;
const msgItem = allMessages.find(item => item.id === msg!.id);
msg = msgItem || allMessages[0] || null;
}

// Update from notifications
if (state.notify) {
const { origin } = state.notify[0];
if (notify) {
const { origin } = notify[0];
const newItem = notifications.find(([item]) => item.origin === origin);
state.notify = newItem || state.notify;
notify = newItem || notifications[0] || null;
}

return state;
return { msg, notify };
}

// To msgList
Expand Down
Loading

0 comments on commit d185827

Please sign in to comment.