refactor: extract totp operations in seperate service

backend/src/auth/auth.controller.ts

@@ -11,6 +11,7 @@ import { Throttle } from "@nestjs/throttler";
 import { User } from "@prisma/client";
 import { ConfigService } from "src/config/config.service";
 import { AuthService } from "./auth.service";
+import { AuthTotpService } from "./authTotp.service";
 import { GetUser } from "./decorator/getUser.decorator";
 import { AuthRegisterDTO } from "./dto/authRegister.dto";
 import { AuthSignInDTO } from "./dto/authSignIn.dto";
@@ -25,6 +26,7 @@ import { JwtGuard } from "./guard/jwt.guard";
 export class AuthController {
   constructor(
     private authService: AuthService,
+    private authTotpService: AuthTotpService,
     private config: ConfigService
   ) {}
 
@@ -47,7 +49,7 @@ export class AuthController {
   @Post("signIn/totp")
   @HttpCode(200)
   signInTotp(@Body() dto: AuthSignInTotpDTO) {
-    return this.authService.signInTotp(dto);
+    return this.authTotpService.signInTotp(dto);
   }
 
   @Patch("password")
@@ -65,23 +67,22 @@ export class AuthController {
     return { accessToken };
   }
 
-  // TODO: Implement recovery codes to disable 2FA just in case someone gets locked out
   @Post("totp/enable")
   @UseGuards(JwtGuard)
   async enableTotp(@GetUser() user: User, @Body() body: EnableTotpDTO) {
-    return this.authService.enableTotp(user, body.password);
+    return this.authTotpService.enableTotp(user, body.password);
   }
 
   @Post("totp/verify")
   @UseGuards(JwtGuard)
   async verifyTotp(@GetUser() user: User, @Body() body: VerifyTotpDTO) {
-    return this.authService.verifyTotp(user, body.password, body.code);
+    return this.authTotpService.verifyTotp(user, body.password, body.code);
   }
 
   @Post("totp/disable")
   @UseGuards(JwtGuard)
   async disableTotp(@GetUser() user: User, @Body() body: VerifyTotpDTO) {
     // Note: We use VerifyTotpDTO here because it has both fields we need: password and totp code
-    return this.authService.disableTotp(user, body.password, body.code);
+    return this.authTotpService.disableTotp(user, body.password, body.code);
   }
 }

backend/src/auth/auth.module.ts

@@ -2,12 +2,13 @@ import { Module } from "@nestjs/common";
 import { JwtModule } from "@nestjs/jwt";
 import { AuthController } from "./auth.controller";
 import { AuthService } from "./auth.service";
+import { AuthTotpService } from "./authTotp.service";
 import { JwtStrategy } from "./strategy/jwt.strategy";
 
 @Module({
   imports: [JwtModule.register({})],
   controllers: [AuthController],
-  providers: [AuthService, JwtStrategy],
+  providers: [AuthService, AuthTotpService, JwtStrategy],
   exports: [AuthService],
 })
 export class AuthModule {}

backend/src/auth/auth.service.ts

@@ -13,10 +13,6 @@ import { ConfigService } from "src/config/config.service";
 import { PrismaService } from "src/prisma/prisma.service";
 import { AuthRegisterDTO } from "./dto/authRegister.dto";
 import { AuthSignInDTO } from "./dto/authSignIn.dto";
-import { authenticator, totp } from "otplib";
-import * as qrcode from "qrcode-svg";
-import * as crypto from "crypto";
-import { AuthSignInTotpDTO } from "./dto/authSignInTotp.dto";
 
 @Injectable()
 export class AuthService {
@@ -81,61 +77,6 @@ export class AuthService {
     return { accessToken, refreshToken };
   }
 
-  async signInTotp(dto: AuthSignInTotpDTO) {
-    if (!dto.email && !dto.username)
-      throw new BadRequestException("Email or username is required");
-
-    const user = await this.prisma.user.findFirst({
-      where: {
-        OR: [{ email: dto.email }, { username: dto.username }],
-      },
-    });
-
-    if (!user || !(await argon.verify(user.password, dto.password)))
-      throw new UnauthorizedException("Wrong email or password");
-
-    const token = await this.prisma.loginToken.findFirst({
-      where: {
-        token: dto.loginToken,
-      },
-    });
-
-    if (!token || token.userId != user.id || token.used)
-      throw new UnauthorizedException("Invalid login token");
-
-    if (token.expiresAt < new Date())
-      throw new UnauthorizedException("Login token expired");
-
-    // Check the TOTP code
-    const { totpSecret } = await this.prisma.user.findUnique({
-      where: { id: user.id },
-      select: { totpSecret: true },
-    });
-
-    if (!totpSecret) {
-      throw new BadRequestException("TOTP is not enabled");
-    }
-
-    const decryptedSecret = this.decryptTotpSecret(totpSecret, dto.password);
-
-    const expected = authenticator.generate(decryptedSecret);
-
-    if (dto.totp !== expected) {
-      throw new BadRequestException("Invalid code");
-    }
-
-    // Set the login token to used
-    await this.prisma.loginToken.update({
-      where: { token: token.token },
-      data: { used: true },
-    });
-
-    const accessToken = await this.createAccessToken(user);
-    const refreshToken = await this.createRefreshToken(user.id);
-
-    return { accessToken, refreshToken };
-  }
-
   async updatePassword(user: User, oldPassword: string, newPassword: string) {
     if (!(await argon.verify(user.password, oldPassword)))
       throw new ForbiddenException("Invalid password");
@@ -192,151 +133,4 @@ export class AuthService {
 
     return loginToken;
   }
-
-  encryptTotpSecret(totpSecret: string, password: string) {
-    let iv = this.config.get("TOTP_SECRET");
-    iv = Buffer.from(iv, "base64");
-    const key = crypto
-      .createHash("sha256")
-      .update(String(password))
-      .digest("base64")
-      .substr(0, 32);
-
-    const cipher = crypto.createCipheriv("aes-256-cbc", key, iv);
-
-    let encrypted = cipher.update(totpSecret);
-
-    encrypted = Buffer.concat([encrypted, cipher.final()]);
-
-    return encrypted.toString("base64");
-  }
-
-  decryptTotpSecret(encryptedTotpSecret: string, password: string) {
-    let iv = this.config.get("TOTP_SECRET");
-    iv = Buffer.from(iv, "base64");
-    const key = crypto
-      .createHash("sha256")
-      .update(String(password))
-      .digest("base64")
-      .substr(0, 32);
-
-    const encryptedText = Buffer.from(encryptedTotpSecret, "base64");
-    const decipher = crypto.createDecipheriv("aes-256-cbc", key, iv);
-    let decrypted = decipher.update(encryptedText);
-    decrypted = Buffer.concat([decrypted, decipher.final()]);
-
-    return decrypted.toString();
-  }
-
-  async enableTotp(user: User, password: string) {
-    if (!(await argon.verify(user.password, password)))
-      throw new ForbiddenException("Invalid password");
-
-    // Check if we have a secret already
-    const { totpVerified } = await this.prisma.user.findUnique({
-      where: { id: user.id },
-      select: { totpVerified: true },
-    });
-
-    if (totpVerified) {
-      throw new BadRequestException("TOTP is already enabled");
-    }
-
-    // TODO: Maybe make the issuer configurable with env vars?
-    const secret = authenticator.generateSecret();
-    const encryptedSecret = this.encryptTotpSecret(secret, password);
-
-    const otpURL = totp.keyuri(
-      user.username || user.email,
-      "pingvin-share",
-      secret
-    );
-
-    await this.prisma.user.update({
-      where: { id: user.id },
-      data: {
-        totpEnabled: true,
-        totpSecret: encryptedSecret,
-      },
-    });
-
-    // TODO: Maybe we should generate the QR code on the client rather than the server?
-    const qrCode = new qrcode({
-      content: otpURL,
-      container: "svg-viewbox",
-      join: true,
-    }).svg();
-
-    return {
-      totpAuthUrl: otpURL,
-      totpSecret: secret,
-      qrCode:
-        "data:image/svg+xml;base64," + Buffer.from(qrCode).toString("base64"),
-    };
-  }
-
-  // TODO: Maybe require a token to verify that the user who started enabling totp is the one who is verifying it?
-  async verifyTotp(user: User, password: string, code: string) {
-    if (!(await argon.verify(user.password, password)))
-      throw new ForbiddenException("Invalid password");
-
-    const { totpSecret } = await this.prisma.user.findUnique({
-      where: { id: user.id },
-      select: { totpSecret: true },
-    });
-
-    if (!totpSecret) {
-      throw new BadRequestException("TOTP is not in progress");
-    }
-
-    const decryptedSecret = this.decryptTotpSecret(totpSecret, password);
-
-    const expected = authenticator.generate(decryptedSecret);
-
-    if (code !== expected) {
-      throw new BadRequestException("Invalid code");
-    }
-
-    await this.prisma.user.update({
-      where: { id: user.id },
-      data: {
-        totpVerified: true,
-      },
-    });
-
-    return true;
-  }
-
-  async disableTotp(user: User, password: string, code: string) {
-    if (!(await argon.verify(user.password, password)))
-      throw new ForbiddenException("Invalid password");
-
-    const { totpSecret } = await this.prisma.user.findUnique({
-      where: { id: user.id },
-      select: { totpSecret: true },
-    });
-
-    if (!totpSecret) {
-      throw new BadRequestException("TOTP is not enabled");
-    }
-
-    const decryptedSecret = this.decryptTotpSecret(totpSecret, password);
-
-    const expected = authenticator.generate(decryptedSecret);
-
-    if (code !== expected) {
-      throw new BadRequestException("Invalid code");
-    }
-
-    await this.prisma.user.update({
-      where: { id: user.id },
-      data: {
-        totpVerified: false,
-        totpEnabled: false,
-        totpSecret: null,
-      },
-    });
-
-    return true;
-  }
 }

backend/src/auth/authTotp.service.ts

@@ -0,0 +1,224 @@
+import {
+  BadRequestException,
+  ForbiddenException,
+  UnauthorizedException,
+} from "@nestjs/common";
+import { ConfigService } from "@nestjs/config";
+import { User } from "@prisma/client";
+import * as argon from "argon2";
+import * as crypto from "crypto";
+import { authenticator, totp } from "otplib";
+import * as qrcode from "qrcode-svg";
+import { PrismaService } from "src/prisma/prisma.service";
+import { AuthService } from "./auth.service";
+import { AuthSignInTotpDTO } from "./dto/authSignInTotp.dto";
+
+export class AuthTotpService {
+  constructor(
+    private config: ConfigService,
+    private prisma: PrismaService,
+    private authService: AuthService
+  ) {}
+
+  async signInTotp(dto: AuthSignInTotpDTO) {
+    if (!dto.email && !dto.username)
+      throw new BadRequestException("Email or username is required");
+
+    const user = await this.prisma.user.findFirst({
+      where: {
+        OR: [{ email: dto.email }, { username: dto.username }],
+      },
+    });
+
+    if (!user || !(await argon.verify(user.password, dto.password)))
+      throw new UnauthorizedException("Wrong email or password");
+
+    const token = await this.prisma.loginToken.findFirst({
+      where: {
+        token: dto.loginToken,
+      },
+    });
+
+    if (!token || token.userId != user.id || token.used)
+      throw new UnauthorizedException("Invalid login token");
+
+    if (token.expiresAt < new Date())
+      throw new UnauthorizedException("Login token expired");
+
+    // Check the TOTP code
+    const { totpSecret } = await this.prisma.user.findUnique({
+      where: { id: user.id },
+      select: { totpSecret: true },
+    });
+
+    if (!totpSecret) {
+      throw new BadRequestException("TOTP is not enabled");
+    }
+
+    const decryptedSecret = this.decryptTotpSecret(totpSecret, dto.password);
+
+    const expected = authenticator.generate(decryptedSecret);
+
+    if (dto.totp !== expected) {
+      throw new BadRequestException("Invalid code");
+    }
+
+    // Set the login token to used
+    await this.prisma.loginToken.update({
+      where: { token: token.token },
+      data: { used: true },
+    });
+
+    const accessToken = await this.authService.createAccessToken(user);
+    const refreshToken = await this.authService.createRefreshToken(user.id);
+
+    return { accessToken, refreshToken };
+  }
+
+  encryptTotpSecret(totpSecret: string, password: string) {
+    let iv = this.config.get("TOTP_SECRET");
+    iv = Buffer.from(iv, "base64");
+    const key = crypto
+      .createHash("sha256")
+      .update(String(password))
+      .digest("base64")
+      .substr(0, 32);
+
+    const cipher = crypto.createCipheriv("aes-256-cbc", key, iv);
+
+    let encrypted = cipher.update(totpSecret);
+
+    encrypted = Buffer.concat([encrypted, cipher.final()]);
+
+    return encrypted.toString("base64");
+  }
+
+  decryptTotpSecret(encryptedTotpSecret: string, password: string) {
+    let iv = this.config.get("TOTP_SECRET");
+    iv = Buffer.from(iv, "base64");
+    const key = crypto
+      .createHash("sha256")
+      .update(String(password))
+      .digest("base64")
+      .substr(0, 32);
+
+    const encryptedText = Buffer.from(encryptedTotpSecret, "base64");
+    const decipher = crypto.createDecipheriv("aes-256-cbc", key, iv);
+    let decrypted = decipher.update(encryptedText);
+    decrypted = Buffer.concat([decrypted, decipher.final()]);
+
+    return decrypted.toString();
+  }
+
+  async enableTotp(user: User, password: string) {
+    if (!(await argon.verify(user.password, password)))
+      throw new ForbiddenException("Invalid password");
+
+    // Check if we have a secret already
+    const { totpVerified } = await this.prisma.user.findUnique({
+      where: { id: user.id },
+      select: { totpVerified: true },
+    });
+
+    if (totpVerified) {
+      throw new BadRequestException("TOTP is already enabled");
+    }
+
+    // TODO: Maybe make the issuer configurable with env vars?
+    const secret = authenticator.generateSecret();
+    const encryptedSecret = this.encryptTotpSecret(secret, password);
+
+    const otpURL = totp.keyuri(
+      user.username || user.email,
+      "pingvin-share",
+      secret
+    );
+
+    await this.prisma.user.update({
+      where: { id: user.id },
+      data: {
+        totpEnabled: true,
+        totpSecret: encryptedSecret,
+      },
+    });
+
+    // TODO: Maybe we should generate the QR code on the client rather than the server?
+    const qrCode = new qrcode({
+      content: otpURL,
+      container: "svg-viewbox",
+      join: true,
+    }).svg();
+
+    return {
+      totpAuthUrl: otpURL,
+      totpSecret: secret,
+      qrCode:
+        "data:image/svg+xml;base64," + Buffer.from(qrCode).toString("base64"),
+    };
+  }
+
+  // TODO: Maybe require a token to verify that the user who started enabling totp is the one who is verifying it?
+  async verifyTotp(user: User, password: string, code: string) {
+    if (!(await argon.verify(user.password, password)))
+      throw new ForbiddenException("Invalid password");
+
+    const { totpSecret } = await this.prisma.user.findUnique({
+      where: { id: user.id },
+      select: { totpSecret: true },
+    });
+
+    if (!totpSecret) {
+      throw new BadRequestException("TOTP is not in progress");
+    }
+
+    const decryptedSecret = this.decryptTotpSecret(totpSecret, password);
+
+    const expected = authenticator.generate(decryptedSecret);
+
+    if (code !== expected) {
+      throw new BadRequestException("Invalid code");
+    }
+
+    await this.prisma.user.update({
+      where: { id: user.id },
+      data: {
+        totpVerified: true,
+      },
+    });
+
+    return true;
+  }
+
+  async disableTotp(user: User, password: string, code: string) {
+    if (!(await argon.verify(user.password, password)))
+      throw new ForbiddenException("Invalid password");
+
+    const { totpSecret } = await this.prisma.user.findUnique({
+      where: { id: user.id },
+      select: { totpSecret: true },
+    });
+
+    if (!totpSecret) {
+      throw new BadRequestException("TOTP is not enabled");
+    }
+
+    const decryptedSecret = this.decryptTotpSecret(totpSecret, password);
+
+    const expected = authenticator.generate(decryptedSecret);
+
+    if (code !== expected) {
+      throw new BadRequestException("Invalid code");
+    }
+
+    await this.prisma.user.update({
+      where: { id: user.id },
+      data: {
+        totpVerified: false,
+        totpEnabled: false,
+        totpSecret: null,
+      },
+    });
+
+    return true;
+  }
+}