Validate authToken from user data file before returning it (#871)

* WIP

* Simplify logic

* Address review feedback

* Display error modal when login process fails

package.json

@@ -142,7 +142,8 @@
 		"url-loader": "^4.1.1",
 		"wpcom": "^5.4.2",
 		"yargs": "17.7.2",
-		"yauzl": "^3.2.0"
+		"yauzl": "^3.2.0",
+		"zod": "^3.24.1"
 	},
 	"optionalDependencies": {
 		"appdmg": "^0.6.6"

src/components/auth-provider.tsx

@@ -1,4 +1,5 @@
 import * as Sentry from '@sentry/electron/renderer';
+import { useI18n } from '@wordpress/react-i18n';
 import { createContext, useState, useEffect, useMemo, useCallback, ReactNode } from 'react';
 import WPCOM from 'wpcom';
 import { useI18nData } from '../hooks/use-i18n-data';
@@ -36,23 +37,26 @@ const AuthProvider: React.FC< AuthProviderProps > = ( { children } ) => {
 	const [ client, setClient ] = useState< WPCOM | undefined >( undefined );
 	const [ user, setUser ] = useState< AuthContextType[ 'user' ] >( undefined );
 	const { locale } = useI18nData();
+	const { __ } = useI18n();
 
 	const authenticate = useCallback( () => getIpcApi().authenticate(), [] );
 
 	useIpcListener( 'auth-updated', ( _event, { token, error } ) => {
 		if ( error ) {
-			Sentry.captureException( error );
+			getIpcApi().showErrorMessageBox( {
+				title: __( 'Authentication error' ),
+				message: __( 'Please try again.' ),
+			} );
 			return;
 		}
+
 		setIsAuthenticated( true );
 		setClient( createWpcomClient( token.accessToken, locale ) );
-		if ( token.id || token.email || token.displayName ) {
-			setUser( {
-				id: token.id || null,
-				email: token.email || '',
-				displayName: token.displayName || '',
-			} );
-		}
+		setUser( {
+			id: token.id,
+			email: token.email,
+			displayName: token.displayName || '',
+		} );
 	} );
 
 	const logout = useCallback( async () => {
@@ -70,22 +74,20 @@ const AuthProvider: React.FC< AuthProviderProps > = ( { children } ) => {
 	useEffect( () => {
 		async function run() {
 			try {
-				const isAuthenticated = await getIpcApi().isAuthenticated();
-				setIsAuthenticated( isAuthenticated );
-				if ( isAuthenticated ) {
-					const token = await getIpcApi().getAuthenticationToken();
-					if ( ! token ) {
-						return;
-					}
-					setClient( createWpcomClient( token.accessToken, locale ) );
-					if ( token.id || token.email || token.displayName ) {
-						setUser( {
-							id: token.id || null,
-							email: token.email || '',
-							displayName: token.displayName || '',
-						} );
-					}
+				const token = await getIpcApi().getAuthenticationToken();
+
+				if ( ! token ) {
+					setIsAuthenticated( false );
+					return;
 				}
+
+				setIsAuthenticated( true );
+				setClient( createWpcomClient( token.accessToken, locale ) );
+				setUser( {
+					id: token.id,
+					email: token.email,
+					displayName: token.displayName || '',
+				} );
 			} catch ( err ) {
 				console.error( err );
 				Sentry.captureException( err );

src/index.ts

@@ -24,7 +24,7 @@ import {
 	executeCLICommand,
 } from './lib/cli';
 import { getUserLocaleWithFallback } from './lib/locale-node';
-import { handleAuthCallback, setUpAuthCallbackHandler } from './lib/oauth';
+import { onOpenUrlCallback } from './lib/oauth';
 import { getSentryReleaseInfo } from './lib/sentry-release';
 import { setupLogging } from './logging';
 import { createMainWindow, getMainWindow } from './main-window';
@@ -72,38 +72,13 @@ if ( gotTheLock && ! isInInstaller ) {
 	}
 }
 
-const onOpenUrlCallback = async ( url: string ) => {
-	const urlObject = new URL( url );
-	const { host, hash, searchParams } = urlObject;
-	if ( host === 'auth' ) {
-		handleAuthCallback( hash ).then( ( authResult ) => {
-			if ( authResult instanceof Error ) {
-				ipcMain.emit( 'auth-callback', null, { error: authResult } );
-			} else {
-				ipcMain.emit( 'auth-callback', null, { token: authResult } );
-			}
-		} );
-	}
-
-	if ( host === 'sync-connect-site' ) {
-		const remoteSiteId = parseInt( searchParams.get( 'remoteSiteId' ) ?? '' );
-		const studioSiteId = searchParams.get( 'studioSiteId' );
-		if ( remoteSiteId && studioSiteId ) {
-			const mainWindow = await getMainWindow();
-			mainWindow.webContents.send( 'sync-connect-site', { remoteSiteId, studioSiteId } );
-		}
-	}
-};
-
 async function appBoot() {
 	app.setName( packageJson.productName );
 
 	Menu.setApplicationMenu( null );
 
 	setupCustomProtocolHandler();
 
-	setUpAuthCallbackHandler();
-
 	setupLogging();
 
 	setupUpdates();

src/ipc-handlers.ts

@@ -17,6 +17,7 @@ import nodePath from 'path';
 import * as Sentry from '@sentry/electron/main';
 import { __, LocaleData, defaultI18n } from '@wordpress/i18n';
 import archiver from 'archiver';
+import { StoredToken } from 'src/lib/oauth';
 import { DEFAULT_PHP_VERSION } from '../vendor/wp-now/src/constants';
 import { MAIN_MIN_WIDTH, SIDEBAR_WIDTH } from './constants';
 import { ACTIVE_SYNC_OPERATIONS } from './lib/active-sync-operations';
@@ -648,7 +649,7 @@ export function authenticate( _event: IpcMainInvokeEvent ) {
 
 export async function getAuthenticationToken(
 	_event: IpcMainInvokeEvent
-): Promise< oauthClient.StoredToken | null > {
+): Promise< StoredToken | null > {
 	return oauthClient.getAuthenticationToken();
 }
 

src/lib/oauth.ts

@@ -1,37 +1,38 @@
-import { ipcMain } from 'electron';
 import * as Sentry from '@sentry/electron/main';
 import wpcom from 'wpcom';
+import { z } from 'zod';
 import { PROTOCOL_PREFIX, WP_AUTHORIZE_ENDPOINT, CLIENT_ID, SCOPES } from '../constants';
 import { shellOpenExternalWrapper } from '../lib/shell-open-external-wrapper';
 import { getMainWindow } from '../main-window';
 import { loadUserData, saveUserData } from '../storage/user-data';
 
-export interface StoredToken {
-	accessToken?: string;
-	expiresIn?: number;
-	expirationTime?: number;
-	id?: number;
-	email?: string;
-	displayName?: string;
-}
 const REDIRECT_URI = `${ PROTOCOL_PREFIX }://auth`;
-const TOKEN_KEY = 'authToken';
+const authTokenSchema = z.object( {
+	accessToken: z.string(),
+	expiresIn: z.number(),
+	expirationTime: z.number(),
+	id: z.number(),
+	email: z.string(),
+	displayName: z.string().optional(),
+} );
+
+export type StoredToken = z.infer< typeof authTokenSchema >;
 
 async function getToken(): Promise< StoredToken | null > {
 	try {
 		const userData = await loadUserData();
-		return userData[ TOKEN_KEY ] ?? null;
+		return authTokenSchema.parse( userData.authToken );
 	} catch ( error ) {
 		return null;
 	}
 }
 
-async function storeToken( tokens: StoredToken ) {
+async function storeToken( token: StoredToken ) {
 	try {
 		const userData = await loadUserData();
 		await saveUserData( {
 			...userData,
-			[ TOKEN_KEY ]: tokens,
+			authToken: token,
 		} );
 	} catch ( error ) {
 		console.error( 'Failed to store token', error );
@@ -43,7 +44,7 @@ export async function clearAuthenticationToken() {
 		const userData = await loadUserData();
 		await saveUserData( {
 			...userData,
-			[ TOKEN_KEY ]: undefined,
+			authToken: undefined,
 		} );
 	} catch ( error ) {
 		return;
@@ -53,10 +54,7 @@ export async function clearAuthenticationToken() {
 export async function getAuthenticationToken(): Promise< StoredToken | null > {
 	// Check if tokens already exist and are valid
 	const existingToken = await getToken();
-	if (
-		existingToken?.accessToken &&
-		new Date().getTime() < ( existingToken?.expirationTime ?? 0 )
-	) {
+	if ( existingToken && new Date().getTime() < existingToken.expirationTime ) {
 		return existingToken;
 	}
 	return null;
@@ -67,32 +65,32 @@ export async function isAuthenticated(): Promise< boolean > {
 	return !! token;
 }
 
-export async function handleAuthCallback( hash: string ): Promise< Error | StoredToken > {
+async function handleAuthCallback( hash: string ): Promise< StoredToken > {
 	const params = new URLSearchParams( hash.substring( 1 ) );
 	const error = params.get( 'error' );
+
 	if ( error ) {
 		// Close the browser if code found or error
-		return new Error( error );
+		throw new Error( error );
 	}
+
 	const accessToken = params.get( 'access_token' ) ?? '';
 	const expiresIn = parseInt( params.get( 'expires_in' ) ?? '0' );
+
 	if ( isNaN( expiresIn ) || expiresIn === 0 || ! accessToken ) {
-		return new Error( 'Error while getting token' );
-	}
-	let response: { ID?: number; email?: string; display_name?: string } = {};
-	try {
-		response = await new wpcom( accessToken ).req.get( '/me?fields=ID,email,display_name' );
-	} catch ( error ) {
-		Sentry.captureException( error );
+		throw new Error( 'Error while getting token' );
 	}
-	return {
+
+	const response = await new wpcom( accessToken ).req.get( '/me?fields=ID,email,display_name' );
+
+	return authTokenSchema.parse( {
 		expiresIn,
 		expirationTime: new Date().getTime() + expiresIn * 1000,
 		accessToken,
 		id: response.ID,
 		email: response.email,
 		displayName: response.display_name,
-	};
+	} );
 }
 
 export function authenticate(): void {
@@ -103,14 +101,27 @@ export function authenticate(): void {
 	shellOpenExternalWrapper( authUrl );
 }
 
-export function setUpAuthCallbackHandler() {
-	ipcMain.on( 'auth-callback', async ( _event, { token, error } ) => {
-		const mainWindow = await getMainWindow();
-		if ( error ) {
+export async function onOpenUrlCallback( url: string ) {
+	const urlObject = new URL( url );
+	const { host, hash, searchParams } = urlObject;
+
+	if ( host === 'auth' ) {
+		try {
+			const authResult = await handleAuthCallback( hash );
+			const mainWindow = await getMainWindow();
+			await storeToken( authResult );
+			mainWindow.webContents.send( 'auth-updated', { token: authResult } );
+		} catch ( error ) {
+			Sentry.captureException( error );
+			const mainWindow = await getMainWindow();
 			mainWindow.webContents.send( 'auth-updated', { error } );
-		} else {
-			await storeToken( token );
-			mainWindow.webContents.send( 'auth-updated', { token } );
 		}
-	} );
+	} else if ( host === 'sync-connect-site' ) {
+		const remoteSiteId = parseInt( searchParams.get( 'remoteSiteId' ) ?? '' );
+		const studioSiteId = searchParams.get( 'studioSiteId' );
+		if ( remoteSiteId && studioSiteId ) {
+			const mainWindow = await getMainWindow();
+			mainWindow.webContents.send( 'sync-connect-site', { remoteSiteId, studioSiteId } );
+		}
+	}
 }