[WIP] feat: Add recovery page

apps/platform/trpc/routers/authRouter/recoveryRouter.ts

@@ -3,16 +3,25 @@ import { Argon2id } from 'oslo/password';
 import { router, publicRateLimitedProcedure } from '../../trpc';
 import { eq } from '@u22n/database/orm';
 import { accounts } from '@u22n/database/schema';
-import { nanoIdToken, zodSchemas } from '@u22n/utils';
+import {
+  nanoIdToken,
+  strongPasswordSchema,
+  typeIdValidator,
+  zodSchemas
+} from '@u22n/utils';
 import { TRPCError } from '@trpc/server';
 import { createLuciaSessionCookie } from '../../../utils/session';
-import { decodeHex } from 'oslo/encoding';
-import { TOTPController } from 'oslo/otp';
-import { setCookie } from 'hono/cookie';
+import { decodeHex, encodeHex } from 'oslo/encoding';
+import { TOTPController, createTOTPKeyURI } from 'oslo/otp';
+import { deleteCookie, getCookie, setCookie } from 'hono/cookie';
 import { env } from '../../../env';
 import { storage } from '../../../storage';
+import { ms } from 'itty-time';
 
 export const recoveryRouter = router({
+  /**
+   * @deprecated use `getRecoveryVerificationToken` instead
+   */
   recoverAccount: publicRateLimitedProcedure.recoverAccount
     .input(
       z.object({
@@ -163,5 +172,253 @@ export const recoveryRouter = router({
         message:
           'Something went wrong, you should never see this message. Please report to team immediately.'
       });
+    }),
+  getRecoveryVerificationToken: publicRateLimitedProcedure.recoverAccount
+    .input(
+      z
+        .object({
+          username: zodSchemas.username(2),
+          recoveryCode: zodSchemas.nanoIdToken()
+        })
+        .and(
+          z.union([
+            z.object({ password: z.string().min(8) }),
+            z.object({ twoFactorCode: z.string().min(6).max(6) })
+          ])
+        )
+    )
+    .query(async ({ input, ctx }) => {
+      const { db } = ctx;
+
+      const account = await db.query.accounts.findFirst({
+        where: eq(accounts.username, input.username),
+        columns: {
+          id: true,
+          publicId: true,
+          username: true,
+          passwordHash: true,
+          twoFactorSecret: true,
+          twoFactorEnabled: true,
+          recoveryCode: true
+        }
+      });
+
+      if (
+        !account ||
+        !account.recoveryCode ||
+        !account.passwordHash ||
+        !account.twoFactorSecret
+      ) {
+        throw new TRPCError({
+          code: 'NOT_FOUND',
+          message:
+            'Either you provided a wrong username or recovery is not enabled for this account'
+        });
+      }
+
+      const isRecoveryCodeValid = await new Argon2id().verify(
+        account.recoveryCode,
+        input.recoveryCode
+      );
+
+      if (!isRecoveryCodeValid) {
+        throw new TRPCError({
+          code: 'UNAUTHORIZED',
+          message: 'Invalid Credentials'
+        });
+      }
+
+      let resetting: 'password' | '2fa' | null = null;
+
+      if ('password' in input) {
+        const validPassword = await new Argon2id().verify(
+          account.passwordHash,
+          input.password
+        );
+
+        if (!validPassword) {
+          throw new TRPCError({
+            code: 'UNAUTHORIZED',
+            message: 'Invalid Credentials'
+          });
+        }
+
+        resetting = '2fa';
+      }
+
+      if ('twoFactorCode' in input) {
+        const secret = decodeHex(account.twoFactorSecret!);
+        const otpValid = await new TOTPController().verify(
+          input.twoFactorCode,
+          secret
+        );
+
+        if (!otpValid) {
+          throw new TRPCError({
+            code: 'UNAUTHORIZED',
+            message: 'Invalid Credentials'
+          });
+        }
+        resetting = 'password';
+      }
+
+      if (!resetting) {
+        throw new TRPCError({
+          code: 'BAD_REQUEST',
+          message: 'Password or 2FA code required'
+        });
+      }
+
+      const resetToken = nanoIdToken();
+      await storage.auth.setItem(
+        `reset-token:${resetting}:${account.publicId}`,
+        resetToken
+      );
+      setCookie(ctx.event, `reset-token_${resetting}`, resetToken, {
+        maxAge: ms('5 minutes'),
+        httpOnly: true,
+        domain: env.PRIMARY_DOMAIN,
+        sameSite: 'Lax',
+        secure: env.NODE_ENV === 'production'
+      });
+
+      // If it is a 2FA reset, return the new URI too
+      if (resetting === '2fa') {
+        const newSecret = crypto.getRandomValues(new Uint8Array(20));
+        const uri = createTOTPKeyURI('UnInbox.com', input.username, newSecret);
+        const hexSecret = encodeHex(newSecret);
+        await storage.auth.setItem(
+          `2fa-reset-secret:${account.publicId}`,
+          hexSecret
+        );
+        return { resetting, accountPublicId: account.publicId, uri };
+      } else {
+        return { resetting, accountPublicId: account.publicId };
+      }
+    }),
+  resetPassword: publicRateLimitedProcedure.completeRecovery
+    .input(
+      z.object({
+        accountPublicId: typeIdValidator('account'),
+        newPassword: strongPasswordSchema
+      })
+    )
+    .mutation(async ({ input, ctx }) => {
+      const { db, event } = ctx;
+
+      const resetToken = getCookie(event, 'reset-token_password');
+      const storedResetToken = await storage.auth.getItem<string>(
+        `reset-token:password:${input.accountPublicId}`
+      );
+
+      if (!resetToken || !storedResetToken || resetToken !== storedResetToken) {
+        throw new TRPCError({
+          code: 'BAD_REQUEST',
+          message: 'Invalid reset token'
+        });
+      }
+
+      const account = await db.query.accounts.findFirst({
+        where: eq(accounts.publicId, input.accountPublicId),
+        columns: {
+          id: true
+        }
+      });
+
+      if (!account) {
+        throw new TRPCError({
+          code: 'NOT_FOUND',
+          message: 'Account not found'
+        });
+      }
+
+      const passwordHash = await new Argon2id().hash(input.newPassword);
+      await db
+        .update(accounts)
+        .set({
+          passwordHash,
+          recoveryCode: null
+        })
+        .where(eq(accounts.id, account.id));
+
+      await storage.auth.removeItem(
+        `reset-token:password:${input.accountPublicId}`
+      );
+
+      deleteCookie(event, 'reset-token_password');
+
+      return { success: true };
+    }),
+  resetTwoFactor: publicRateLimitedProcedure.completeRecovery
+    .input(
+      z.object({
+        accountPublicId: typeIdValidator('account'),
+        twoFactorCode: z.string().min(6).max(6)
+      })
+    )
+    .mutation(async ({ input, ctx }) => {
+      const { db, event } = ctx;
+
+      const resetToken = getCookie(event, 'reset-token_2fa');
+      const storedResetToken = await storage.auth.getItem<string>(
+        `reset-token:2fa:${input.accountPublicId}`
+      );
+
+      if (!resetToken || !storedResetToken || resetToken !== storedResetToken) {
+        throw new TRPCError({
+          code: 'BAD_REQUEST',
+          message: 'Invalid reset token'
+        });
+      }
+
+      const account = await db.query.accounts.findFirst({
+        where: eq(accounts.publicId, input.accountPublicId),
+        columns: {
+          id: true
+        }
+      });
+
+      if (!account) {
+        throw new TRPCError({
+          code: 'NOT_FOUND',
+          message: 'Account not found'
+        });
+      }
+
+      const storedSecret = await storage.auth.getItem<string>(
+        `2fa-reset-secret:${input.accountPublicId}`
+      );
+      if (!storedSecret) {
+        throw new TRPCError({
+          code: 'NOT_FOUND',
+          message: '2FA Secret not found, please try again after some time'
+        });
+      }
+
+      const secret = decodeHex(storedSecret);
+      const isValid = await new TOTPController().verify(
+        input.twoFactorCode,
+        secret
+      );
+      if (!isValid) {
+        throw new TRPCError({
+          code: 'UNAUTHORIZED',
+          message: '2FA code is not valid'
+        });
+      }
+
+      await db.update(accounts).set({
+        twoFactorEnabled: true,
+        twoFactorSecret: storedSecret,
+        recoveryCode: null
+      });
+
+      await storage.auth.removeItem(
+        `2fa-reset-secret:${input.accountPublicId}`
+      );
+      await storage.auth.removeItem(`reset-token:2fa:${input.accountPublicId}`);
+      deleteCookie(event, 'reset-token_2fa');
+
+      return { success: true };
     })
 });

apps/platform/trpc/trpc.ts

@@ -55,6 +55,7 @@ const publicRateLimits = {
   signUpWithPassword: [10, '1h'],
   signInWithPassword: [20, '1h'],
   recoverAccount: [10, '1h'],
+  completeRecovery: [20, '1h'],
   validateInvite: [10, '1h']
 } satisfies Record<string, [number, Duration]>;
 

apps/web/.eslintrc.cjs

@@ -34,7 +34,8 @@ const config = {
           attributes: false
         }
       }
-    ]
+    ],
+    'react/no-children-prop': ['warn', { allowFunctions: true }]
   }
 };
 module.exports = config;

apps/web/package.json

@@ -18,9 +18,11 @@
     "@radix-ui/themes": "^3.0.2",
     "@simplewebauthn/browser": "^10.0.0",
     "@tailwindcss/typography": "^0.5.13",
+    "@tanstack/react-form": "^0.19.5",
     "@tanstack/react-query": "^4.36.1",
     "@tanstack/react-table": "^8.16.0",
     "@tanstack/react-virtual": "^3.5.0",
+    "@tanstack/zod-form-adapter": "^0.19.5",
     "@trpc/client": "10.45.2",
     "@trpc/react-query": "10.45.2",
     "@trpc/server": "10.45.2",

apps/web/src/app/(login)/_components/password-login.tsx

@@ -59,6 +59,7 @@ export default function PasswordLoginButton() {
         description: 'Redirecting you to create an organization'
       });
       router.push('/join/org');
+      return;
     }
     toast.success('Sign in successful!', {
       description: 'Redirecting you to your conversations'

apps/web/src/app/(login)/page.tsx

@@ -2,7 +2,6 @@ import { Flex, Box, Heading, Separator, Badge, Button } from '@radix-ui/themes';
 import PasskeyLoginButton from './_components/passkey-login';
 import PasswordLoginButton from './_components/password-login';
 import Link from 'next/link';
-import RecoveryButton from './_components/recovery-button';
 
 export default async function Page() {
   return (
@@ -52,7 +51,12 @@ export default async function Page() {
             </Button>
           </Flex>
         </Box>
-        <RecoveryButton />
+        <Button
+          size="3"
+          className="w-fit cursor-pointer font-semibold"
+          variant="ghost">
+          <Link href="/recovery">Recover your Account</Link>
+        </Button>
       </Box>
     </Flex>
   );