Skip to content

Commit

Permalink
Replace Promise.withResolvers for compatiblity with older browers; Fi…
Browse files Browse the repository at this point in the history
…x bug where MFA couldn't be retried after a failed attempt; Add extra tests.
  • Loading branch information
Joerger committed Dec 19, 2024
1 parent 90ef02a commit 5ee3d03
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 14 deletions.
47 changes: 44 additions & 3 deletions web/packages/teleport/src/lib/useMfa.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,13 @@ const mockChallengeReq: CreateAuthenticateChallengeRequest = {
};

describe('useMfa', () => {
beforeEach(() => jest.spyOn(console, 'error').mockImplementation());
beforeEach(() => {
jest.spyOn(console, 'error').mockImplementation();
});

afterEach(() => {
jest.clearAllMocks();
});

test('mfa required', async () => {
jest.spyOn(auth, 'getMfaChallenge').mockResolvedValueOnce(mockChallenge);
Expand Down Expand Up @@ -184,15 +190,18 @@ describe('useMfa', () => {
throw err;
});

const { result: mfa } = renderHook(() => useMfa({}));
const { result: mfa } = renderHook(() =>
useMfa({
req: mockChallengeReq,
})
);

const respPromise = mfa.current.getChallengeResponse();
await waitFor(() => {
expect(auth.getMfaChallenge).toHaveBeenCalledWith(mockChallengeReq);
});
await mfa.current.submit('webauthn');

await expect(respPromise).rejects.toThrow(err);
await waitFor(() => {
expect(mfa.current.attempt).toEqual({
status: 'error',
Expand All @@ -201,5 +210,37 @@ describe('useMfa', () => {
data: null,
});
});

// After an error, the mfa response promise remains in an unresolved state,
// allowing for retries.
jest
.spyOn(auth, 'getMfaChallengeResponse')
.mockResolvedValueOnce(mockResponse);
await mfa.current.submit('webauthn');
expect(await respPromise).toEqual(mockResponse);
});

test('reset mfa attempt', async () => {
jest.spyOn(auth, 'getMfaChallenge').mockResolvedValue(mockChallenge);
const { result: mfa } = renderHook(() =>
useMfa({
req: mockChallengeReq,
})
);

const respPromise = mfa.current.getChallengeResponse();
await waitFor(() => {
expect(auth.getMfaChallenge).toHaveBeenCalled();
});

mfa.current.resetAttempt();

await expect(respPromise).rejects.toThrow(
new Error('MFA attempt cancelled by user')
);

await waitFor(() => {
expect(mfa.current.attempt.status).toEqual('error');
});
});
});
47 changes: 36 additions & 11 deletions web/packages/teleport/src/lib/useMfa.ts
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,12 @@ export type MfaProps = {
isMfaRequired?: boolean | null;
};

type mfaResponsePromiseWithResolvers = {
promise: Promise<MfaChallengeResponse>;
resolve: (v: MfaChallengeResponse) => void;
reject: (v?: any) => void;
};

/**
* Use the returned object to request MFA checks with a shared state.
* When MFA authentication is in progress, the object's properties can
Expand All @@ -48,10 +54,10 @@ export type MfaProps = {
export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
const [mfaRequired, setMfaRequired] = useState<boolean>();
const [options, setMfaOptions] = useState<MfaOption[]>();

const [challenge, setMfaChallenge] = useState<MfaAuthenticateChallenge>();
const mfaResponsePromise =
useRef<PromiseWithResolvers<MfaChallengeResponse>>();

const mfaResponsePromiseWithResolvers =
useRef<mfaResponsePromiseWithResolvers>();

useEffect(() => {
setMfaRequired(isMfaRequired);
Expand All @@ -70,6 +76,7 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
// 4. Receive the mfa response through the mfaResponsePromise ref and return it.
//
// The caller should also display errors seen in attempt.

const [attempt, getResponse, setMfaAttempt] = useAsync(
useCallback(
async (challenge?: MfaAuthenticateChallenge) => {
Expand All @@ -88,22 +95,36 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {

// Prepare a new promise to collect the mfa response retrieved
// through the submit function.
mfaResponsePromise.current = Promise.withResolvers();
let resolve, reject;
const promise = new Promise<MfaChallengeResponse>((res, rej) => {
resolve = res;
reject = rej;
});

mfaResponsePromiseWithResolvers.current = {
promise,
resolve,
reject,
};

setMfaChallenge(challenge);
try {
return await mfaResponsePromise.current.promise;
return await promise;
} finally {
mfaResponsePromise.current = null;
mfaResponsePromiseWithResolvers.current = null;
setMfaChallenge(null);
}
},
[req, mfaResponsePromise, options, mfaRequired]
[req, mfaResponsePromiseWithResolvers, options, mfaRequired]
)
);

const resetAttempt = () => {
if (mfaResponsePromise.current) mfaResponsePromise.current.reject();
mfaResponsePromise.current = null;
if (mfaResponsePromiseWithResolvers.current)
mfaResponsePromiseWithResolvers.current.reject(
new Error('MFA attempt cancelled by user')
);
mfaResponsePromiseWithResolvers.current = null;
setMfaChallenge(null);
setMfaAttempt(makeEmptyAttempt());
};
Expand All @@ -119,8 +140,12 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {

const submit = useCallback(
async (mfaType?: DeviceType, totpCode?: string) => {
if (!mfaResponsePromiseWithResolvers.current) {
throw new Error('submit called without an in flight MFA attempt');
}

try {
await mfaResponsePromise.current.resolve(
await mfaResponsePromiseWithResolvers.current.resolve(
await auth.getMfaChallengeResponse(challenge, mfaType, totpCode)
);
} catch (err) {
Expand All @@ -132,7 +157,7 @@ export function useMfa({ req, isMfaRequired }: MfaProps): MfaState {
});
}
},
[challenge, mfaResponsePromise, setMfaAttempt]
[challenge, mfaResponsePromiseWithResolvers, setMfaAttempt]
);

return {
Expand Down

0 comments on commit 5ee3d03

Please sign in to comment.