Add more edge-cases to MFA

backend/src/controllers/v2/authController.ts

@@ -243,7 +243,6 @@ export const sendMfaToken = async (req: Request, res: Response) => {
  * @param res 
  */
 export const verifyMfaToken = async (req: Request, res: Response) => {
-  try {
     const { email, mfaToken } = req.body;
 
     await TokenService.validateToken({
@@ -296,15 +295,6 @@ export const verifyMfaToken = async (req: Request, res: Response) => {
       resObj.protectedKeyTag = user.protectedKeyTag;
     }
     
-    // case: user does not have MFA enabled
-    // return (access) token in response
     return res.status(200).send(resObj); 
-  } catch (err) {
-    Sentry.setUser(null);
-    Sentry.captureException(err);
-    return res.status(400).send({
-      message: 'Failed to authenticate. Try again?'
-    });
-  }
 }
 

backend/src/helpers/token.ts

@@ -12,7 +12,7 @@ import {
 import {
     SALT_ROUNDS
 } from '../config';
-import { ForbiddenRequestError } from '../utils/errors';
+import { UnauthorizedRequestError } from '../utils/errors';
 
 /**
  * Create and store a token in the database for purpose [type]
@@ -34,7 +34,7 @@ const createTokenHelper = async ({
     phoneNumber?: string;
     organizationId?: Types.ObjectId
 }) => {
-    let token, expiresAt;
+    let token, expiresAt, triesLeft;
     try {
         // generate random token based on specified token use-case
         // type [type]
@@ -47,6 +47,7 @@ const createTokenHelper = async ({
             case TOKEN_EMAIL_MFA:
                 // generate random 6-digit code
                 token = String(crypto.randomInt(Math.pow(10, 5), Math.pow(10, 6) - 1));
+                triesLeft = 5;
                 expiresAt = new Date((new Date()).getTime() + 300000);
                 break;
             case TOKEN_EMAIL_ORG_INVITATION:
@@ -78,6 +79,7 @@ const createTokenHelper = async ({
             phoneNumber?: string;
             organization?: Types.ObjectId;
             tokenHash: string;
+            triesLeft?: number;
             expiresAt: Date;
         }
 
@@ -100,6 +102,10 @@ const createTokenHelper = async ({
             query.organization = organizationId 
             update.organization = organizationId 
         } 
+        
+        if (triesLeft) {
+            update.triesLeft = triesLeft;
+        }
        
         await TokenData.findOneAndUpdate(
             query,
@@ -157,19 +163,51 @@ const validateTokenHelper = async ({
     if (!tokenData) throw new Error('Failed to find token to validate');
     
     if (tokenData.expiresAt < new Date()) {
+        // case: token expired
         await TokenData.findByIdAndDelete(tokenData._id);
-        throw ForbiddenRequestError({
+        throw UnauthorizedRequestError({
-            message: 'Failed token data validation due to token is no longer valid'
+            message: 'MFA session expired. Please log in again',
+            context: {
+                code: 'mfa_expired'
+            }
         });
     }
 
     const isValid = await bcrypt.compare(token, tokenData.tokenHash);
     if (!isValid) {
-        throw ForbiddenRequestError({
+        // case: token is not valid
-            message: 'Failed token data validation due to incorrect token'
+        if (tokenData?.triesLeft !== undefined) {
+            // case: token has a try-limit
+            if (tokenData.triesLeft === 1) {
+                // case: token is out of tries
+                await TokenData.findByIdAndDelete(tokenData._id);
+            } else {
+                // case: token has more than 1 try left
+                await TokenData.findByIdAndUpdate(tokenData._id, {
+                    triesLeft: tokenData.triesLeft - 1
+                }, {
+                    new: true
+                });
+            }
+
+            throw UnauthorizedRequestError({
+                message: 'MFA code is invalid',
+                context: {
+                    code: 'mfa_invalid',
+                    triesLeft: tokenData.triesLeft - 1
+                }
+            });
+        }
+
+        throw UnauthorizedRequestError({
+            message: 'MFA code is invalid',
+            context: {
+                code: 'mfa_invalid'
+            }
         });
     }
 
+    // case: token is valid
     await TokenData.findByIdAndDelete(tokenData._id);
 }
 

backend/src/middleware/requestErrorHandler.ts

@@ -17,9 +17,7 @@ export const requestErrorHandler: ErrorRequestHandler = (error: RequestError | E
     }
 
     //TODO: Find better way to type check for error. In current setting you need to cast type to get the functions and variables from RequestError
-    if (error instanceof TokenExpiredError) {
+    if (!(error instanceof RequestError)) {
-        error = UnauthorizedRequestError({ stack: error.stack, message: 'Token expired' });
-    } else if (!(error instanceof RequestError)) {
         error = InternalServerError({ context: { exception: error.message }, stack: error.stack })
         getLogger('backend-main').log((<RequestError>error).levelName.toLowerCase(), (<RequestError>error).message)
     }

backend/src/models/tokenData.ts

@@ -6,6 +6,7 @@ export interface ITokenData {
   phoneNumber?: string;
   organization?: Types.ObjectId;
   tokenHash: string;
+  triesLeft?: number;
   expiresAt: Date;
   createdAt: Date;
   updatedAt: Date;
@@ -37,6 +38,9 @@ const tokenDataSchema = new Schema<ITokenData>({
     select: false,
     required: true
   },
+  triesLeft: {
+    type: Number
+  },
   expiresAt: {
     type: Date,
     expires: 0,

frontend/src/components/login/LoginStep.tsx

@@ -63,7 +63,6 @@ export default function LoginStep ({
             }
 
         } catch (err) {
-            console.error(err);
             setLoginError(true);
         }
         

frontend/src/components/login/MFAStep.tsx

@@ -30,6 +30,18 @@ const props = {
   }
 } as const;
 
+interface VerifyMfaTokenError {
+  response: {
+    data: {
+      context: {
+        code: string;
+        triesLeft: number;
+      }
+    },
+    status: number;
+  }
+}
+
 /**
  * 2nd step of login - users enter their MFA code
  * @param {Object} obj
@@ -73,7 +85,15 @@ export default function MFAStep({
       }
       
     } catch (err) {
-      console.error(err);
+      const error = err as VerifyMfaTokenError;
+      
+      if (error?.response?.status === 500) {
+        window.location.reload();
+      } else if (error?.response?.data?.context?.triesLeft === 0) {
+        window.location.reload();
+        router.push('/login'); 
+      }
+
       setIsLoading(false);
       setCodeError(true);
     }