2fa for oauth, rework 2fa system

src/app/pages/login/2fa.vue

@@ -35,18 +35,20 @@
 <script setup lang="ts">
 const toast = useToast();
 const { t } = useI18n();
-const loginStore = useLoginStore();
 
 const authenticating = ref(false);
 const totp = ref<string>('');
 
-if (!loginStore.hasPendingLogin) {
-  await navigateTo('/login');
-}
+const { error } = await useFetch('/api/auth/pending');
+watchEffect(() => {
+  if (error.value) {
+    navigateTo('/login');
+  }
+});
 
 const _submit = useSubmit(
   (data) =>
-    $fetch('/api/auth/password', {
+    $fetch('/api/auth/verify-2fa', {
       method: 'post',
       body: data,
     }),
@@ -54,8 +56,8 @@ const _submit = useSubmit(
     revert: async (success, data) => {
       if (success) {
         if (data?.status === 'success') {
-          loginStore.clearPendingLogin();
           await navigateTo('/');
+          return;
         } else if (data?.status === 'INVALID_TOTP_CODE') {
           authenticating.value = false;
           totp.value = '';
@@ -65,15 +67,6 @@ const _submit = useSubmit(
             type: 'error',
           });
           return;
-        } else if (data?.status === 'TOTP_REQUIRED') {
-          authenticating.value = false;
-          totp.value = '';
-          toast.showToast({
-            title: t('general.2fa'),
-            message: t('login.2faRequired'),
-            type: 'error',
-          });
-          return;
         }
       }
       authenticating.value = false;
@@ -83,24 +76,27 @@ const _submit = useSubmit(
 );
 
 async function submit() {
-  const challenge = loginStore.pendingChallenge;
-
-  if (!challenge || !totp.value || authenticating.value) return;
+  if (!totp.value || authenticating.value) return;
 
   authenticating.value = true;
+  return _submit({ totpCode: totp.value });
+}
 
-  if (challenge.type === 'password') {
-    return _submit({
-      username: challenge.username,
-      password: challenge.password,
-      remember: challenge.remember,
-      totpCode: totp.value,
-    });
+const _cancel = useSubmit(
+  (data) =>
+    $fetch('/api/auth/cancel', {
+      method: 'post',
+      body: data,
+    }),
+  {
+    revert: async () => {
+      await navigateTo('/login');
+    },
+    noSuccessToast: true,
   }
-}
+);
 
 async function cancel() {
-  loginStore.clearPendingLogin();
-  await navigateTo('/login');
+  return _cancel({});
 }
 </script>

src/app/pages/login/index.vue

@@ -69,7 +69,6 @@
 <script setup lang="ts">
 const toast = useToast();
 const { t } = useI18n();
-const loginStore = useLoginStore();
 
 const authenticating = ref(false);
 const remember = ref(false);
@@ -88,15 +87,9 @@ const _submit = useSubmit(
     revert: async (success, data) => {
       if (success) {
         if (data?.status === 'success') {
-          loginStore.clearPendingLogin();
           await navigateTo('/');
         } else if (data?.status === 'TOTP_REQUIRED') {
           authenticating.value = false;
-          loginStore.setPendingPasswordLogin(
-            username.value,
-            password.value,
-            remember.value
-          );
           await navigateTo('/login/2fa');
           return;
         } else if (data?.status === 'INVALID_TOTP_CODE') {

src/app/stores/login.ts

@@ -1,36 +0,0 @@
-type PendingLoginChallenge = {
-  type: 'password';
-  username: string;
-  password: string;
-  remember: boolean;
-};
-
-export const useLoginStore = defineStore('Login', () => {
-  const pendingChallenge = ref<PendingLoginChallenge | null>(null);
-
-  const hasPendingLogin = computed(() => pendingChallenge.value !== null);
-
-  function setPendingPasswordLogin(
-    username: string,
-    password: string,
-    remember: boolean
-  ) {
-    pendingChallenge.value = {
-      type: 'password',
-      username,
-      password,
-      remember,
-    };
-  }
-
-  function clearPendingLogin() {
-    pendingChallenge.value = null;
-  }
-
-  return {
-    pendingChallenge,
-    hasPendingLogin,
-    setPendingPasswordLogin,
-    clearPendingLogin,
-  };
-});

src/server/api/auth/[provider]/callback.get.ts

@@ -34,6 +34,18 @@ export default defineEventHandler(async (event) => {
 
   if (!result.success) {
     switch (result.error) {
+      case 'TOTP_REQUIRED':
+        await session.update({
+          pendingLogin: {
+            type: 'oauth',
+            userId: result.userId,
+            remember: false,
+          },
+          oauth_nonce: undefined,
+          oauth_state: undefined,
+          oauth_verifier: undefined,
+        });
+        return sendRedirect(event, '/login/2fa');
       case 'USER_DISABLED':
         throw createError({
           statusCode: 401,
@@ -56,7 +68,7 @@ export default defineEventHandler(async (event) => {
           statusMessage: 'Unexpected error',
         });
     }
-    assertUnreachable(result.error);
+    assertUnreachable(result);
   }
 
   // Create session

src/server/api/auth/cancel.post.ts

@@ -0,0 +1,10 @@
+export default defineEventHandler(async (event) => {
+  const session = await useWGSession(event, false);
+  await session.update({
+    pendingLogin: undefined,
+    oauth_nonce: undefined,
+    oauth_state: undefined,
+    oauth_verifier: undefined,
+  });
+  return { success: true as const };
+});

src/server/api/auth/password.post.ts

@@ -8,12 +8,14 @@ export default defineEventHandler(async (event) => {
     });
   }
 
-  const { username, password, remember, totpCode } = await readValidatedBody(
+  const { username, password, remember } = await readValidatedBody(
     event,
     validateZod(UserLoginSchema, event)
   );
 
-  const result = await Database.users.login(username, password, totpCode);
+  const result = await Database.users.login(username, password);
+
+  const session = await useWGSession(event, remember);
 
   // TODO: add localization support
 
@@ -25,6 +27,13 @@ export default defineEventHandler(async (event) => {
           statusMessage: 'Invalid username or password',
         });
       case 'TOTP_REQUIRED':
+        await session.update({
+          pendingLogin: {
+            type: 'password',
+            userId: result.userId,
+            remember,
+          },
+        });
         return { status: 'TOTP_REQUIRED' as const };
       case 'INVALID_TOTP_CODE':
         return { status: 'INVALID_TOTP_CODE' as const };
@@ -39,13 +48,11 @@ export default defineEventHandler(async (event) => {
           statusMessage: 'Unexpected error',
         });
     }
-    assertUnreachable(result.error);
+    assertUnreachable(result);
   }
 
   const user = result.user;
 
-  const session = await useWGSession(event, remember);
-
   const data = await session.update({
     userId: user.id,
   });

src/server/api/auth/pending.get.ts

@@ -0,0 +1,14 @@
+export default defineEventHandler(async (event) => {
+  const session = await useWGSession(event);
+
+  if (!session.data.pendingLogin) {
+    throw createError({
+      statusCode: 401,
+      statusMessage: 'No pending authentication',
+    });
+  }
+
+  return {
+    type: session.data.pendingLogin.type,
+  };
+});

src/server/api/auth/verify-2fa.post.ts

@@ -0,0 +1,40 @@
+import { z } from 'zod';
+
+const Verify2faSchema = z.object({
+  totpCode: z.string().min(6).max(6),
+});
+
+export default defineEventHandler(async (event) => {
+  const { totpCode } = await readValidatedBody(
+    event,
+    validateZod(Verify2faSchema, event)
+  );
+  const session = await useWGSession(event, false);
+
+  const pendingLogin = session.data.pendingLogin;
+  if (!pendingLogin) {
+    throw createError({
+      statusCode: 401,
+      statusMessage: 'No pending authentication',
+    });
+  }
+
+  const isValid = await Database.users.validateTotpCode(
+    pendingLogin.userId,
+    totpCode
+  );
+
+  if (!isValid) {
+    return { status: 'INVALID_TOTP_CODE' as const };
+  }
+
+  await session.update({
+    userId: pendingLogin.userId,
+    pendingLogin: undefined,
+    oauth_nonce: undefined,
+    oauth_state: undefined,
+    oauth_verifier: undefined,
+  });
+
+  return { status: 'success' as const };
+});

src/server/database/repositories/user/service.ts

@@ -13,10 +13,14 @@ type LoginResult =
       success: false;
       error:
         | 'INCORRECT_CREDENTIALS'
-        | 'TOTP_REQUIRED'
         | 'USER_DISABLED'
         | 'INVALID_TOTP_CODE'
         | 'UNEXPECTED_ERROR';
+    }
+  | {
+      success: false;
+      error: 'TOTP_REQUIRED';
+      userId: ID;
     };
 
 type LoginWithOAuthResult =
@@ -31,6 +35,11 @@ type LoginWithOAuthResult =
         | 'USER_ALREADY_LINKED'
         | 'UNEXPECTED_ERROR'
         | 'AUTO_REGISTER_DISABLED';
+    }
+  | {
+      success: false;
+      error: 'TOTP_REQUIRED';
+      userId: ID;
     };
 
 function createPreparedStatement(db: DBType) {
@@ -172,7 +181,7 @@ export class UserService {
     return this.#statements.updateKey.execute({ id, key });
   }
 
-  login(username: string, password: string, code: string | undefined) {
+  login(username: string, password: string) {
     return this.#db.transaction(async (tx): Promise<LoginResult> => {
       const txUser = await tx.query.user
         .findFirst({ where: eq(user.username, username) })
@@ -187,19 +196,11 @@ export class UserService {
       }
 
       if (txUser.totpVerified) {
-        if (!code) {
-          return { success: false, error: 'TOTP_REQUIRED' };
-        } else {
-          const totpKey = txUser.totpKey;
-          if (!totpKey) {
-            return { success: false, error: 'UNEXPECTED_ERROR' };
-          }
-
-          const totp = this.#createTotp({ username: txUser.username, totpKey });
-          if (totp.validate({ token: code, window: 1 }) === null) {
-            return { success: false, error: 'INVALID_TOTP_CODE' };
-          }
-        }
+        return {
+          success: false,
+          error: 'TOTP_REQUIRED',
+          userId: txUser.id,
+        };
       }
 
       if (!txUser.enabled) {
@@ -265,6 +266,22 @@ export class UserService {
     });
   }
 
+  async validateTotpCode(id: ID, code: string) {
+    const txUser = await this.#db.query.user
+      .findFirst({ where: eq(user.id, id) })
+      .execute();
+
+    if (!txUser || !txUser.totpVerified || !txUser.totpKey) {
+      return false;
+    }
+
+    const totp = this.#createTotp({
+      username: txUser.username,
+      totpKey: txUser.totpKey,
+    });
+    return totp.validate({ token: code, window: 1 }) !== null;
+  }
+
   /**
    * Login or register user with OAuth provider.
    * If user with the same email already exists, link account with OAuth provider.
@@ -288,6 +305,13 @@ export class UserService {
         .execute();
 
       if (userById) {
+        if (userById.totpVerified) {
+          return {
+            success: false,
+            error: 'TOTP_REQUIRED',
+            userId: userById.id,
+          };
+        }
         if (!userById.enabled) {
           return { success: false, error: 'USER_DISABLED' };
         }
@@ -317,6 +341,14 @@ export class UserService {
           .where(eq(user.id, userByEmail.id))
           .execute();
 
+        if (userByEmail.totpVerified) {
+          return {
+            success: false,
+            error: 'TOTP_REQUIRED',
+            userId: userByEmail.id,
+          };
+        }
+
         // TODO: return updated user
         return { success: true, user: userByEmail };
       }

src/server/database/repositories/user/types.ts

@@ -25,7 +25,6 @@ export const UserLoginSchema = z.object({
   username: username,
   password: password,
   remember: remember,
-  totpCode: totpCode.optional(),
 });
 
 export const UserSetupSchema = z

src/server/utils/session.ts

@@ -3,6 +3,12 @@ import type { UserType } from '#db/repositories/user/types';
 
 export type WGSession = Partial<{
   userId: ID;
+  // TODO: add pending login expiration
+  pendingLogin: {
+    type: 'password' | 'oauth';
+    userId: ID;
+    remember: boolean;
+  };
   oauth_verifier: string;
   oauth_nonce: string;
   oauth_state: string;