regression: 2fa inline errors for API calls

apps/meteor/app/utils/client/lib/RestApiClient.ts

@@ -1,7 +1,7 @@
 /* eslint-disable react-hooks/rules-of-hooks */
 import { RestClient } from '@rocket.chat/api-client';
 
-import { invokeTwoFactorModal } from '../../../../client/lib/2fa/process2faReturn';
+import { challenge2fa } from '../../../../client/lib/2fa/challenge2fa';
 import { baseURI } from '../../../../client/lib/baseURI';
 import { STORAGE_KEYS, getStoredItem } from '../../../../client/lib/sdk/storage';
 
@@ -28,7 +28,7 @@ export const APIClient = new RestApiClient({
 	baseUrl: baseURI.replace(/\/$/, ''),
 });
 
-APIClient.handleTwoFactorChallenge(invokeTwoFactorModal);
+APIClient.handleTwoFactorChallenge(challenge2fa);
 
 /**
  * The original rest api code throws the Response object, which is very useful

apps/meteor/client/lib/2fa/challenge2fa.ts

@@ -0,0 +1,161 @@
+import type { TwoFactorMethod } from '@rocket.chat/api-client';
+import { isTotpInvalidError, isTotpRequiredError, hasRequiredTwoFactorMethod, isTotpMaxAttemptsError } from '@rocket.chat/api-client';
+import { SHA256 } from '@rocket.chat/sha256';
+import { imperativeModal } from '@rocket.chat/ui-client';
+import { lazy } from 'react';
+
+import { getUser } from '../user';
+
+const TwoFactorModal = lazy(() => import('../../components/TwoFactorModal'));
+
+type RejectChallenge = (error: unknown) => void;
+type OnConfirm = (twoFactorCode: string) => Promise<void>;
+type OnClose = () => void;
+type UnresolvedChallenge = {
+	onConfirm: OnConfirm;
+	rejectChallenge: RejectChallenge;
+	onClose: OnClose;
+};
+
+// This is a "store" for callbacks that need to be accessed outside of the challenge2fa function's scope
+// e.g. We cannot update modal props, so callbacks that can change are stored here, and then wrapped in another function
+// that will call whatever is present here.
+let unresolvedChallenge: undefined | UnresolvedChallenge = undefined;
+
+const saveChallenge = (challenge: UnresolvedChallenge) => {
+	unresolvedChallenge = challenge;
+};
+
+const endChallenge = (rejectError?: unknown) => {
+	if (rejectError) {
+		unresolvedChallenge?.rejectChallenge(rejectError);
+	}
+	unresolvedChallenge = undefined;
+	imperativeModal.close();
+};
+
+const wrapWithEndChallenge = (cb: () => void) => () => {
+	endChallenge();
+	cb();
+};
+
+// Helper to extract the resolvers from a promise without TS complaining the functions might not exist
+const makeResolvablePromise = <T>() => {
+	let resolve: undefined | ((value: T) => void) = undefined;
+	let reject: undefined | ((error: unknown) => void) = undefined;
+	const promise = new Promise<T>((_resolve, _reject) => {
+		resolve = _resolve;
+		reject = _reject;
+	});
+
+	if (!resolve || !reject) {
+		throw new Error('Failed to create resolvable promise');
+	}
+
+	return [promise, resolve as (value: T | PromiseLike<T>) => void, reject as (error: unknown) => void] as const;
+};
+
+const getProps = (
+	method: 'totp' | 'email' | 'password',
+	emailOrUsername?: { username: string } | { email: string } | { id: string } | string,
+) => {
+	switch (method) {
+		case 'totp':
+			return { method };
+		case 'email':
+			return {
+				method,
+				emailOrUsername: typeof emailOrUsername === 'string' ? emailOrUsername : getUser()?.username,
+			};
+		case 'password':
+			return { method };
+	}
+};
+
+function assertModalProps(props: {
+	method: TwoFactorMethod;
+	emailOrUsername?: string;
+}): asserts props is { method: 'totp' } | { method: 'password' } | { method: 'email'; emailOrUsername: string } {
+	if (props.method === 'email' && typeof props.emailOrUsername !== 'string') {
+		throw new Error('Invalid Two Factor method');
+	}
+}
+
+type Request2faPromptOptions = {
+	error: unknown;
+	emailOrUsername?: string;
+};
+
+export const challenge2fa = ({ error, emailOrUsername }: Request2faPromptOptions): [Promise<string>, () => void] | undefined => {
+	if (isTotpMaxAttemptsError(error)) {
+		Promise.all([import('../../../app/utils/lib/i18n'), import('../toast')]).then(([{ t }, { dispatchToastMessage }]) => {
+			dispatchToastMessage({
+				type: 'error',
+				message: t('totp-max-attempts'),
+			});
+		});
+
+		if (unresolvedChallenge) {
+			endChallenge(error);
+		}
+
+		throw error;
+	}
+	if ((!isTotpRequiredError(error) && !isTotpInvalidError(error)) || !hasRequiredTwoFactorMethod(error)) {
+		if (unresolvedChallenge) {
+			endChallenge(error);
+		}
+
+		throw error;
+	}
+
+	const twoFactorMethod = 'details' in error ? error.details.method : 'password';
+
+	// `code` promise is resolved with the actual code
+	const [code, resolveCode, rejectCode] = makeResolvablePromise<string>();
+	// `challengePromise` is resolved async when the user completes the challenge
+	// or rejected with the new error in case this is a retry
+	const [challengePromise, resolveChallenge, rejectChallengePromise] = makeResolvablePromise<void>();
+
+	const onConfirm = async (twoFactorCode: string): Promise<void> => {
+		const actualCode = twoFactorMethod === 'password' ? SHA256(twoFactorCode) : twoFactorCode;
+		resolveCode(actualCode);
+		await challengePromise;
+		endChallenge();
+	};
+
+	const onClose = () => {
+		endChallenge();
+		rejectCode(new Error('Two-factor_authentication_cancelled'));
+	};
+
+	const rejectChallenge = (error: unknown) => {
+		rejectChallengePromise(error);
+	};
+
+	if (unresolvedChallenge) {
+		// This is a retry, the modal will catch this error in order to show inline information
+		unresolvedChallenge.rejectChallenge(error);
+		saveChallenge({ onConfirm, rejectChallenge, onClose });
+		return [code, wrapWithEndChallenge(resolveChallenge)];
+	}
+
+	saveChallenge({ onConfirm, rejectChallenge, onClose });
+
+	const method = error.details.method ? error.details.method : 'password';
+
+	const props = getProps(method, emailOrUsername || error.details.emailOrUsername || getUser()?.username);
+
+	assertModalProps(props);
+
+	imperativeModal.open({
+		component: TwoFactorModal,
+		props: {
+			...props,
+			onConfirm: (code) => unresolvedChallenge?.onConfirm(code),
+			onClose: () => unresolvedChallenge?.onClose(),
+		},
+	});
+
+	return [code, wrapWithEndChallenge(resolveChallenge)];
+};

apps/meteor/client/lib/2fa/process2faReturn.ts

@@ -88,7 +88,7 @@ export async function process2faAsyncReturn<TResult>({
 	onCode,
 	emailOrUsername,
 }: {
-	error: unknown;
+	error: MeteorErrorLike | undefined;
 	onCode: (code: string, method: string) => TResult | Promise<TResult>;
 	emailOrUsername: string | null | undefined;
 }): Promise<TResult> {
@@ -138,12 +138,13 @@ export const invokeTwoFactorModal = async (
 			props: {
 				...props,
 				onConfirm: async (code: string, method: string): Promise<void> => {
+					const actualCode = method === 'password' ? SHA256(code) : code;
 					if (validateCode) {
-						await validateCode(code, method);
+						await validateCode(actualCode, method);
 					}
 					isResolved = true;
 					imperativeModal.close();
-					resolve(method === 'password' ? SHA256(code) : code);
+					resolve(actualCode);
 				},
 				onClose: (): void => {
 					if (isClosed) {

apps/meteor/client/lib/2fa/utils.ts

@@ -1,17 +1,6 @@
 import type { MeteorErrorLike } from './types';
 
-export const isTotpRequiredError = (
-	error: unknown,
-): error is MeteorErrorLike & ({ error: 'totp-required' } | { errorType: 'totp-required' }) =>
-	typeof error === 'object' &&
-	((error as { error?: unknown } | undefined)?.error === 'totp-required' ||
-		(error as { errorType?: unknown } | undefined)?.errorType === 'totp-required');
-
-export const isTotpInvalidError = (
-	error: unknown,
-): error is MeteorErrorLike & ({ error: 'totp-invalid' } | { errorType: 'totp-invalid' }) =>
-	(error as { error?: unknown } | undefined)?.error === 'totp-invalid' ||
-	(error as { errorType?: unknown } | undefined)?.errorType === 'totp-invalid';
+export { isTotpInvalidError, isTotpRequiredError } from '@rocket.chat/api-client';
 
 export const isTotpMaxAttemptsError = (
 	error: unknown,

packages/api-client/src/RestClientInterface.ts

@@ -87,6 +87,6 @@ export interface RestClientInterface {
 	send(endpoint: string, method: string, options?: Omit<RequestInit, 'method'>): Promise<Response>;
 
 	handleTwoFactorChallenge(
-		cb: (args: { method: 'totp' | 'email' | 'password'; emailOrUsername?: string; invalidAttempt?: boolean }) => Promise<string>,
+		challenge: (args: { error: unknown }) => [code: Promise<string>, resolveChallenge: () => void] | undefined,
 	): void;
 }

packages/api-client/src/errors.ts

@@ -1,20 +1,44 @@
 const twoFactorMethods = ['totp', 'email', 'password'] as const;
 
-type TwoFactorMethod = (typeof twoFactorMethods)[number];
+export type TwoFactorMethod = (typeof twoFactorMethods)[number];
 
-export const isTotpRequiredError = (error: unknown): error is { error: 'totp-required' } | { errorType: 'totp-required' } =>
-	typeof error === 'object' &&
-	((error as { error?: unknown } | undefined)?.error === 'totp-required' ||
-		(error as { errorType?: unknown } | undefined)?.errorType === 'totp-required');
+const INVALID_TOTP = 'totp-invalid';
+const REQUIRED_TOTP = 'totp-required';
+const MAX_ATTEMPTS_TOTP = 'totp-max-attempts';
 
-export const isTotpInvalidError = (error: unknown): error is { error: 'totp-invalid' } | { errorType: 'totp-invalid' } =>
-	(error as { error?: unknown } | undefined)?.error === 'totp-invalid' ||
-	(error as { errorType?: unknown } | undefined)?.errorType === 'totp-invalid';
+type TwoFactorErrorTypes = typeof INVALID_TOTP | typeof REQUIRED_TOTP | typeof MAX_ATTEMPTS_TOTP;
+type TwoFactorError<T extends TwoFactorErrorTypes> = {
+	details: { method: TwoFactorMethod; emailOrUsername?: string };
+} & ({ error: T } | { errorType: T });
+
+export type TwoFactorErrorResponse = TwoFactorError<TwoFactorErrorTypes>;
+
+const isTotpError = <TErrorType extends TwoFactorErrorTypes>(error: unknown, type: TErrorType): error is TwoFactorError<TErrorType> => {
+	if (typeof error !== 'object' || error === null) {
+		return false;
+	}
+
+	if ('error' in error && error.error === type) {
+		return true;
+	}
+
+	if ('errorType' in error && error.errorType === type) {
+		return true;
+	}
+
+	return false;
+};
+
+export const isTotpRequiredError = (error: unknown) => isTotpError(error, REQUIRED_TOTP);
+
+export const isTotpInvalidError = (error: unknown) => isTotpError(error, INVALID_TOTP);
+
+export const isTotpMaxAttemptsError = (error: unknown) => isTotpError(error, MAX_ATTEMPTS_TOTP);
 
 export const isTwoFactorMethod = (method: string): method is TwoFactorMethod => twoFactorMethods.includes(method as TwoFactorMethod);
 
 export const hasRequiredTwoFactorMethod = (error: unknown): error is { details: { method: TwoFactorMethod; emailOrUsername?: string } } => {
-	const details = error && typeof error === 'object' && 'details' in error && (error.details as unknown);
+	const details = error && typeof error === 'object' && 'details' in error && error.details;
 	return (
 		typeof details === 'object' &&
 		details !== null &&

packages/api-client/src/index.ts

@@ -11,9 +11,9 @@ import { stringify } from 'query-string';
 
 import type { Credentials } from './Credentials';
 import type { Middleware, RestClientInterface } from './RestClientInterface';
-import { hasRequiredTwoFactorMethod, isTotpInvalidError, isTotpRequiredError } from './errors';
 
 export type { RestClientInterface, Credentials };
+export * from './errors';
 
 const pipe =
 	<T extends (...args: any[]) => any>(fn: T) =>
@@ -51,11 +51,7 @@ const checkIfIsFormData = (data: any = {}): boolean => {
 };
 
 export class RestClient implements RestClientInterface {
-	private twoFactorHandler?: (args: {
-		method: 'totp' | 'email' | 'password';
-		emailOrUsername?: string;
-		invalidAttempt?: boolean;
-	}) => Promise<string>;
+	private challenge2fa?: (args: { error: unknown }) => [code: Promise<string>, resolveChallenge: () => void] | undefined;
 
 	private readonly baseUrl: string;
 
@@ -223,28 +219,35 @@ export class RestClient implements RestClientInterface {
 
 			const error = await clone.json();
 
-			if ((isTotpRequiredError(error) || isTotpInvalidError(error)) && hasRequiredTwoFactorMethod(error) && this.twoFactorHandler) {
-				const method2fa = 'details' in error ? error.details.method : 'password';
-
-				const code = await this.twoFactorHandler({
-					method: method2fa,
-					emailOrUsername: error.details.emailOrUsername,
-					invalidAttempt: isTotpInvalidError(error),
-				});
-
-				return this.send(endpoint, method, {
-					...options,
-					headers: {
-						...this.getCredentialsAsHeaders(),
-						...this.headers,
-						...headers,
-						'x-2fa-code': code,
-						'x-2fa-method': method2fa,
-					},
-				});
+			const method2fa = 'details' in error ? error.details.method : 'password';
+
+			if (!this.challenge2fa) {
+				return Promise.reject(error);
 			}
 
-			return Promise.reject(response);
+			const challengeResult = this.challenge2fa({ error });
+
+			if (!challengeResult) {
+				return Promise.reject(error);
+			}
+
+			const [code, resolveChallenge] = challengeResult;
+
+			const resolvedCode = await code;
+
+			return this.send(endpoint, method, {
+				...options,
+				headers: {
+					...this.getCredentialsAsHeaders(),
+					...this.headers,
+					...headers,
+					'x-2fa-code': resolvedCode,
+					'x-2fa-method': method2fa,
+				},
+			}).then((response) => {
+				resolveChallenge();
+				return response;
+			});
 		});
 	}
 
@@ -297,9 +300,7 @@ export class RestClient implements RestClientInterface {
 		} as RestClientInterface['send'];
 	}
 
-	handleTwoFactorChallenge(
-		cb: (args: { method: 'totp' | 'email' | 'password'; emailOrUsername?: string; invalidAttempt?: boolean }) => Promise<string>,
-	): void {
-		this.twoFactorHandler = cb;
+	handleTwoFactorChallenge(challenge2fa: (args: { error: unknown }) => [Promise<string>, () => void] | undefined): void {
+		this.challenge2fa = challenge2fa;
 	}
 }