feat(server): verificationToken model (#9655)

packages/backend/server/src/__tests__/models/verification-token.spec.ts

@@ -0,0 +1,193 @@
+import { TestingModule } from '@nestjs/testing';
+import { PrismaClient } from '@prisma/client';
+import ava, { TestFn } from 'ava';
+
+import {
+  TokenType,
+  VerificationTokenModel,
+} from '../../models/verification-token';
+import { createTestingModule, initTestingDB } from '../utils';
+
+interface Context {
+  module: TestingModule;
+  verificationToken: VerificationTokenModel;
+  db: PrismaClient;
+}
+
+const test = ava as TestFn<Context>;
+
+test.before(async t => {
+  const module = await createTestingModule({
+    providers: [VerificationTokenModel],
+  });
+
+  t.context.verificationToken = module.get(VerificationTokenModel);
+  t.context.db = module.get(PrismaClient);
+  t.context.module = module;
+});
+
+test.beforeEach(async t => {
+  await initTestingDB(t.context.db);
+});
+
+test.after(async t => {
+  await t.context.module.close();
+});
+
+test('should be able to create token', async t => {
+  const { verificationToken } = t.context;
+  const token = await verificationToken.create(
+    TokenType.SignIn,
+    'user@affine.pro'
+  );
+
+  t.truthy(
+    await verificationToken.verify(TokenType.SignIn, token, {
+      credential: 'user@affine.pro',
+    })
+  );
+});
+
+test('should be able to get token', async t => {
+  const { verificationToken } = t.context;
+  const token = await verificationToken.create(
+    TokenType.SignIn,
+    'user@affine.pro'
+  );
+
+  t.truthy(await verificationToken.get(TokenType.SignIn, token));
+  // will be delete after the first time of verification
+  t.falsy(await verificationToken.get(TokenType.SignIn, token));
+});
+
+test('should be able to get token and keep work', async t => {
+  const { verificationToken } = t.context;
+  const token = await verificationToken.create(
+    TokenType.SignIn,
+    'user@affine.pro'
+  );
+
+  t.truthy(await verificationToken.get(TokenType.SignIn, token, true));
+  t.truthy(await verificationToken.get(TokenType.SignIn, token));
+  t.falsy(await verificationToken.get(TokenType.SignIn, token));
+});
+
+test('should fail the verification if the token is invalid', async t => {
+  const { verificationToken } = t.context;
+  const token = await verificationToken.create(
+    TokenType.SignIn,
+    'user@affine.pro'
+  );
+
+  // wrong type
+  t.falsy(
+    await verificationToken.verify(TokenType.ChangeEmail, token, {
+      credential: 'user@affine.pro',
+    })
+  );
+
+  // no credential
+  t.falsy(await verificationToken.verify(TokenType.SignIn, token));
+
+  // wrong credential
+  t.falsy(
+    await verificationToken.verify(TokenType.SignIn, token, {
+      credential: 'wrong@affine.pro',
+    })
+  );
+});
+
+test('should fail if the token expired', async t => {
+  const { verificationToken, db } = t.context;
+  const token = await verificationToken.create(
+    TokenType.SignIn,
+    'user@affine.pro'
+  );
+
+  await db.verificationToken.updateMany({
+    data: {
+      expiresAt: new Date(Date.now() - 1000),
+    },
+  });
+
+  t.falsy(
+    await verificationToken.verify(TokenType.SignIn, token, {
+      credential: 'user@affine.pro',
+    })
+  );
+});
+
+test('should be able to verify without credential', async t => {
+  const { verificationToken } = t.context;
+  const token = await verificationToken.create(TokenType.SignIn);
+
+  t.truthy(await verificationToken.verify(TokenType.SignIn, token));
+
+  // will be invalid after the first time of verification
+  t.falsy(await verificationToken.verify(TokenType.SignIn, token));
+});
+
+test('should be able to verify only once', async t => {
+  const { verificationToken } = t.context;
+  const token = await verificationToken.create(
+    TokenType.SignIn,
+    'user@affine.pro'
+  );
+
+  t.truthy(
+    await verificationToken.verify(TokenType.SignIn, token, {
+      credential: 'user@affine.pro',
+    })
+  );
+
+  // will be invalid after the first time of verification
+  t.falsy(
+    await verificationToken.verify(TokenType.SignIn, token, {
+      credential: 'user@affine.pro',
+    })
+  );
+});
+
+test('should be able to verify and keep work', async t => {
+  const { verificationToken } = t.context;
+  const token = await verificationToken.create(
+    TokenType.SignIn,
+    'user@affine.pro'
+  );
+
+  t.truthy(
+    await verificationToken.verify(TokenType.SignIn, token, {
+      credential: 'user@affine.pro',
+      keep: true,
+    })
+  );
+
+  t.truthy(
+    await verificationToken.verify(TokenType.SignIn, token, {
+      credential: 'user@affine.pro',
+    })
+  );
+
+  // will be invalid without keep
+  t.falsy(
+    await verificationToken.verify(TokenType.SignIn, token, {
+      credential: 'user@affine.pro',
+    })
+  );
+});
+
+test('should cleanup expired tokens', async t => {
+  const { verificationToken, db } = t.context;
+  await verificationToken.create(TokenType.SignIn, 'user@affine.pro');
+
+  await db.verificationToken.updateMany({
+    data: {
+      expiresAt: new Date(Date.now() - 1000),
+    },
+  });
+
+  let count = await verificationToken.cleanExpired();
+  t.is(count, 1);
+  count = await verificationToken.cleanExpired();
+  t.is(count, 0);
+});

packages/backend/server/src/models/index.ts

@@ -2,14 +2,18 @@ import { Global, Injectable, Module } from '@nestjs/common';
 
 import { SessionModel } from './session';
 import { UserModel } from './user';
+import { VerificationTokenModel } from './verification-token';
 
-const models = [UserModel, SessionModel] as const;
+export * from './verification-token';
+
+const models = [UserModel, SessionModel, VerificationTokenModel] as const;
 
 @Injectable()
 export class Models {
   constructor(
     public readonly user: UserModel,
-    public readonly session: SessionModel
+    public readonly session: SessionModel,
+    public readonly verificationToken: VerificationTokenModel
   ) {}
 }
 

packages/backend/server/src/models/verification-token.ts

@@ -0,0 +1,144 @@
+import { randomUUID } from 'node:crypto';
+
+import { Injectable, Logger } from '@nestjs/common';
+import { PrismaClient, type VerificationToken } from '@prisma/client';
+
+import { CryptoHelper } from '../base/helpers';
+
+export type { VerificationToken };
+
+export enum TokenType {
+  SignIn,
+  VerifyEmail,
+  ChangeEmail,
+  ChangePassword,
+  Challenge,
+}
+
+@Injectable()
+export class VerificationTokenModel {
+  private readonly logger = new Logger(VerificationTokenModel.name);
+  constructor(
+    private readonly db: PrismaClient,
+    private readonly crypto: CryptoHelper
+  ) {}
+
+  /**
+   * create token by type and credential (optional) with ttl in seconds (default 30 minutes)
+   */
+  async create(
+    type: TokenType,
+    credential?: string,
+    ttlInSec: number = 30 * 60
+  ) {
+    const plaintextToken = randomUUID();
+    const { token } = await this.db.verificationToken.create({
+      data: {
+        type,
+        token: plaintextToken,
+        credential,
+        expiresAt: new Date(Date.now() + ttlInSec * 1000),
+      },
+    });
+    return this.crypto.encrypt(token);
+  }
+
+  /**
+   * get token by type
+   *
+   * token will be deleted if expired or keep is not set
+   */
+  async get(type: TokenType, token: string, keep?: boolean) {
+    token = this.crypto.decrypt(token);
+    const record = await this.db.verificationToken.findUnique({
+      where: {
+        type_token: {
+          token,
+          type,
+        },
+      },
+    });
+
+    if (!record) {
+      return null;
+    }
+
+    const isExpired = record.expiresAt <= new Date();
+
+    // always delete expired token
+    // or if keep is not set for one time token
+    if (isExpired || !keep) {
+      const count = await this.delete(type, token);
+
+      // already deleted, means token has been used
+      if (!count) {
+        return null;
+      }
+    }
+
+    return !isExpired ? record : null;
+  }
+
+  /**
+   * get token and verify credential
+   *
+   * if credential is not provided, it will be failed
+   *
+   * token will be deleted if expired or keep is not set
+   */
+  async verify(
+    type: TokenType,
+    token: string,
+    {
+      credential,
+      keep,
+    }: {
+      credential?: string;
+      keep?: boolean;
+    } = {}
+  ) {
+    const record = await this.get(type, token, true);
+    if (!record) {
+      return null;
+    }
+
+    const valid = !record.credential || record.credential === credential;
+    // keep is not set for one time valid token
+    if (valid && !keep) {
+      const count = await this.delete(type, record.token);
+
+      // already deleted, means token has been used
+      if (!count) {
+        return null;
+      }
+    }
+
+    return valid ? record : null;
+  }
+
+  async delete(type: TokenType, token: string) {
+    const { count } = await this.db.verificationToken.deleteMany({
+      where: {
+        token,
+        type,
+      },
+    });
+    this.logger.log(`Deleted token success by type ${type} and token ${token}`);
+    return count;
+  }
+
+  /**
+   * clean expired tokens
+   */
+  async cleanExpired() {
+    const { count } = await this.db.verificationToken.deleteMany({
+      where: {
+        expiresAt: {
+          lte: new Date(),
+        },
+      },
+    });
+    this.logger.log(`Cleaned ${count} expired tokens`);
+    return count;
+  }
+}