feat: rate limiter (#4011)

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],

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
    *

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',

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;
 

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,

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,

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<DeleteAccount> {
     await this.prisma.user.delete({
@@ -172,6 +183,7 @@ export class UserResolver {
     };
   }
 
+  @Throttle(10, 60)
   @Mutation(() => AddToNewFeaturesWaitingList)
   async addToNewFeaturesWaitingList(
     @CurrentUser() user: UserType,

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',

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);

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));
   });
 });

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 {

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');
 

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, any>): string {
+    return req?.get('CF-Connecting-IP') ?? req?.get('CF-ray') ?? req?.ip;
+  }
+}
+
+export { Throttle };