fix: refactor forgot password system (#329)

* refactor forgot password system

* ready

apps/server/src/core/auth/auth.constants.ts

@@ -0,0 +1,3 @@
+export enum UserTokenType {
+  FORGOT_PASSWORD = 'forgot-password',
+}

apps/server/src/core/auth/auth.controller.ts

@@ -10,7 +10,6 @@ import {
 } from '@nestjs/common';
 import { LoginDto } from './dto/login.dto';
 import { AuthService } from './services/auth.service';
-import { CreateUserDto } from './dto/create-user.dto';
 import { SetupGuard } from './guards/setup.guard';
 import { EnvironmentService } from '../../integrations/environment/environment.service';
 import { CreateAdminUserDto } from './dto/create-admin-user.dto';
@@ -20,6 +19,8 @@ import { User, Workspace } from '@docmost/db/types/entity.types';
 import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
 import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
 import { ForgotPasswordDto } from './dto/forgot-password.dto';
+import { PasswordResetDto } from './dto/password-reset.dto';
+import { VerifyUserTokenDto } from './dto/verify-user-token.dto';
 
 @Controller('auth')
 export class AuthController {
@@ -34,18 +35,6 @@ export class AuthController {
     return this.authService.login(loginInput, req.raw.workspaceId);
   }
 
-  @HttpCode(HttpStatus.OK)
-  @Post('forgot-password')
-  async forgotPassword(
-    @Req() req,
-    @Body() forgotPasswordDto: ForgotPasswordDto,
-  ) {
-    return this.authService.forgotPassword(
-      forgotPasswordDto,
-      req.raw.workspaceId,
-    );
-  }
-
   /* @HttpCode(HttpStatus.OK)
   @Post('register')
   async register(@Req() req, @Body() createUserDto: CreateUserDto) {
@@ -74,4 +63,31 @@ export class AuthController {
   ) {
     return this.authService.changePassword(dto, user.id, workspace.id);
   }
+
+  @HttpCode(HttpStatus.OK)
+  @Post('forgot-password')
+  async forgotPassword(
+    @Body() forgotPasswordDto: ForgotPasswordDto,
+    @AuthWorkspace() workspace: Workspace,
+  ) {
+    return this.authService.forgotPassword(forgotPasswordDto, workspace.id);
+  }
+
+  @HttpCode(HttpStatus.OK)
+  @Post('password-reset')
+  async passwordReset(
+    @Body() passwordResetDto: PasswordResetDto,
+    @AuthWorkspace() workspace: Workspace,
+  ) {
+    return this.authService.passwordReset(passwordResetDto, workspace.id);
+  }
+
+  @HttpCode(HttpStatus.OK)
+  @Post('verify-token')
+  async verifyResetToken(
+    @Body() verifyUserTokenDto: VerifyUserTokenDto,
+    @AuthWorkspace() workspace: Workspace,
+  ) {
+    return this.authService.verifyUserToken(verifyUserTokenDto, workspace.id);
+  }
 }

apps/server/src/core/auth/dto/forgot-password.dto.ts

@@ -1,16 +1,7 @@
-import { IsEmail, IsNotEmpty, IsOptional, IsString, MinLength } from 'class-validator';
+import { IsEmail, IsNotEmpty } from 'class-validator';
 
 export class ForgotPasswordDto {
   @IsNotEmpty()
   @IsEmail()
   email: string;
-
-  @IsOptional()
-  @IsString()
-  token: string;
-
-  @IsOptional()
-  @IsString()
-  @MinLength(8)
-  newPassword: string;
 }

apps/server/src/core/auth/dto/password-reset.dto.ts

@@ -0,0 +1,10 @@
+import { IsString, MinLength } from 'class-validator';
+
+export class PasswordResetDto {
+  @IsString()
+  token: string;
+
+  @IsString()
+  @MinLength(8)
+  newPassword: string;
+}

apps/server/src/core/auth/dto/verify-user-token.dto.ts

@@ -0,0 +1,9 @@
+import { IsString, MinLength } from 'class-validator';
+
+export class VerifyUserTokenDto {
+  @IsString()
+  token: string;
+
+  @IsString()
+  type: string;
+}

apps/server/src/core/auth/services/auth.service.ts

@@ -11,13 +11,25 @@ import { TokensDto } from '../dto/tokens.dto';
 import { SignupService } from './signup.service';
 import { CreateAdminUserDto } from '../dto/create-admin-user.dto';
 import { UserRepo } from '@docmost/db/repos/user/user.repo';
-import { comparePasswordHash, hashPassword, nanoIdGen } from '../../../common/helpers';
+import {
+  comparePasswordHash,
+  hashPassword,
+  nanoIdGen,
+} from '../../../common/helpers';
 import { ChangePasswordDto } from '../dto/change-password.dto';
 import { MailService } from '../../../integrations/mail/mail.service';
 import ChangePasswordEmail from '@docmost/transactional/emails/change-password-email';
 import { ForgotPasswordDto } from '../dto/forgot-password.dto';
 import ForgotPasswordEmail from '@docmost/transactional/emails/forgot-password-email';
-import { UserTokensRepo } from '@docmost/db/repos/user-tokens/user-tokens.repo';
+import { UserTokenRepo } from '@docmost/db/repos/user-token/user-token.repo';
+import { PasswordResetDto } from '../dto/password-reset.dto';
+import { UserToken } from '@docmost/db/types/entity.types';
+import { UserTokenType } from '../auth.constants';
+import { KyselyDB } from '@docmost/db/types/kysely.types';
+import { InjectKysely } from 'nestjs-kysely';
+import { executeTx } from '@docmost/db/utils';
+import { VerifyUserTokenDto } from '../dto/verify-user-token.dto';
+import { EnvironmentService } from 'src/integrations/environment/environment.service';
 
 @Injectable()
 export class AuthService {
@@ -25,8 +37,10 @@ export class AuthService {
     private signupService: SignupService,
     private tokenService: TokenService,
     private userRepo: UserRepo,
-    private userTokensRepo: UserTokensRepo,
+    private userTokenRepo: UserTokenRepo,
     private mailService: MailService,
+    private environmentService: EnvironmentService,
+    @InjectKysely() private readonly db: KyselyDB,
   ) {}
 
   async login(loginDto: LoginDto, workspaceId: string) {
@@ -50,94 +64,6 @@ export class AuthService {
     return { tokens };
   }
 
-  async forgotPassword(
-    forgotPasswordDto: ForgotPasswordDto,
-    workspaceId: string,
-  ) {
-    const user = await this.userRepo.findByEmail(
-      forgotPasswordDto.email,
-      workspaceId,
-      true,
-    );
-    if (!user) {
-      return;
-    }
-
-    if (
-      forgotPasswordDto.token == null ||
-      forgotPasswordDto.newPassword == null
-    ) {
-      // Generate 5-character user token
-      const code = nanoIdGen(5).toUpperCase();
-      const hashedToken = await hashPassword(code);
-      await this.userTokensRepo.insertUserToken({
-        token: hashedToken,
-        user_id: user.id,
-        workspace_id: user.workspaceId,
-        expires_at: new Date(new Date().getTime() + 3_600_000), // should expires in 1 hour
-        type: "forgot-password",
-      });
-
-      const emailTemplate = ForgotPasswordEmail({
-        username: user.name,
-        code: code,
-      });
-      await this.mailService.sendToQueue({
-        to: user.email,
-        subject: 'Reset your password',
-        template: emailTemplate,
-      });
-
-      return;
-    }
-
-    // Get all user tokens that are not expired
-    const userTokens = await this.userTokensRepo.findByUserId(
-      user.id,
-      user.workspaceId,
-      "forgot-password"
-    );
-    // Limit to the last 3 token, so we have a total time window of 15 minutes
-    const validUserTokens = userTokens
-      .filter((token) => token.expires_at > new Date() && token.used_at == null)
-      .slice(0, 3);
-
-    for (const token of validUserTokens) {
-      const validated = await comparePasswordHash(
-        forgotPasswordDto.token,
-        token.token,
-      );
-      if (validated) {
-        await Promise.all([
-          this.userTokensRepo.deleteUserToken(user.id, user.workspaceId, "forgot-password"),
-          this.userTokensRepo.deleteExpiredUserTokens(),
-        ]);
-
-        const newPasswordHash = await hashPassword(
-          forgotPasswordDto.newPassword,
-        );
-        await this.userRepo.updateUser(
-          {
-            password: newPasswordHash,
-          },
-          user.id,
-          workspaceId,
-        );
-
-        const emailTemplate = ChangePasswordEmail({ username: user.name });
-        await this.mailService.sendToQueue({
-          to: user.email,
-          subject: 'Your password has been changed',
-          template: emailTemplate,
-        });
-
-        return;
-      }
-    }
-
-    throw new BadRequestException('Incorrect code');
-  }
-
   async register(createUserDto: CreateUserDto, workspaceId: string) {
     const user = await this.signupService.signup(createUserDto, workspaceId);
 
@@ -192,4 +118,108 @@ export class AuthService {
       template: emailTemplate,
     });
   }
+
+  async forgotPassword(
+    forgotPasswordDto: ForgotPasswordDto,
+    workspaceId: string,
+  ): Promise<void> {
+    const user = await this.userRepo.findByEmail(
+      forgotPasswordDto.email,
+      workspaceId,
+    );
+
+    if (!user) {
+      return;
+    }
+
+    const token = nanoIdGen(16);
+    const resetLink = `${this.environmentService.getAppUrl()}/password-reset?token=${token}`;
+
+    await this.userTokenRepo.insertUserToken({
+      token: token,
+      userId: user.id,
+      workspaceId: user.workspaceId,
+      expiresAt: new Date(new Date().getTime() + 60 * 60 * 1000), // 1 hour
+      type: UserTokenType.FORGOT_PASSWORD,
+    });
+
+    const emailTemplate = ForgotPasswordEmail({
+      username: user.name,
+      resetLink: resetLink,
+    });
+
+    await this.mailService.sendToQueue({
+      to: user.email,
+      subject: 'Reset your password',
+      template: emailTemplate,
+    });
+  }
+
+  async passwordReset(passwordResetDto: PasswordResetDto, workspaceId: string) {
+    const userToken = await this.userTokenRepo.findById(
+      passwordResetDto.token,
+      workspaceId,
+    );
+
+    if (
+      !userToken ||
+      userToken.type !== UserTokenType.FORGOT_PASSWORD ||
+      userToken.expiresAt < new Date()
+    ) {
+      throw new BadRequestException('Invalid or expired token');
+    }
+
+    const user = await this.userRepo.findById(userToken.userId, workspaceId);
+    if (!user) {
+      throw new NotFoundException('User not found');
+    }
+
+    const newPasswordHash = await hashPassword(passwordResetDto.newPassword);
+
+    await executeTx(this.db, async (trx) => {
+      await this.userRepo.updateUser(
+        {
+          password: newPasswordHash,
+        },
+        user.id,
+        workspaceId,
+        trx,
+      );
+
+      trx
+        .deleteFrom('userTokens')
+        .where('userId', '=', user.id)
+        .where('type', '=', UserTokenType.FORGOT_PASSWORD)
+        .execute();
+    });
+
+    const emailTemplate = ChangePasswordEmail({ username: user.name });
+    await this.mailService.sendToQueue({
+      to: user.email,
+      subject: 'Your password has been changed',
+      template: emailTemplate,
+    });
+
+    const tokens: TokensDto = await this.tokenService.generateTokens(user);
+
+    return { tokens };
+  }
+
+  async verifyUserToken(
+    userTokenDto: VerifyUserTokenDto,
+    workspaceId: string,
+  ): Promise<void> {
+    const userToken = await this.userTokenRepo.findById(
+      userTokenDto.token,
+      workspaceId,
+    );
+
+    if (
+      !userToken ||
+      userToken.type !== userTokenDto.type ||
+      userToken.expiresAt < new Date()
+    ) {
+      throw new BadRequestException('Invalid or expired token');
+    }
+  }
 }