Update token service and guards

Author: thomtrp

packages/twenty-server/src/engine/core-modules/auth/auth.resolver.ts

@@ -1,10 +1,4 @@
-import {
-  BadRequestException,
-  ForbiddenException,
-  InternalServerErrorException,
-  NotFoundException,
-  UseGuards,
-} from '@nestjs/common';
+import { UseGuards } from '@nestjs/common';
 import { Args, Mutation, Query, Resolver } from '@nestjs/graphql';
 import { InjectRepository } from '@nestjs/typeorm';
 
@@ -25,15 +19,13 @@ import { UpdatePasswordViaResetTokenInput } from 'src/engine/core-modules/auth/d
 import { ValidatePasswordResetToken } from 'src/engine/core-modules/auth/dto/validate-password-reset-token.entity';
 import { ValidatePasswordResetTokenInput } from 'src/engine/core-modules/auth/dto/validate-password-reset-token.input';
 import { authGraphqlApiExceptionHandler } from 'src/engine/core-modules/auth/utils/auth-graphql-api-exception-handler.util';
-import { NotFoundError } from 'src/engine/core-modules/graphql/utils/graphql-errors.util';
 import { UserService } from 'src/engine/core-modules/user/services/user.service';
 import { User } from 'src/engine/core-modules/user/user.entity';
 import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
 import { AuthUser } from 'src/engine/decorators/auth/auth-user.decorator';
 import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
 import { JwtAuthGuard } from 'src/engine/guards/jwt.auth.guard';
 import { CaptchaGuard } from 'src/engine/integrations/captcha/captcha.guard';
-import { assert } from 'src/utils/assert';
 
 import { ChallengeInput } from './dto/challenge.input';
 import { ImpersonateInput } from './dto/impersonate.input';
@@ -91,81 +83,99 @@ export class AuthResolver {
   @Query(() => Workspace)
   async findWorkspaceFromInviteHash(
     @Args() workspaceInviteHashValidInput: WorkspaceInviteHashValidInput,
-  ): Promise<Workspace> {
-    const workspace = await this.workspaceRepository.findOneBy({
-      inviteHash: workspaceInviteHashValidInput.inviteHash,
-    });
-
-    if (!workspace) {
-      throw new NotFoundError('Workspace does not exist');
+  ): Promise<Workspace | void> {
+    try {
+      return await this.authService.findWorkspaceFromInviteHash(
+        workspaceInviteHashValidInput.inviteHash,
+      );
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
     }
-
-    return workspace;
   }
 
   @UseGuards(CaptchaGuard)
   @Mutation(() => LoginToken)
-  async challenge(@Args() challengeInput: ChallengeInput): Promise<LoginToken> {
-    const user = await this.authService.challenge(challengeInput);
-    const loginToken = await this.tokenService.generateLoginToken(user.email);
+  async challenge(
+    @Args() challengeInput: ChallengeInput,
+  ): Promise<LoginToken | void> {
+    try {
+      const user = await this.authService.challenge(challengeInput);
+      const loginToken = await this.tokenService.generateLoginToken(user.email);
 
-    return { loginToken };
+      return { loginToken };
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 
   @UseGuards(CaptchaGuard)
   @Mutation(() => LoginToken)
-  async signUp(@Args() signUpInput: SignUpInput): Promise<LoginToken> {
-    const user = await this.authService.signInUp({
-      ...signUpInput,
-      fromSSO: false,
-    });
+  async signUp(@Args() signUpInput: SignUpInput): Promise<LoginToken | void> {
+    try {
+      const user = await this.authService.signInUp({
+        ...signUpInput,
+        fromSSO: false,
+      });
 
-    const loginToken = await this.tokenService.generateLoginToken(user.email);
+      const loginToken = await this.tokenService.generateLoginToken(user.email);
 
-    return { loginToken };
+      return { loginToken };
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 
   @Mutation(() => ExchangeAuthCode)
   async exchangeAuthorizationCode(
     @Args() exchangeAuthCodeInput: ExchangeAuthCodeInput,
   ) {
-    const tokens = await this.tokenService.verifyAuthorizationCode(
-      exchangeAuthCodeInput,
-    );
+    try {
+      const tokens = await this.tokenService.verifyAuthorizationCode(
+        exchangeAuthCodeInput,
+      );
 
-    return tokens;
+      return tokens;
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 
   @Mutation(() => TransientToken)
   @UseGuards(JwtAuthGuard)
   async generateTransientToken(
     @AuthUser() user: User,
   ): Promise<TransientToken | void> {
-    const workspaceMember = await this.userService.loadWorkspaceMember(user);
+    try {
+      const workspaceMember = await this.userService.loadWorkspaceMember(user);
+
+      if (!workspaceMember) {
+        return;
+      }
+      const transientToken = await this.tokenService.generateTransientToken(
+        workspaceMember.id,
+        user.id,
+        user.defaultWorkspace.id,
+      );
 
-    if (!workspaceMember) {
-      return;
+      return { transientToken };
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
     }
-    const transientToken = await this.tokenService.generateTransientToken(
-      workspaceMember.id,
-      user.id,
-      user.defaultWorkspace.id,
-    );
-
-    return { transientToken };
   }
 
   @Mutation(() => Verify)
-  async verify(@Args() verifyInput: VerifyInput): Promise<Verify> {
-    const email = await this.tokenService.verifyLoginToken(
-      verifyInput.loginToken,
-    );
-
-    assert(email, 'Invalid token', ForbiddenException);
+  async verify(@Args() verifyInput: VerifyInput): Promise<Verify | void> {
+    try {
+      const email = await this.tokenService.verifyLoginToken(
+        verifyInput.loginToken,
+      );
 
-    const result = await this.authService.verify(email);
+      const result = await this.authService.verify(email);
 
-    return result;
+      return result;
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 
   @Mutation(() => AuthorizeApp)
@@ -191,35 +201,40 @@ export class AuthResolver {
   async generateJWT(
     @AuthUser() user: User,
     @Args() args: GenerateJwtInput,
-  ): Promise<AuthTokens> {
-    const token = await this.tokenService.generateSwitchWorkspaceToken(
-      user,
-      args.workspaceId,
-    );
+  ): Promise<AuthTokens | void> {
+    try {
+      const token = await this.tokenService.generateSwitchWorkspaceToken(
+        user,
+        args.workspaceId,
+      );
 
-    return token;
+      return token;
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 
   @Mutation(() => AuthTokens)
-  async renewToken(@Args() args: AppTokenInput): Promise<AuthTokens> {
-    if (!args.appToken) {
-      throw new BadRequestException('Refresh token is mendatory');
-    }
-
-    const tokens = await this.tokenService.generateTokensFromRefreshToken(
-      args.appToken,
-    );
+  async renewToken(@Args() args: AppTokenInput): Promise<AuthTokens | void> {
+    try {
+      const tokens = await this.tokenService.generateTokensFromRefreshToken(
+        args.appToken,
+      );
 
-    return { tokens: tokens };
+      return { tokens: tokens };
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 
   @UseGuards(JwtAuthGuard)
   @Mutation(() => Verify)
   async impersonate(
     @Args() impersonateInput: ImpersonateInput,
+    @AuthUser() user: User,
   ): Promise<Verify | void> {
     try {
-      return await this.authService.impersonate(impersonateInput.userId);
+      return await this.authService.impersonate(impersonateInput.userId, user);
     } catch (error) {
       authGraphqlApiExceptionHandler(error);
     }
@@ -231,53 +246,62 @@ export class AuthResolver {
     @Args() args: ApiKeyTokenInput,
     @AuthWorkspace() { id: workspaceId }: Workspace,
   ): Promise<ApiKeyToken | undefined> {
-    return await this.tokenService.generateApiKeyToken(
-      workspaceId,
-      args.apiKeyId,
-      args.expiresAt,
-    );
+    try {
+      return await this.tokenService.generateApiKeyToken(
+        workspaceId,
+        args.apiKeyId,
+        args.expiresAt,
+      );
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 
   @Mutation(() => EmailPasswordResetLink)
   async emailPasswordResetLink(
     @Args() emailPasswordResetInput: EmailPasswordResetLinkInput,
-  ): Promise<EmailPasswordResetLink> {
-    const resetToken = await this.tokenService.generatePasswordResetToken(
-      emailPasswordResetInput.email,
-    );
-
-    return await this.tokenService.sendEmailPasswordResetLink(
-      resetToken,
-      emailPasswordResetInput.email,
-    );
+  ): Promise<EmailPasswordResetLink | void> {
+    try {
+      const resetToken = await this.tokenService.generatePasswordResetToken(
+        emailPasswordResetInput.email,
+      );
+
+      return await this.tokenService.sendEmailPasswordResetLink(
+        resetToken,
+        emailPasswordResetInput.email,
+      );
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 
   @Mutation(() => InvalidatePassword)
   async updatePasswordViaResetToken(
     @Args() args: UpdatePasswordViaResetTokenInput,
-  ): Promise<InvalidatePassword> {
-    const { id } = await this.tokenService.validatePasswordResetToken(
-      args.passwordResetToken,
-    );
-
-    assert(id, 'User not found', NotFoundException);
-
-    const { success } = await this.authService.updatePassword(
-      id,
-      args.newPassword,
-    );
+  ): Promise<InvalidatePassword | void> {
+    try {
+      const { id } = await this.tokenService.validatePasswordResetToken(
+        args.passwordResetToken,
+      );
 
-    assert(success, 'Password update failed', InternalServerErrorException);
+      await this.authService.updatePassword(id, args.newPassword);
 
-    return await this.tokenService.invalidatePasswordResetToken(id);
+      return await this.tokenService.invalidatePasswordResetToken(id);
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 
   @Query(() => ValidatePasswordResetToken)
   async validatePasswordResetToken(
     @Args() args: ValidatePasswordResetTokenInput,
-  ): Promise<ValidatePasswordResetToken> {
-    return this.tokenService.validatePasswordResetToken(
-      args.passwordResetToken,
-    );
+  ): Promise<ValidatePasswordResetToken | void> {
+    try {
+      return this.tokenService.validatePasswordResetToken(
+        args.passwordResetToken,
+      );
+    } catch (error) {
+      authGraphqlApiExceptionHandler(error);
+    }
   }
 }

packages/twenty-server/src/engine/core-modules/auth/controllers/google-apis-auth.controller.ts

@@ -1,14 +1,11 @@
-import {
-  Controller,
-  Get,
-  Req,
-  Res,
-  UnauthorizedException,
-  UseGuards,
-} from '@nestjs/common';
+import { Controller, Get, Req, Res, UseGuards } from '@nestjs/common';
 
 import { Response } from 'express';
 
+import {
+  AuthException,
+  AuthExceptionCode,
+} from 'src/engine/core-modules/auth/auth.exception';
 import { GoogleAPIsOauthExchangeCodeForTokenGuard } from 'src/engine/core-modules/auth/guards/google-apis-oauth-exchange-code-for-token.guard';
 import { GoogleAPIsOauthRequestCodeGuard } from 'src/engine/core-modules/auth/guards/google-apis-oauth-request-code.guard';
 import { GoogleAPIsService } from 'src/engine/core-modules/auth/services/google-apis.service';
@@ -59,13 +56,17 @@ export class GoogleAPIsAuthController {
     const demoWorkspaceIds = this.environmentService.get('DEMO_WORKSPACE_IDS');
 
     if (demoWorkspaceIds.includes(workspaceId)) {
-      throw new UnauthorizedException(
+      throw new AuthException(
         'Cannot connect Google account to demo workspace',
+        AuthExceptionCode.FORBIDDEN_EXCEPTION,
       );
     }
 
     if (!workspaceId) {
-      throw new Error('Workspace not found');
+      throw new AuthException(
+        'Workspace not found',
+        AuthExceptionCode.INVALID_INPUT,
+      );
     }
 
     const handle = emails[0].value;

packages/twenty-server/src/engine/core-modules/auth/guards/google-apis-oauth-exchange-code-for-token.guard.ts

@@ -1,13 +1,13 @@
-import {
-  ExecutionContext,
-  Injectable,
-  NotFoundException,
-} from '@nestjs/common';
+import { ExecutionContext, Injectable } from '@nestjs/common';
 import { AuthGuard } from '@nestjs/passport';
 import { InjectRepository } from '@nestjs/typeorm';
 
 import { Repository } from 'typeorm';
 
+import {
+  AuthException,
+  AuthExceptionCode,
+} from 'src/engine/core-modules/auth/auth.exception';
 import { TokenService } from 'src/engine/core-modules/auth/services/token.service';
 import {
   GoogleAPIScopeConfig,
@@ -39,7 +39,10 @@ export class GoogleAPIsOauthExchangeCodeForTokenGuard extends AuthGuard(
       !this.environmentService.get('MESSAGING_PROVIDER_GMAIL_ENABLED') &&
       !this.environmentService.get('CALENDAR_PROVIDER_GOOGLE_ENABLED')
     ) {
-      throw new NotFoundException('Google apis auth is not enabled');
+      throw new AuthException(
+        'Google apis auth is not enabled',
+        AuthExceptionCode.FORBIDDEN_EXCEPTION,
+      );
     }
 
     const { workspaceId } = await this.tokenService.verifyTransientToken(