Refactor login (#748)

* wip refactor login

* wip refactor login

* Fix lint conflicts

* Complete Sign In only

* Feature complete

* Fix test

* Fix test

server/.env.example

@@ -7,7 +7,7 @@ LOGIN_TOKEN_EXPIRES_IN=15m
 REFRESH_TOKEN_SECRET=secret_refresh_token
 REFRESH_TOKEN_EXPIRES_IN=90d
 PG_DATABASE_URL=postgres://postgres:postgrespassword@localhost:5432/default?connection_limit=1
-FRONT_AUTH_CALLBACK_URL=http://localhost:3001/auth/callback
+FRONT_AUTH_CALLBACK_URL=http://localhost:3001/verify
 STORAGE_TYPE=local
 STORAGE_LOCAL_PATH=.local-storage
 

server/src/core/auth/controllers/google-auth.controller.ts

@@ -1,12 +1,4 @@
-import {
-  Controller,
-  Get,
-  InternalServerErrorException,
-  Req,
-  Res,
-  UseGuards,
-} from '@nestjs/common';
-import { AuthGuard } from '@nestjs/passport';
+import { Controller, Get, Req, Res, UseGuards } from '@nestjs/common';
 
 import { Response } from 'express';
 
@@ -14,45 +6,64 @@ import { GoogleRequest } from 'src/core/auth/strategies/google.auth.strategy';
 import { UserService } from 'src/core/user/user.service';
 import { TokenService } from 'src/core/auth/services/token.service';
 import { GoogleProviderEnabledGuard } from 'src/core/auth/guards/google-provider-enabled.guard';
+import { GoogleOauthGuard } from 'src/core/auth/guards/google-oauth.guard';
+import { WorkspaceService } from 'src/core/workspace/services/workspace.service';
+import { EnvironmentService } from 'src/integrations/environment/environment.service';
 
 @Controller('auth/google')
 export class GoogleAuthController {
   constructor(
     private readonly tokenService: TokenService,
     private readonly userService: UserService,
+    private readonly workspaceService: WorkspaceService,
+    private readonly environmentService: EnvironmentService,
   ) {}
 
   @Get()
-  @UseGuards(GoogleProviderEnabledGuard, AuthGuard('google'))
+  @UseGuards(GoogleProviderEnabledGuard, GoogleOauthGuard)
   async googleAuth() {
     // As this method is protected by Google Auth guard, it will trigger Google SSO flow
     return;
   }
 
   @Get('redirect')
-  @UseGuards(GoogleProviderEnabledGuard, AuthGuard('google'))
+  @UseGuards(GoogleProviderEnabledGuard, GoogleOauthGuard)
   async googleAuthRedirect(@Req() req: GoogleRequest, @Res() res: Response) {
-    const { firstName, lastName, email } = req.user;
+    const { firstName, lastName, email, workspaceInviteHash } = req.user;
 
-    const user = await this.userService.createUser({
-      data: {
-        email,
-        firstName: firstName ?? '',
-        lastName: lastName ?? '',
-        locale: 'en',
-        settings: {
-          create: {
-            locale: 'en',
+    let workspaceId: string | undefined = undefined;
+    if (workspaceInviteHash) {
+      const workspace = await this.workspaceService.findFirst({
+        where: {
+          inviteHash: workspaceInviteHash,
+        },
+      });
+
+      if (!workspace) {
+        return res.redirect(
+          `${this.environmentService.getFrontAuthCallbackUrl()}`,
+        );
+      }
+
+      workspaceId = workspace.id;
+    }
+
+    const user = await this.userService.createUser(
+      {
+        data: {
+          email,
+          firstName: firstName ?? '',
+          lastName: lastName ?? '',
+          locale: 'en',
+          settings: {
+            create: {
+              locale: 'en',
+            },
           },
         },
       },
-    });
-
-    if (!user) {
-      throw new InternalServerErrorException(
-        'User email domain does not match an existing workspace',
-      );
-    }
+      workspaceId,
+    );
 
     const loginToken = await this.tokenService.generateLoginToken(user.email);
 

server/src/core/auth/guards/google-oauth.guard.ts

@@ -0,0 +1,27 @@
+import { ExecutionContext, Injectable } from '@nestjs/common';
+import { AuthGuard } from '@nestjs/passport';
+
+@Injectable()
+export class GoogleOauthGuard extends AuthGuard('google') {
+  constructor() {
+    super({
+      prompt: 'select_account',
+    });
+  }
+
+  async canActivate(context: ExecutionContext) {
+    try {
+      const request = context.switchToHttp().getRequest();
+      const workspaceInviteHash = request.query.inviteHash;
+
+      if (workspaceInviteHash && typeof workspaceInviteHash === 'string') {
+        request.params.workspaceInviteHash = workspaceInviteHash;
+      }
+      const activate = (await super.canActivate(context)) as boolean;
+
+      return activate;
+    } catch (ex) {
+      throw ex;
+    }
+  }
+}

server/src/core/auth/strategies/google.auth.strategy.ts

@@ -8,9 +8,10 @@ import { EnvironmentService } from 'src/integrations/environment/environment.ser
 
 export type GoogleRequest = Request & {
   user: {
-    firstName: string | undefined | null;
-    lastName: string | undefined | null;
+    firstName?: string | null;
+    lastName?: string | null;
     email: string;
+    workspaceInviteHash?: string;
   };
 };
 
@@ -22,23 +23,39 @@ export class GoogleStrategy extends PassportStrategy(Strategy, 'google') {
       clientSecret: environmentService.getAuthGoogleClientSecret(),
       callbackURL: environmentService.getAuthGoogleCallbackUrl(),
       scope: ['email', 'profile'],
+      passReqToCallback: true,
     });
   }
 
+  authenticate(req: any, options: any) {
+    options = {
+      ...options,
+      state: JSON.stringify({
+        workspaceInviteHash: req.params.workspaceInviteHash,
+      }),
+    };
+
+    return super.authenticate(req, options);
+  }
+
   async validate(
+    request: GoogleRequest,
     accessToken: string,
     refreshToken: string,
     profile: any,
     done: VerifyCallback,
-  ): Promise<any> {
-    const { name, emails, photos } = profile;
+  ): Promise<void> {
+    const { name, emails } = profile;
+    const state =
+      typeof request.query.state === 'string'
+        ? JSON.parse(request.query.state)
+        : undefined;
+
     const user = {
       email: emails[0].value,
       firstName: name.givenName,
       lastName: name.familyName,
-      picture: photos[0].value,
-      refreshToken,
-      accessToken,
+      workspaceInviteHash: state.workspaceInviteHash,
     };
     done(null, user);
   }

server/src/core/company/company.service.ts

@@ -46,6 +46,6 @@ export class CompanyService {
       data: companies,
     });
 
-    return this.findMany({ where: { workspaceId }});
+    return this.findMany({ where: { workspaceId } });
   }
 }