From 08a446fefd51a182c329ac125baee005a7ce894c Mon Sep 17 00:00:00 2001
From: Lucas Smith <me@lucasjamessmith.me>
Date: Sat, 16 Nov 2024 09:17:45 +1100
Subject: [PATCH] feat: support windows for 2fa tokens (#1478)

## Description

When using 2fa enabled authentication on direct templates we run into an
issue where a 2fa token has been attached to a field but it's submitted
at a later point.

To better facilitate this we have introduced the ability to have a
window of valid tokens.

This won't affect other signing methods since tokens are verified
immediately after they're entered.

## Related Issue

N/A

## Changes Made

- Updated our validate2FAToken method to use a window based approach
rather than the default verify method.

## Testing Performed

- Created a series of tokens and tested upon different intervals and
windows to confirm functionality works as expected.
---
 .../lib/server-only/2fa/verify-2fa-token.ts   | 28 +++++++++++++++----
 .../document/is-recipient-authorized.ts       |  1 +
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/packages/lib/server-only/2fa/verify-2fa-token.ts b/packages/lib/server-only/2fa/verify-2fa-token.ts
index 0e8ec6afc..c5a4c3c95 100644
--- a/packages/lib/server-only/2fa/verify-2fa-token.ts
+++ b/packages/lib/server-only/2fa/verify-2fa-token.ts
@@ -1,21 +1,25 @@
 import { base32 } from '@scure/base';
-import { TOTPController } from 'oslo/otp';
+import { generateHOTP } from 'oslo/otp';
 
 import type { User } from '@documenso/prisma/client';
 
 import { DOCUMENSO_ENCRYPTION_KEY } from '../../constants/crypto';
 import { symmetricDecrypt } from '../../universal/crypto';
 
-const totp = new TOTPController();
-
 type VerifyTwoFactorAuthenticationTokenOptions = {
   user: User;
   totpCode: string;
+  // The number of windows to look back
+  window?: number;
+  // The duration that the token is valid for in seconds
+  period?: number;
 };
 
 export const verifyTwoFactorAuthenticationToken = async ({
   user,
   totpCode,
+  window = 1,
+  period = 30_000,
 }: VerifyTwoFactorAuthenticationTokenOptions) => {
   const key = DOCUMENSO_ENCRYPTION_KEY;
 
@@ -27,7 +31,21 @@ export const verifyTwoFactorAuthenticationToken = async ({
     'utf-8',
   );
 
-  const isValidToken = await totp.verify(totpCode, base32.decode(secret));
+  const decodedSecret = base32.decode(secret);
 
-  return isValidToken;
+  let now = Date.now();
+
+  for (let i = 0; i < window; i++) {
+    const counter = Math.floor(now / period);
+
+    const hotp = await generateHOTP(decodedSecret, counter);
+
+    if (totpCode === hotp) {
+      return true;
+    }
+
+    now -= period;
+  }
+
+  return false;
 };
diff --git a/packages/lib/server-only/document/is-recipient-authorized.ts b/packages/lib/server-only/document/is-recipient-authorized.ts
index 151235fe2..e9596211a 100644
--- a/packages/lib/server-only/document/is-recipient-authorized.ts
+++ b/packages/lib/server-only/document/is-recipient-authorized.ts
@@ -112,6 +112,7 @@ export const isRecipientAuthorized = async ({
       return await verifyTwoFactorAuthenticationToken({
         user,
         totpCode: token,
+        window: 10, // 5 minutes worth of tokens
       });
     })
     .exhaustive();