Re-inject site scripts when background process wakes up from idle (#40)

* Re-inject site scripts when background process wakes up from idle

* Make the onDOMContentLoaded even listener a global listener so it persists

* Inject GK buttons into current tabs on startup and on install

* Clear injectionDomains cache instead of reloading extension when permissions are granted

Author: jdgarcia

src/background.ts

@@ -1,11 +1,10 @@
-import type { WebNavigation } from 'webextension-polyfill';
-import { runtime, scripting, tabs, webNavigation } from 'webextension-polyfill';
+import { runtime, scripting, storage, tabs, webNavigation } from 'webextension-polyfill';
 import { fetchUser } from './gkApi';
 import { injectionScope as inject_azureDevops } from './hosts/azureDevops';
 import { injectionScope as inject_bitbucket } from './hosts/bitbucket';
 import { injectionScope as inject_github } from './hosts/github';
 import { injectionScope as inject_gitlab } from './hosts/gitlab';
-import { refreshPermissions } from './permissions-helper';
+import { domainToMatchPattern, refreshPermissions } from './permissions-helper';
 import { getEnterpriseConnections, GKDotDevUrl, PermissionsGrantedMessage, PopupInitMessage } from './shared';
 import type { CacheContext } from './types';
 
@@ -23,6 +22,19 @@ const DefaultInjectionDomains: InjectionDomains = {
 	azureDevops: ['dev.azure.com'],
 };
 
+webNavigation.onDOMContentLoaded.addListener(async details => {
+	const injectionDomains = await getInjectionDomains();
+
+	const injectionFn = getInjectionFn(details.url, injectionDomains);
+	if (injectionFn) {
+		void scripting.executeScript({
+			target: { tabId: details.tabId },
+			func: injectionFn,
+			args: [details.url, GKDotDevUrl],
+		});
+	}
+});
+
 webNavigation.onHistoryStateUpdated.addListener(details => {
 	// used to detect when the user navigates to a different page in the same tab
 	const url = new URL(details.url);
@@ -39,14 +51,55 @@ runtime.onMessage.addListener(async msg => {
 		const context: CacheContext = {};
 		return refreshPermissions(context);
 	} else if (msg === PermissionsGrantedMessage) {
-		// Reload extension to update injection listener
-		runtime.reload();
+		await storage.session.remove('injectionDomains');
 		return undefined;
 	}
 	console.error('Recevied unknown runtime message', msg);
 	return undefined;
 });
 
+runtime.onInstalled.addListener(injectIntoCurrentTabs);
+runtime.onStartup.addListener(injectIntoCurrentTabs);
+
+async function injectIntoCurrentTabs() {
+	const injectionDomains = await getInjectionDomains();
+	const allDomains = Object.values<string[]>(injectionDomains as any).flat();
+
+	const currentTabs = await tabs.query({
+		url: allDomains.map(domainToMatchPattern),
+		status: 'complete',
+		discarded: false,
+	});
+	currentTabs.forEach(tab => {
+		if (tab.id && tab.url) {
+			const injectionFn = getInjectionFn(tab.url, injectionDomains);
+			if (injectionFn) {
+				void scripting.executeScript({
+					target: { tabId: tab.id },
+					func: injectionFn,
+					args: [tab.url, GKDotDevUrl],
+				});
+			}
+		}
+	});
+}
+
+async function getInjectionDomains() {
+	let { injectionDomains } = (await storage.session.get('injectionDomains')) as {
+		injectionDomains?: InjectionDomains;
+	};
+	if (!injectionDomains) {
+		const context: CacheContext = {};
+		// This removes unneded permissions
+		await refreshPermissions(context);
+
+		injectionDomains = await computeInjectionDomains(context);
+		await storage.session.set({ injectionDomains: injectionDomains });
+	}
+
+	return injectionDomains;
+}
+
 async function computeInjectionDomains(context: CacheContext) {
 	const injectionDomains = structuredClone(DefaultInjectionDomains);
 	const enterpriseConnections = await getEnterpriseConnections(context);
@@ -63,33 +116,14 @@ async function computeInjectionDomains(context: CacheContext) {
 	return injectionDomains;
 }
 
-async function addInjectionListener(context: CacheContext) {
-	const injectionDomains = await computeInjectionDomains(context);
-	const allDomains = Object.values<string[]>(injectionDomains as any).flat();
-
-	// note: This is a closure over injectionDomains
-	const injectScript = (details: WebNavigation.OnDOMContentLoadedDetailsType) => {
-		void scripting.executeScript({
-			target: { tabId: details.tabId },
-			// injectImmediately: true,
-			func: getInjectionFn(details.url, injectionDomains),
-			args: [details.url, GKDotDevUrl],
-		});
-	};
-
-	webNavigation.onDOMContentLoaded.addListener(injectScript, {
-		url: allDomains.map(domain => ({ hostContains: domain })),
-	});
-}
-
 function urlHostHasDomain(url: URL, domains: string[]): boolean {
 	return domains.some(domain => url.hostname.endsWith(domain));
 }
 
 function getInjectionFn(
 	rawUrl: string,
 	injectionDomains: InjectionDomains,
-): (url: string, gkDotDevUrl: string) => void {
+): ((url: string, gkDotDevUrl: string) => void) | null {
 	const url = new URL(rawUrl);
 	if (urlHostHasDomain(url, injectionDomains.github)) {
 		return inject_github;
@@ -107,20 +141,12 @@ function getInjectionFn(
 		return inject_azureDevops;
 	}
 
-	console.error('Unsupported host');
-	throw new Error('Unsupported host');
+	return null;
 }
 
 async function main() {
 	// The fetchUser function also updates the extension icon if the user is logged in
 	await fetchUser();
-
-	const context: CacheContext = {};
-	// This removes unneded permissions
-	await refreshPermissions(context);
-	// NOTE: This may request hosts that we may not have permissions for, which will log errors for the extension
-	// This does not cause any issues, and eliminating the errors requires more logic
-	await addInjectionListener(context);
 }
 
 void main();

src/permissions-helper.ts

@@ -3,7 +3,7 @@ import { permissions } from 'webextension-polyfill';
 import { arrayDifference, CloudProviders, getEnterpriseConnections } from './shared';
 import type { CacheContext } from './types';
 
-function domainToMatchPattern(domain: string): string {
+export function domainToMatchPattern(domain: string): string {
 	return `*://*.${domain}/*`;
 }
 

src/popup/components/RequestPermissionsBanner.tsx

@@ -33,6 +33,7 @@ export const RequestPermissionsBanner = ({ permissionsRequest }: { permissionsRe
 					const granted = await permissions.request(permissionsRequest.request);
 					if (granted) {
 						await sendPermissionsGranted();
+						window.close();
 					}
 				}}
 			>