Merge pull request #3 from Holo-Host/feat/rewrite-password-to-memory

(WIP) Feat/rewrite password to memory
Holo-Host · May 19, 2024 · 2b9972d · 2b9972d
2 parents 23a47b4 + 1c225e1
commit 2b9972d
Show file tree

Hide file tree

Showing 27 changed files with 634 additions and 548 deletions.
diff --git a/holo-key-manager-extension/scripts/background.ts b/holo-key-manager-extension/scripts/background.ts
@@ -1,24 +1,32 @@
 import {
+	APP_NOT_AUTHENTICATED,
 	BACKGROUND_SCRIPT_RECEIVED_DATA,
+	EXTENSION_NOT_AUTHENTICATED,
+	EXTENSION_SESSION_INFO,
 	GENERIC_ERROR,
+	GET_EXTENSION_SESSION,
 	NEEDS_SETUP,
 	NO_KEY_FOR_HAPP,
-	NOT_AUTHENTICATED,
 	SENDER_BACKGROUND_SCRIPT,
 	SENDER_EXTENSION,
+	SENDER_WEBAPP,
+	SETUP_EXTENSION_SESSION,
 	SIGN_IN,
+	SIGN_IN_SUCCESS,
 	SIGN_MESSAGE,
 	SIGN_MESSAGE_SUCCESS,
 	SIGN_OUT,
 	SIGN_OUT_SUCCESS,
 	SIGN_UP,
+	SIGN_UP_SUCCESS,
 	UNKNOWN_ACTION
 } from '@shared/const';
 import { createQueryParams } from '@shared/helpers';
 import { isAppSignUpComplete, isAuthenticated, isSetupComplete, signOut } from '@shared/services';
 import {
 	type ActionPayload,
 	type Message,
+	type MessageWithId,
 	MessageWithIdSchema,
 	type WindowProperties
 } from '@shared/types';
@@ -31,6 +39,8 @@ import { signMessageLogic } from './helpers';
 
 let windowId: number | undefined;
 
+let session: string | undefined;
+
 type SendResponse = (response?: Message) => void;
 type SendResponseWithSender = (response: ActionPayload) => void;
 
@@ -39,8 +49,8 @@ const handleError = (sendResponse: SendResponseWithSender) => {
 	sendResponse({ action: GENERIC_ERROR });
 };
 
-const createWindowProperties = (actionPayload?: ActionPayload): WindowProperties => ({
-	url: `webapp-extension/setup.html${actionPayload ? `?${createQueryParams(actionPayload)}` : ''}`,
+const createWindowProperties = (parsedMessage?: MessageWithId): WindowProperties => ({
+	url: `webapp-extension/setup.html${parsedMessage ? `?${createQueryParams(parsedMessage)}` : ''}`,
 	type: 'popup',
 	height: 500,
 	width: 375,
@@ -81,31 +91,30 @@ const manageWindow = (
 
 const updateOrCreateWindowCommon = (
 	handleWindowUpdateOrCreate: () => Promise<void>,
-	actionPayload?: ActionPayload
-) => manageWindow(createWindowProperties(actionPayload), handleWindowUpdateOrCreate);
+	parsedMessage?: MessageWithId
+) => manageWindow(createWindowProperties(parsedMessage), handleWindowUpdateOrCreate);
 
 const updateOrCreateWindow = (
 	successAction: typeof NEEDS_SETUP,
 	sendResponse: SendResponseWithSender
 ) =>
 	updateOrCreateWindowCommon(async () => {
 		if (chrome.runtime.lastError) return handleError(sendResponse);
-
 		try {
 			sendResponse({ action: successAction });
 		} catch (error) {
 			handleError(sendResponse);
 		}
 	});
 
-const waitForFormSubmission = (): Promise<Message> =>
+const waitForFormSubmission = (id: string): Promise<Message> =>
 	new Promise((resolve) => {
 		const messageListener = (
-			message: Message,
+			message: MessageWithId,
 			sender: chrome.runtime.MessageSender,
 			sendResponse: SendResponse
 		) => {
-			if (message.sender !== SENDER_EXTENSION) return;
+			if (message.sender !== SENDER_EXTENSION || message.id !== id) return;
 			sendResponse({
 				action: BACKGROUND_SCRIPT_RECEIVED_DATA,
 				sender: SENDER_BACKGROUND_SCRIPT
@@ -118,54 +127,48 @@ const waitForFormSubmission = (): Promise<Message> =>
 
 const createOrUpdateDataResponseWindow = (
 	sendResponse: SendResponseWithSender,
-	actionPayload: ActionPayload
+	parsedMessage: MessageWithId
 ) =>
 	updateOrCreateWindowCommon(async () => {
 		if (chrome.runtime.lastError) return handleError(sendResponse);
 
 		try {
-			const message = await waitForFormSubmission();
+			const message = await waitForFormSubmission(parsedMessage.id);
 			sendResponse(message);
 		} catch (error) {
 			handleError(sendResponse);
 		}
-	}, actionPayload);
-
-const processMessage = async (message: Message, sendResponse: SendResponse) => {
-	const sendResponseWithSender = (response: ActionPayload) =>
-		sendResponse({ ...response, sender: SENDER_BACKGROUND_SCRIPT });
-	const parsedMessage = MessageWithIdSchema.safeParse(message);
-	if (!parsedMessage.success) return;
+	}, parsedMessage);
 
+const processMessageWebApp = async (
+	parsedMessage: MessageWithId,
+	sendResponseWithSender: SendResponseWithSender
+) => {
 	try {
 		if (!(await isSetupComplete())) {
 			return updateOrCreateWindow(NEEDS_SETUP, sendResponseWithSender);
 		}
-
-		switch (parsedMessage.data.action) {
+		switch (parsedMessage.action) {
 			case SIGN_UP:
-				return createOrUpdateDataResponseWindow(sendResponseWithSender, {
-					action: parsedMessage.data.action,
-					payload: parsedMessage.data.payload
-				});
+				return createOrUpdateDataResponseWindow(sendResponseWithSender, parsedMessage);
 			case SIGN_IN:
-				return (await isAppSignUpComplete(parsedMessage.data.payload.happId))
-					? createOrUpdateDataResponseWindow(sendResponseWithSender, {
-							action: parsedMessage.data.action,
-							payload: parsedMessage.data.payload
-						})
+				return (await isAppSignUpComplete(parsedMessage.payload.happId))
+					? createOrUpdateDataResponseWindow(sendResponseWithSender, parsedMessage)
 					: sendResponseWithSender({ action: NO_KEY_FOR_HAPP });
 			case SIGN_MESSAGE:
-				if (await isAuthenticated(parsedMessage.data.payload.happId)) {
-					const signedMessage = await signMessageLogic(parsedMessage.data.payload);
+				if (!session) {
+					return sendResponseWithSender({ action: EXTENSION_NOT_AUTHENTICATED });
+				}
+				if (await isAuthenticated(parsedMessage.payload.happId)) {
+					const signature = await signMessageLogic({ ...parsedMessage.payload, session });
 					return sendResponseWithSender({
 						action: SIGN_MESSAGE_SUCCESS,
-						payload: signedMessage
+						payload: signature
 					});
 				}
-				return sendResponseWithSender({ action: NOT_AUTHENTICATED });
+				return sendResponseWithSender({ action: APP_NOT_AUTHENTICATED });
 			case SIGN_OUT:
-				signOut(parsedMessage.data.payload.happId);
+				signOut(parsedMessage.payload.happId);
 				return sendResponseWithSender({ action: SIGN_OUT_SUCCESS });
 			default:
 				return sendResponseWithSender({ action: UNKNOWN_ACTION });
@@ -175,7 +178,54 @@ const processMessage = async (message: Message, sendResponse: SendResponse) => {
 	}
 };
 
+const processMessageExtension = async (
+	parsedMessage: MessageWithId,
+	sendResponseWithSender: SendResponseWithSender
+) => {
+	try {
+		switch (parsedMessage.action) {
+			case SETUP_EXTENSION_SESSION:
+				session = parsedMessage.payload;
+				return sendResponseWithSender({
+					action: BACKGROUND_SCRIPT_RECEIVED_DATA
+				});
+			case GET_EXTENSION_SESSION:
+				return sendResponseWithSender({
+					action: EXTENSION_SESSION_INFO,
+					payload: session
+				});
+			case SIGN_IN_SUCCESS:
+			case SIGN_UP_SUCCESS:
+				break;
+			default:
+				return sendResponseWithSender({ action: UNKNOWN_ACTION });
+		}
+	} catch (error) {
+		handleError(sendResponseWithSender);
+	}
+};
+
 chrome.runtime.onMessage.addListener((message: Message, sender, sendResponse: SendResponse) => {
-	processMessage(message, sendResponse);
+	const sendResponseWithSender = (response: ActionPayload) =>
+		sendResponse({ ...response, sender: SENDER_BACKGROUND_SCRIPT });
+
+	const parsedMessage = MessageWithIdSchema.safeParse(message);
+	if (!parsedMessage.success) return;
+
+	const processMessage = (sender: string) => {
+		switch (sender) {
+			case SENDER_WEBAPP:
+				return processMessageWebApp;
+			case SENDER_EXTENSION:
+				return processMessageExtension;
+			default:
+				return null;
+		}
+	};
+
+	const handler = processMessage(parsedMessage.data.sender);
+	if (handler) {
+		handler(parsedMessage.data, sendResponseWithSender);
+	}
 	return true;
 });
diff --git a/holo-key-manager-extension/scripts/helpers.ts b/holo-key-manager-extension/scripts/helpers.ts
@@ -1,15 +1,17 @@
 import { AUTHENTICATED_APPS_LIST, SESSION } from '@shared/const';
 import { base64ToUint8Array, uint8ArrayToBase64 } from '@shared/helpers';
-import { getSessionKey, storageService } from '@shared/services';
+import { getDeviceKey, storageService } from '@shared/services';
 import {
 	AuthenticatedAppsListSchema,
-	type MessageToSign,
+	type SignMessage,
 	SuccessMessageSignedSchema
 } from '@shared/types';
 // @ts-expect-error no types for hcSeedBundle
 import * as hcSeedBundle from 'hcSeedBundle';
 
-export const signMessageLogic = async ({ message, happId }: MessageToSign) => {
+export const signMessageLogic = async ({ message, happId, session }: SignMessage) => {
+	const encryptedDeviceKey = await getDeviceKey();
+
 	const authenticatedAppsListData = await storageService.getWithoutCallback({
 		key: AUTHENTICATED_APPS_LIST,
 		area: SESSION
@@ -18,28 +20,23 @@ export const signMessageLogic = async ({ message, happId }: MessageToSign) => {
 	const parsedAuthenticatedAppsListData =
 		AuthenticatedAppsListSchema.safeParse(authenticatedAppsListData);
 
-	if (!parsedAuthenticatedAppsListData.success) {
-		throw new Error('Failed to parse authenticated apps list data');
+	if (!parsedAuthenticatedAppsListData.success || !session) {
+		throw new Error('Authentication failed: Unable to parse apps list or session missing');
 	}
 
 	const index = parsedAuthenticatedAppsListData.data[happId];
-	const sessionKey = await getSessionKey();
-
-	if (!sessionKey.success) {
-		throw new Error('Session data not found');
-	}
 
 	await hcSeedBundle.seedBundleReady;
 
 	const cipherList = hcSeedBundle.UnlockedSeedBundle.fromLocked(
-		base64ToUint8Array(sessionKey.data)
+		base64ToUint8Array(encryptedDeviceKey)
 	);
 
 	if (!(cipherList[0] instanceof hcSeedBundle.LockedSeedCipherPwHash)) {
 		throw new Error('Expecting PwHash');
 	}
 
-	const pw = new TextEncoder().encode(SESSION);
+	const pw = new TextEncoder().encode(session);
 	const keyUnlocked = cipherList[0].unlock(hcSeedBundle.parseSecret(pw));
 
 	const appKey = keyUnlocked.derive(index);

diff --git a/holo-key-manager-extension/src/lib/helpers/navigation.ts b/holo-key-manager-extension/src/lib/helpers/navigation.ts
@@ -13,7 +13,8 @@ export const extractDetailsFromUrl = derived(page, ($page) => {
 		happUiUrl: '',
 		message: 'Unknown Message',
 		requireEmail: false,
-		requireRegistrationCode: false
+		requireRegistrationCode: false,
+		messageId: ''
 	};
 
 	const params = new URLSearchParams(new URL($page.url.href).search);

diff --git a/holo-key-manager-extension/src/lib/helpers/queries.ts b/holo-key-manager-extension/src/lib/helpers/queries.ts
@@ -4,12 +4,21 @@ import { unlockKey } from '$services';
 import {
 	AUTHENTICATED_APPS_LIST,
 	BACKGROUND_SCRIPT_RECEIVED_DATA,
+	EXTENSION_SESSION_INFO,
+	GET_EXTENSION_SESSION,
 	LOCAL,
 	PASSWORD,
+	SENDER_EXTENSION,
 	SESSION
 } from '$shared/const';
 import { parseMessageSchema, uint8ArrayToBase64 } from '$shared/helpers';
-import { getSessionKey, sendMessage, storageService } from '$shared/services';
+import {
+	createMessageWithId,
+	getDeviceKey,
+	responseToMessage,
+	sendMessage,
+	storageService
+} from '$shared/services';
 import {
 	AuthenticatedAppsListSchema,
 	HashSaltSchema,
@@ -48,23 +57,45 @@ export const fetchAuthenticatedAppsList = async (happId?: string) => {
 	return parsedAuthenticatedAppsListData.data;
 };
 
-export const sendMessageAndHandleResponse = async (message: Message) => {
-	const response = await sendMessage(message);
+export const sendMessageAndHandleResponse = async (message: Message, id?: string) => {
+	const messageWithId = id ? responseToMessage(message, id) : createMessageWithId(message);
+
+	const response = await sendMessage(messageWithId);
 
 	const parsedResponse = parseMessageSchema(response);
 
 	if (parsedResponse.data.action !== BACKGROUND_SCRIPT_RECEIVED_DATA)
 		throw new Error('Error sending data to webapp');
 };
 
+export const getExtensionSession = async () => {
+	const messageWithId = createMessageWithId({
+		sender: SENDER_EXTENSION,
+		action: GET_EXTENSION_SESSION
+	});
+
+	const response = await sendMessage(messageWithId);
+
+	const parsedResponse = parseMessageSchema(response);
+
+	if (!parsedResponse.success || parsedResponse.data.action !== EXTENSION_SESSION_INFO) {
+		throw new Error('Error getting extension session');
+	}
+
+	return parsedResponse.data.payload;
+};
+
 export const deriveSignPubKey = async (newIndex: number) => {
-	const sessionKey = await getSessionKey();
+	const session = await getExtensionSession();
 
-	if (!sessionKey.success) {
+	if (!session) {
 		throw new Error('Session data not found');
 	}
 
-	const keyUnlocked = await unlockKey(sessionKey.data, SESSION);
+	const encryptedDeviceKey = await getDeviceKey();
+
+	const keyUnlocked = await unlockKey(encryptedDeviceKey, session);
+
 	const { signPubKey } = keyUnlocked.derive(newIndex);
 
 	keyUnlocked.zero();