diff --git a/packages/backend/server/src/__tests__/models/verification-token.spec.ts b/packages/backend/server/src/__tests__/models/verification-token.spec.ts new file mode 100644 index 0000000000..03680abef1 --- /dev/null +++ b/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; + +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); +}); diff --git a/packages/backend/server/src/models/index.ts b/packages/backend/server/src/models/index.ts index 32f50132ba..0edcc2e90a 100644 --- a/packages/backend/server/src/models/index.ts +++ b/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 ) {} } diff --git a/packages/backend/server/src/models/verification-token.ts b/packages/backend/server/src/models/verification-token.ts new file mode 100644 index 0000000000..85be612b5e --- /dev/null +++ b/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; + } +}