Skip to content

Commit

Permalink
fix(server): skip throttle for currentUser (toeverything#6700)
Browse files Browse the repository at this point in the history
  • Loading branch information
forehalo committed Apr 25, 2024
1 parent 6237bf1 commit 3297486
Show file tree
Hide file tree
Showing 5 changed files with 45 additions and 21 deletions.
9 changes: 2 additions & 7 deletions packages/backend/server/src/core/auth/resolver.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ import {
} from '@nestjs/graphql';
import type { Request, Response } from 'express';

import { Config, Throttle } from '../../fundamentals';
import { Config, SkipThrottle, Throttle } from '../../fundamentals';
import { UserService } from '../user';
import { UserType } from '../user/types';
import { validators } from '../utils/validators';
Expand All @@ -33,12 +33,6 @@ export class ClientTokenType {
sessionToken?: string;
}

/**
* Auth resolver
* Token rate limit: 20 req/m
* Sign up/in rate limit: 10 req/m
* Other rate limit: 5 req/m
*/
@Throttle('strict')
@Resolver(() => UserType)
export class AuthResolver {
Expand All @@ -49,6 +43,7 @@ export class AuthResolver {
private readonly token: TokenService
) {}

@SkipThrottle()
@Public()
@Query(() => UserType, {
name: 'currentUser',
Expand Down
2 changes: 1 addition & 1 deletion packages/backend/server/src/fundamentals/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ export {
export type { PrismaTransaction } from './prisma';
export * from './storage';
export { type StorageProvider, StorageProviderFactory } from './storage';
export { CloudThrottlerGuard, Throttle } from './throttler';
export { CloudThrottlerGuard, SkipThrottle, Throttle } from './throttler';
export {
getRequestFromHost,
getRequestResponseFromContext,
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { applyDecorators, SetMetadata } from '@nestjs/common';
import { SkipThrottle, Throttle as RawThrottle } from '@nestjs/throttler';

export type Throttlers = 'default' | 'strict';
export type Throttlers = 'default' | 'strict' | 'authenticated';
export const THROTTLER_PROTECTED = 'affine_throttler:protected';

/**
Expand All @@ -10,8 +10,9 @@ export const THROTTLER_PROTECTED = 'affine_throttler:protected';
* If a Controller or Query do not protected behind a Throttler,
* it will never be rate limited.
*
* - Ease: 120 calls within 60 seconds
* - Strict: 10 calls within 60 seconds
* - default: 120 calls within 60 seconds
* - strict: 10 calls within 60 seconds
* - authenticated: no rate limit for authenticated users, apply [default] throttler for unauthenticated users
*
* @example
*
Expand Down
4 changes: 3 additions & 1 deletion packages/backend/server/src/fundamentals/throttler/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -166,10 +166,12 @@ export class CloudThrottlerGuard extends ThrottlerGuard {
}

getSpecifiedThrottler(context: ExecutionContext) {
return this.reflector.getAllAndOverride<Throttlers | undefined>(
const throttler = this.reflector.getAllAndOverride<Throttlers | undefined>(
THROTTLER_PROTECTED,
[context.getHandler(), context.getClass()]
);

return throttler === 'authenticated' ? undefined : throttler;
}
}

Expand Down
44 changes: 35 additions & 9 deletions packages/backend/server/tests/nestjs/throttler.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,13 @@ class ThrottledController {
return 'default3';
}

@Public()
@Get('/authenticated')
@Throttle('authenticated')
none() {
return 'none';
}

@Throttle('strict')
@Get('/strict')
strict() {
Expand Down Expand Up @@ -156,7 +163,6 @@ test('should use default throttler for unauthenticated user when not specified',

t.is(headers.limit, '120');
t.is(headers.remaining, '119');
t.regex(headers.reset, /59|60/);
});

test('should skip throttler for unauthenticated user when specified', async t => {
Expand Down Expand Up @@ -192,7 +198,6 @@ test('should use specified throttler for unauthenticated user', async t => {

t.is(headers.limit, '20');
t.is(headers.remaining, '19');
t.regex(headers.reset, /59|60/);
});

// ==== authenticated user visits ====
Expand Down Expand Up @@ -223,7 +228,6 @@ test('should use default throttler for authenticated user when not specified', a

t.is(headers.limit, '120');
t.is(headers.remaining, '119');
t.regex(headers.reset, /59|60/);
});

test('should use same throttler for multiple routes', async t => {
Expand All @@ -238,7 +242,6 @@ test('should use same throttler for multiple routes', async t => {

t.is(headers.limit, '120');
t.is(headers.remaining, '119');
t.regex(headers.reset, /59|60/);

res = await request(app.getHttpServer())
.get('/throttled/default2')
Expand All @@ -263,7 +266,6 @@ test('should use different throttler if specified', async t => {

t.is(headers.limit, '120');
t.is(headers.remaining, '119');
t.regex(headers.reset, /59|60/);

res = await request(app.getHttpServer())
.get('/throttled/default3')
Expand All @@ -274,7 +276,34 @@ test('should use different throttler if specified', async t => {

t.is(headers.limit, '10');
t.is(headers.remaining, '9');
t.regex(headers.reset, /59|60/);
});

test('should skip throttler for authenticated if `authenticated` throttler used', async t => {
const { app, cookie } = t.context;

const res = await request(app.getHttpServer())
.get('/throttled/authenticated')
.set('Cookie', cookie)
.expect(200);

const headers = rateLimitHeaders(res);

t.is(headers.limit, undefined!);
t.is(headers.remaining, undefined!);
t.is(headers.reset, undefined!);
});

test('should apply `default` throttler for authenticated user if `authenticated` throttler used', async t => {
const { app } = t.context;

const res = await request(app.getHttpServer())
.get('/throttled/authenticated')
.expect(200);

const headers = rateLimitHeaders(res);

t.is(headers.limit, '120');
t.is(headers.remaining, '119');
});

test('should skip throttler for authenticated user when specified', async t => {
Expand Down Expand Up @@ -304,7 +333,6 @@ test('should use specified throttler for authenticated user', async t => {

t.is(headers.limit, '20');
t.is(headers.remaining, '19');
t.regex(headers.reset, /59|60/);
});

test('should separate anonymous and authenticated user throttlers', async t => {
Expand All @@ -323,9 +351,7 @@ test('should separate anonymous and authenticated user throttlers', async t => {

t.is(authenticatedResHeaders.limit, '120');
t.is(authenticatedResHeaders.remaining, '119');
t.regex(authenticatedResHeaders.reset, /59|60/);

t.is(unauthenticatedResHeaders.limit, '120');
t.is(unauthenticatedResHeaders.remaining, '119');
t.regex(unauthenticatedResHeaders.reset, /59|60/);
});

0 comments on commit 3297486

Please sign in to comment.