diff --git a/apps/server/package.json b/apps/server/package.json index a7343cfce3..e940e24009 100644 --- a/apps/server/package.json +++ b/apps/server/package.json @@ -12,7 +12,9 @@ "start": "node --loader ts-node/esm.mjs --es-module-specifier-resolution node ./src/index.ts", "dev": "nodemon ./src/index.ts", "test": "yarn exec ts-node-esm ./scripts/run-test.ts all", + "test:select": "yarn exec ts-node-esm ./scripts/run-test.ts", "test:watch": "yarn exec ts-node-esm ./scripts/run-test.ts all --watch", + "test:select:watch": "yarn exec ts-node-esm ./scripts/run-test.ts --watch", "test:coverage": "c8 yarn exec ts-node-esm ./scripts/run-test.ts all", "postinstall": "prisma generate" }, @@ -28,6 +30,7 @@ "@nestjs/graphql": "^12.0.8", "@nestjs/platform-express": "^10.2.2", "@nestjs/platform-socket.io": "^10.2.2", + "@nestjs/throttler": "^4.2.1", "@nestjs/websockets": "^10.2.2", "@node-rs/argon2": "^1.5.2", "@node-rs/crc32": "^1.7.2", @@ -55,6 +58,7 @@ "graphql-upload": "^16.0.2", "ioredis": "^5.3.2", "lodash-es": "^4.17.21", + "nestjs-throttler-storage-redis": "^0.3.3", "next-auth": "4.22.5", "nodemailer": "^6.9.4", "on-headers": "^1.0.2", diff --git a/apps/server/src/app.ts b/apps/server/src/app.ts index db3b692626..f85878e660 100644 --- a/apps/server/src/app.ts +++ b/apps/server/src/app.ts @@ -6,6 +6,7 @@ import { MetricsModule } from './metrics'; import { BusinessModules } from './modules'; import { PrismaModule } from './prisma'; import { StorageModule } from './storage'; +import { RateLimiterModule } from './throttler'; @Module({ imports: [ @@ -13,6 +14,7 @@ import { StorageModule } from './storage'; ConfigModule.forRoot(), StorageModule.forRoot(), MetricsModule, + RateLimiterModule, ...BusinessModules, ], controllers: [AppController], diff --git a/apps/server/src/config/def.ts b/apps/server/src/config/def.ts index b313d55be6..469e338b12 100644 --- a/apps/server/src/config/def.ts +++ b/apps/server/src/config/def.ts @@ -187,6 +187,25 @@ export interface AFFiNEConfig { path: string; }; }; + + /** + * Rate limiter config + */ + rateLimiter: { + /** + * How long each request will be throttled (seconds) + * @default 60 + * @env THROTTLE_TTL + */ + ttl: number; + /** + * How many requests can be made in the given time frame + * @default 60 + * @env THROTTLE_LIMIT + */ + limit: number; + }; + /** * Redis Config * diff --git a/apps/server/src/config/default.ts b/apps/server/src/config/default.ts index 478f4ab8cb..b1df181d6b 100644 --- a/apps/server/src/config/default.ts +++ b/apps/server/src/config/default.ts @@ -72,6 +72,8 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => { OAUTH_EMAIL_SERVER: 'auth.email.server', OAUTH_EMAIL_PORT: ['auth.email.port', 'int'], OAUTH_EMAIL_PASSWORD: 'auth.email.password', + THROTTLE_TTL: ['rateLimiter.ttl', 'int'], + THROTTLE_LIMIT: ['rateLimiter.limit', 'int'], REDIS_SERVER_ENABLED: ['redis.enabled', 'boolean'], REDIS_SERVER_HOST: 'redis.host', REDIS_SERVER_PORT: ['redis.port', 'int'], @@ -169,6 +171,10 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => { path: join(homedir(), '.affine-storage'), }, }, + rateLimiter: { + ttl: 60, + limit: 60, + }, redis: { enabled: false, host: '127.0.0.1', diff --git a/apps/server/src/modules/auth/mailer/mail.service.ts b/apps/server/src/modules/auth/mailer/mail.service.ts index 5c9505d364..58d445264f 100644 --- a/apps/server/src/modules/auth/mailer/mail.service.ts +++ b/apps/server/src/modules/auth/mailer/mail.service.ts @@ -42,8 +42,6 @@ export class MailService { }; } ) { - console.log('invitationInfo', invitationInfo); - const buttonUrl = `${this.config.baseUrl}/invite/${inviteId}`; const workspaceAvatar = invitationInfo.workspace.avatar; diff --git a/apps/server/src/modules/auth/next-auth.controller.ts b/apps/server/src/modules/auth/next-auth.controller.ts index 014b8e75ba..65e542de84 100644 --- a/apps/server/src/modules/auth/next-auth.controller.ts +++ b/apps/server/src/modules/auth/next-auth.controller.ts @@ -9,6 +9,7 @@ import { Query, Req, Res, + UseGuards, } from '@nestjs/common'; import { hash, verify } from '@node-rs/argon2'; import type { User } from '@prisma/client'; @@ -19,6 +20,7 @@ import { AuthHandler } from 'next-auth/core'; import { Config } from '../../config'; import { PrismaService } from '../../prisma/service'; +import { CloudThrottlerGuard, Throttle } from '../../throttler'; import { NextAuthOptionsProvide } from './next-auth-options'; import { AuthService } from './service'; @@ -41,6 +43,8 @@ export class NextAuthController { this.callbackSession = nextAuthOptions.callbacks!.session; } + @UseGuards(CloudThrottlerGuard) + @Throttle(20, 60) @All('*') async auth( @Req() req: Request, diff --git a/apps/server/src/modules/auth/resolver.ts b/apps/server/src/modules/auth/resolver.ts index 478808358e..68b74b4ab3 100644 --- a/apps/server/src/modules/auth/resolver.ts +++ b/apps/server/src/modules/auth/resolver.ts @@ -1,4 +1,4 @@ -import { ForbiddenException } from '@nestjs/common'; +import { ForbiddenException, UseGuards } from '@nestjs/common'; import { Args, Context, @@ -12,6 +12,7 @@ import { import type { Request } from 'express'; import { Config } from '../../config'; +import { CloudThrottlerGuard, Throttle } from '../../throttler'; import { UserType } from '../users/resolver'; import { CurrentUser } from './guard'; import { AuthService } from './service'; @@ -25,6 +26,13 @@ export class TokenType { refresh!: string; } +/** + * Auth resolver + * Token rate limit: 20 req/m + * Sign up/in rate limit: 10 req/m + * Other rate limit: 5 req/m + */ +@UseGuards(CloudThrottlerGuard) @Resolver(() => UserType) export class AuthResolver { constructor( @@ -32,6 +40,7 @@ export class AuthResolver { private auth: AuthService ) {} + @Throttle(20, 60) @ResolveField(() => TokenType) token(@CurrentUser() currentUser: UserType, @Parent() user: UserType) { if (user.id !== currentUser.id) { @@ -44,6 +53,7 @@ export class AuthResolver { }; } + @Throttle(10, 60) @Mutation(() => UserType) async signUp( @Context() ctx: { req: Request }, @@ -56,6 +66,7 @@ export class AuthResolver { return user; } + @Throttle(10, 60) @Mutation(() => UserType) async signIn( @Context() ctx: { req: Request }, @@ -67,6 +78,7 @@ export class AuthResolver { return user; } + @Throttle(5, 60) @Mutation(() => UserType) async changePassword( @Context() ctx: { req: Request }, @@ -78,6 +90,7 @@ export class AuthResolver { return user; } + @Throttle(5, 60) @Mutation(() => UserType) async changeEmail( @Context() ctx: { req: Request }, @@ -89,6 +102,7 @@ export class AuthResolver { return user; } + @Throttle(5, 60) @Mutation(() => Boolean) async sendChangePasswordEmail( @Args('email') email: string, @@ -99,6 +113,7 @@ export class AuthResolver { return !res.rejected.length; } + @Throttle(5, 60) @Mutation(() => Boolean) async sendSetPasswordEmail( @Args('email') email: string, @@ -109,6 +124,7 @@ export class AuthResolver { return !res.rejected.length; } + @Throttle(5, 60) @Mutation(() => Boolean) async sendChangeEmail( @Args('email') email: string, diff --git a/apps/server/src/modules/users/resolver.ts b/apps/server/src/modules/users/resolver.ts index 0d3e28bff9..4d982eee9c 100644 --- a/apps/server/src/modules/users/resolver.ts +++ b/apps/server/src/modules/users/resolver.ts @@ -2,6 +2,7 @@ import { BadRequestException, ForbiddenException, HttpException, + UseGuards, } from '@nestjs/common'; import { Args, @@ -19,6 +20,7 @@ import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs'; import { Config } from '../../config'; import { PrismaService } from '../../prisma/service'; +import { CloudThrottlerGuard, Throttle } from '../../throttler'; import type { FileUpload } from '../../types'; import { Auth, CurrentUser, Public } from '../auth/guard'; import { StorageService } from '../storage/storage.service'; @@ -69,6 +71,11 @@ export class AddToNewFeaturesWaitingList { type!: NewFeaturesKind; } +/** + * User resolver + * All op rate limit: 10 req/m + */ +@UseGuards(CloudThrottlerGuard) @Auth() @Resolver(() => UserType) export class UserResolver { @@ -78,6 +85,7 @@ export class UserResolver { private readonly config: Config ) {} + @Throttle(10, 60) @Query(() => UserType, { name: 'currentUser', description: 'Get current user', @@ -100,6 +108,7 @@ export class UserResolver { }; } + @Throttle(10, 60) @Query(() => UserType, { name: 'user', description: 'Get user by email', @@ -135,6 +144,7 @@ export class UserResolver { return user; } + @Throttle(10, 60) @Mutation(() => UserType, { name: 'uploadAvatar', description: 'Upload user avatar', @@ -155,6 +165,7 @@ export class UserResolver { }); } + @Throttle(10, 60) @Mutation(() => DeleteAccount) async deleteAccount(@CurrentUser() user: UserType): Promise { await this.prisma.user.delete({ @@ -172,6 +183,7 @@ export class UserResolver { }; } + @Throttle(10, 60) @Mutation(() => AddToNewFeaturesWaitingList) async addToNewFeaturesWaitingList( @CurrentUser() user: UserType, diff --git a/apps/server/src/modules/workspaces/resolver.ts b/apps/server/src/modules/workspaces/resolver.ts index d95e0d1e6e..8ee0d8a944 100644 --- a/apps/server/src/modules/workspaces/resolver.ts +++ b/apps/server/src/modules/workspaces/resolver.ts @@ -1,5 +1,10 @@ import type { Storage } from '@affine/storage'; -import { ForbiddenException, Inject, NotFoundException } from '@nestjs/common'; +import { + ForbiddenException, + Inject, + NotFoundException, + UseGuards, +} from '@nestjs/common'; import { Args, Field, @@ -24,6 +29,7 @@ import { applyUpdate, Doc } from 'yjs'; import { PrismaService } from '../../prisma'; import { StorageProvide } from '../../storage'; +import { CloudThrottlerGuard, Throttle } from '../../throttler'; import type { FileUpload } from '../../types'; import { Auth, CurrentUser, Public } from '../auth'; import { MailService } from '../auth/mailer'; @@ -113,6 +119,12 @@ export class UpdateWorkspaceInput extends PickType( id!: string; } +/** + * Workspace resolver + * Public apis rate limit: 10 req/m + * Other rate limit: 120 req/m + */ +@UseGuards(CloudThrottlerGuard) @Auth() @Resolver(() => WorkspaceType) export class WorkspaceResolver { @@ -258,10 +270,11 @@ export class WorkspaceResolver { }); } + @Throttle(10, 30) + @Public() @Query(() => WorkspaceType, { description: 'Get public workspace by id', }) - @Public() async publicWorkspace(@Args('id') id: string) { const workspace = await this.prisma.workspace.findUnique({ where: { id }, @@ -463,6 +476,7 @@ export class WorkspaceResolver { } } + @Throttle(10, 30) @Public() @Query(() => InvitationType, { description: 'Update workspace', diff --git a/apps/server/src/tests/auth.spec.ts b/apps/server/src/tests/auth.spec.ts index dc36fd2e66..f3a700d4a5 100644 --- a/apps/server/src/tests/auth.spec.ts +++ b/apps/server/src/tests/auth.spec.ts @@ -11,6 +11,7 @@ import { MetricsModule } from '../metrics'; import { AuthModule } from '../modules/auth'; import { AuthService } from '../modules/auth/service'; import { PrismaModule } from '../prisma'; +import { RateLimiterModule } from '../throttler'; let auth: AuthService; let module: TestingModule; @@ -36,6 +37,7 @@ beforeEach(async () => { GqlModule, AuthModule, MetricsModule, + RateLimiterModule, ], }).compile(); auth = module.get(AuthService); diff --git a/apps/server/src/tests/user.spec.ts b/apps/server/src/tests/user.spec.ts index 8b468fb4c2..dd55c1d8ad 100644 --- a/apps/server/src/tests/user.spec.ts +++ b/apps/server/src/tests/user.spec.ts @@ -1,4 +1,4 @@ -import { ok } from 'node:assert'; +import { ok, rejects } from 'node:assert'; import { afterEach, beforeEach, describe, it } from 'node:test'; import type { INestApplication } from '@nestjs/common'; @@ -71,7 +71,6 @@ describe('User Module', () => { `, }) .expect(200); - const current = await currentUser(app, user.token.token); - ok(current == null); + rejects(currentUser(app, user.token.token)); }); }); diff --git a/apps/server/src/tests/utils.ts b/apps/server/src/tests/utils.ts index 9be2dc13f2..addef78f10 100644 --- a/apps/server/src/tests/utils.ts +++ b/apps/server/src/tests/utils.ts @@ -43,13 +43,14 @@ async function currentUser(app: INestApplication, token: string) { query: ` query { currentUser { - id, name, email, emailVerified, avatarUrl, createdAt, hasPassword + id, name, email, emailVerified, avatarUrl, createdAt, hasPassword, + token { token } } } `, }) .expect(200); - return res.body?.data?.currentUser; + return res.body.data.currentUser; } async function createWorkspace( @@ -440,7 +441,7 @@ async function getInviteInfo( `, }) .expect(200); - return res.body.data.workspace; + return res.body.data.getInviteInfo; } export { diff --git a/apps/server/src/tests/workspace.spec.ts b/apps/server/src/tests/workspace.spec.ts index 3035fbab15..3f2298cf9a 100644 --- a/apps/server/src/tests/workspace.spec.ts +++ b/apps/server/src/tests/workspace.spec.ts @@ -12,6 +12,7 @@ import { AppModule } from '../app'; import { acceptInvite, createWorkspace, + currentUser, getPublicWorkspace, getWorkspaceSharedPages, inviteUser, @@ -61,6 +62,18 @@ describe('Workspace Module', () => { ok(user.email === 'u1@affine.pro', 'user.email is not valid'); }); + it('should be throttled at call signUp', async () => { + let token = ''; + for (let i = 0; i < 10; i++) { + token = (await signUp(app, `u${i}`, `u${i}@affine.pro`, `${i}`)).token + .token; + // throttles are applied to each endpoint separately + await currentUser(app, token); + } + await rejects(signUp(app, 'u11', 'u11@affine.pro', '11')); + await rejects(currentUser(app, token)); + }); + it('should create a workspace', async () => { const user = await signUp(app, 'u1', 'u1@affine.pro', '1'); diff --git a/apps/server/src/throttler.ts b/apps/server/src/throttler.ts new file mode 100644 index 0000000000..3a1db66b6f --- /dev/null +++ b/apps/server/src/throttler.ts @@ -0,0 +1,54 @@ +import { ExecutionContext, Injectable, Logger } from '@nestjs/common'; +import { Global, Module } from '@nestjs/common'; +import { + Throttle, + ThrottlerGuard, + ThrottlerModule, + ThrottlerModuleOptions, +} from '@nestjs/throttler'; +import Redis from 'ioredis'; +import { ThrottlerStorageRedisService } from 'nestjs-throttler-storage-redis'; + +import { Config, ConfigModule } from './config'; +import { getRequestResponseFromContext } from './utils/nestjs'; + +@Global() +@Module({ + imports: [ + ThrottlerModule.forRootAsync({ + imports: [ConfigModule], + inject: [Config], + useFactory: (config: Config): ThrottlerModuleOptions => { + const options: ThrottlerModuleOptions = { + ttl: config.rateLimiter.ttl, + limit: config.rateLimiter.limit, + }; + if (config.redis.enabled) { + new Logger(RateLimiterModule.name).log('Use Redis'); + options.storage = new ThrottlerStorageRedisService( + new Redis(config.redis.port, config.redis.host, { + username: config.redis.username, + password: config.redis.password, + db: config.redis.database + 1, + }) + ); + } + return options; + }, + }), + ], +}) +export class RateLimiterModule {} + +@Injectable() +export class CloudThrottlerGuard extends ThrottlerGuard { + override getRequestResponse(context: ExecutionContext) { + return getRequestResponseFromContext(context) as any; + } + + protected override getTracker(req: Record): string { + return req?.get('CF-Connecting-IP') ?? req?.get('CF-ray') ?? req?.ip; + } +} + +export { Throttle }; diff --git a/yarn.lock b/yarn.lock index 29acb90037..af3411783a 100644 --- a/yarn.lock +++ b/yarn.lock @@ -669,6 +669,7 @@ __metadata: "@nestjs/platform-express": ^10.2.2 "@nestjs/platform-socket.io": ^10.2.2 "@nestjs/testing": ^10.2.2 + "@nestjs/throttler": ^4.2.1 "@nestjs/websockets": ^10.2.2 "@node-rs/argon2": ^1.5.2 "@node-rs/crc32": ^1.7.2 @@ -708,6 +709,7 @@ __metadata: graphql-upload: ^16.0.2 ioredis: ^5.3.2 lodash-es: ^4.17.21 + nestjs-throttler-storage-redis: ^0.3.3 next-auth: 4.22.5 nodemailer: ^6.9.4 nodemon: ^3.0.1 @@ -7297,6 +7299,19 @@ __metadata: languageName: node linkType: hard +"@nestjs/throttler@npm:^4.2.1": + version: 4.2.1 + resolution: "@nestjs/throttler@npm:4.2.1" + dependencies: + md5: ^2.2.1 + peerDependencies: + "@nestjs/common": ^7.0.0 || ^8.0.0 || ^9.0.0 || ^10.0.0 + "@nestjs/core": ^7.0.0 || ^8.0.0 || ^9.0.0 || ^10.0.0 + reflect-metadata: ^0.1.13 + checksum: 28b1b1b3ff0623e67ecd51f237bb4a8c22c39134baa5556918a83771c9ebae9b095d0953f3e0e6ba46cc5f7279e4c9cd5c357cbdc00eccffe0e295dc180bcd76 + languageName: node + linkType: hard + "@nestjs/websockets@npm:^10.2.2": version: 10.2.2 resolution: "@nestjs/websockets@npm:10.2.2" @@ -25723,7 +25738,7 @@ __metadata: languageName: node linkType: hard -"md5@npm:^2.3.0": +"md5@npm:^2.2.1, md5@npm:^2.3.0": version: 2.3.0 resolution: "md5@npm:2.3.0" dependencies: @@ -26522,6 +26537,19 @@ __metadata: languageName: node linkType: hard +"nestjs-throttler-storage-redis@npm:^0.3.3": + version: 0.3.3 + resolution: "nestjs-throttler-storage-redis@npm:0.3.3" + peerDependencies: + "@nestjs/common": ">=9.4.1" + "@nestjs/core": ">=9.4.1" + "@nestjs/throttler": ">=4.0.0" + ioredis: ^5.3.2 + reflect-metadata: ^0.1.13 + checksum: 23c14767ea77e850a4656a3a1023315d3144639040e1172bc89a090c09498300fa2b13c18f3ef6aabbdf77a98a3c2b4bd84d2da64397fcfa94457ae525c24c13 + languageName: node + linkType: hard + "next-auth@npm:4.22.5": version: 4.22.5 resolution: "next-auth@npm:4.22.5"