Throttling in common api (#15338)

2025-10-24 17:32:51 +02:00
parent d7bda9576f
commit 93ab1c4118
11 changed files with 295 additions and 14 deletions
@@ -31,6 +31,8 @@ import { WorkspacePreQueryHookPayload } from 'src/engine/api/graphql/workspace-q
 import { WorkspaceQueryHookService } from 'src/engine/api/graphql/workspace-query-runner/workspace-query-hook/workspace-query-hook.service';
 import { ApiKeyRoleService } from 'src/engine/core-modules/api-key/api-key-role.service';
 import { AuthContext } from 'src/engine/core-modules/auth/types/auth-context.type';
+import { ThrottlerService } from 'src/engine/core-modules/throttler/throttler.service';
+import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';
 import { type PermissionFlagType } from 'src/engine/metadata-modules/permissions/constants/permission-flag-type.constants';
 import {
  PermissionsException,
@@ -71,6 +73,10 @@ export abstract class CommonBaseQueryRunnerService<
  protected readonly workspacePermissionsCacheService: WorkspacePermissionsCacheService;
  @Inject()
  protected readonly commonResultGettersService: CommonResultGettersService;
+  @Inject()
+  protected readonly throttlerService: ThrottlerService;
+  @Inject()
+  protected readonly twentyConfigService: TwentyConfigService;

  protected abstract readonly operationName: CommonQueryNames;

@@ -109,6 +115,8 @@ export abstract class CommonBaseQueryRunnerService<
      commonQueryParser,
    );

+    await this.throttleQueryExecution(authContext.workspace.id);
+
    const extendedQueryRunnerContext =
      await this.prepareExtendedQueryRunnerContext(
        authContext,
@@ -322,4 +330,36 @@ export abstract class CommonBaseQueryRunnerService<
      repository,
    };
  }
+
+  private async throttleQueryExecution(workspaceId: string) {
+    const shortConfig = {
+      key: `api:throttler:${workspaceId}-short-limit`,
+      maxTokens: this.twentyConfigService.get('API_RATE_LIMITING_SHORT_LIMIT'),
+      timeWindow: this.twentyConfigService.get(
+        'API_RATE_LIMITING_SHORT_TTL_IN_MS',
+      ),
+    };
+
+    const longConfig = {
+      key: `api:throttler:${workspaceId}-long-limit`,
+      maxTokens: this.twentyConfigService.get('API_RATE_LIMITING_LONG_LIMIT'),
+      timeWindow: this.twentyConfigService.get(
+        'API_RATE_LIMITING_LONG_TTL_IN_MS',
+      ),
+    };
+
+    await this.throttlerService.tokenBucketThrottle(
+      shortConfig.key,
+      1,
+      shortConfig.maxTokens,
+      shortConfig.timeWindow,
+    );
+
+    await this.throttlerService.tokenBucketThrottle(
+      longConfig.key,
+      1,
+      longConfig.maxTokens,
+      longConfig.timeWindow,
+    );
+  }
 }
@@ -12,6 +12,7 @@ import { WorkspaceQueryHookModule } from 'src/engine/api/graphql/workspace-query
 import { WorkspaceQueryRunnerModule } from 'src/engine/api/graphql/workspace-query-runner/workspace-query-runner.module';
 import { ApiKeyModule } from 'src/engine/core-modules/api-key/api-key.module';
 import { FileModule } from 'src/engine/core-modules/file/file.module';
+import { ThrottlerModule } from 'src/engine/core-modules/throttler/throttler.module';
 import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module';
 import { RoleTargetsEntity } from 'src/engine/metadata-modules/role/role-targets.entity';
 import { UserRoleModule } from 'src/engine/metadata-modules/user-role/user-role.module';
@@ -33,6 +34,7 @@ import { WorkspacePermissionsCacheModule } from 'src/engine/metadata-modules/wor
    ViewModule,
    ViewFilterModule,
    ViewFilterGroupModule,
+    ThrottlerModule,
  ],
  providers: [
    ProcessNestedRelationsHelper,
@@ -17,7 +17,6 @@ import { JsonWebTokenError, TokenExpiredError } from 'jsonwebtoken';

 import { NodeEnvironment } from 'src/engine/core-modules/twenty-config/interfaces/node-environment.interface';

-import { useThrottler } from 'src/engine/api/graphql/graphql-config/hooks/use-throttler';
 import { WorkspaceSchemaFactory } from 'src/engine/api/graphql/workspace-schema.factory';
 import { type AuthContext } from 'src/engine/core-modules/auth/types/auth-context.type';
 import { CoreEngineModule } from 'src/engine/core-modules/core-engine.module';
@@ -55,13 +54,6 @@ export class GraphQLConfigService
    const isDebugMode =
      this.twentyConfigService.get('NODE_ENV') === NodeEnvironment.DEVELOPMENT;
    const plugins = [
-      useThrottler({
-        ttl: this.twentyConfigService.get('API_RATE_LIMITING_TTL'),
-        limit: this.twentyConfigService.get('API_RATE_LIMITING_LIMIT'),
-        identifyFn: (context) => {
-          return context.req.user?.id ?? context.req.ip ?? 'anonymous';
-        },
-      }),
      useGraphQLErrorHandlerHook({
        metricsService: this.metricsService,
        exceptionHandlerService: this.exceptionHandlerService,
@@ -32,8 +32,8 @@ export const metadataModuleFactory = async (
    resolvers: { JSON: GraphQLJSON },
    plugins: [
      useThrottler({
-        ttl: twentyConfigService.get('API_RATE_LIMITING_TTL'),
-        limit: twentyConfigService.get('API_RATE_LIMITING_LIMIT'),
+        ttl: twentyConfigService.get('API_RATE_LIMITING_LONG_TTL_IN_MS') / 1000,
+        limit: twentyConfigService.get('API_RATE_LIMITING_LONG_LIMIT'),
        identifyFn: (context) => {
          return context.req.user?.id ?? context.req.ip ?? 'anonymous';
        },
@@ -12,6 +12,8 @@ import { AuthException } from 'src/engine/core-modules/auth/auth.exception';
 import { authGraphqlApiExceptionHandler } from 'src/engine/core-modules/auth/utils/auth-graphql-api-exception-handler.util';
 import { RecordTransformerException } from 'src/engine/core-modules/record-transformer/record-transformer.exception';
 import { recordTransformerGraphqlApiExceptionHandler } from 'src/engine/core-modules/record-transformer/utils/record-transformer-graphql-api-exception-handler.util';
+import { ThrottlerException } from 'src/engine/core-modules/throttler/throttler.exception';
+import { throttlerToGraphqlApiExceptionHandler } from 'src/engine/core-modules/throttler/utils/throttler-to-graphql-api-exception-handler.util';
 import { PermissionsException } from 'src/engine/metadata-modules/permissions/permissions.exception';
 import { permissionGraphqlApiExceptionHandler } from 'src/engine/metadata-modules/permissions/utils/permission-graphql-api-exception-handler.util';
 import { TwentyORMException } from 'src/engine/twenty-orm/exceptions/twenty-orm.exception';
@@ -42,6 +44,8 @@ export const workspaceQueryRunnerGraphqlApiExceptionHandler = (
      return authGraphqlApiExceptionHandler(error);
    case error instanceof ApiKeyException:
      return apiKeyGraphqlApiExceptionHandler(error);
+    case error instanceof ThrottlerException:
+      return throttlerToGraphqlApiExceptionHandler(error);
    default:
      throw error;
  }
@@ -5,6 +5,8 @@ import { type QueryFailedError } from 'typeorm';
 import { CommonQueryRunnerException } from 'src/engine/api/common/common-query-runners/errors/common-query-runner.exception';
 import { commonQueryRunnerToRestApiExceptionHandler } from 'src/engine/api/common/common-query-runners/utils/common-query-runner-to-rest-api-exception-handler.util';
 import { RestInputRequestParserException } from 'src/engine/api/rest/input-request-parsers/rest-input-request-parser.exception';
+import { ThrottlerException } from 'src/engine/core-modules/throttler/throttler.exception';
+import { throttlerToRestApiExceptionHandler } from 'src/engine/core-modules/throttler/utils/throttler-to-rest-api-exception-handler.util';

 interface QueryFailedErrorWithCode extends QueryFailedError {
  code: string;
@@ -18,6 +20,8 @@ export const workspaceQueryRunnerRestApiExceptionHandler = (
      return commonQueryRunnerToRestApiExceptionHandler(error);
    case error instanceof RestInputRequestParserException:
      throw new BadRequestException(error.message);
+    case error instanceof ThrottlerException:
+      return throttlerToRestApiExceptionHandler(error);
    default:
      throw error;
  }
@@ -0,0 +1,148 @@
+import { Test, type TestingModule } from '@nestjs/testing';
+
+import { type CacheStorageService } from 'src/engine/core-modules/cache-storage/services/cache-storage.service';
+import { CacheStorageNamespace } from 'src/engine/core-modules/cache-storage/types/cache-storage-namespace.enum';
+import { ThrottlerException } from 'src/engine/core-modules/throttler/throttler.exception';
+import { ThrottlerService } from 'src/engine/core-modules/throttler/throttler.service';
+
+describe('ThrottlerService', () => {
+  let service: ThrottlerService;
+  let cacheStorageService: jest.Mocked<CacheStorageService>;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [
+        ThrottlerService,
+        {
+          provide: CacheStorageNamespace.EngineWorkspace,
+          useValue: {
+            get: jest.fn(),
+            set: jest.fn(),
+          },
+        },
+      ],
+    }).compile();
+
+    service = module.get<ThrottlerService>(ThrottlerService);
+    cacheStorageService = module.get(CacheStorageNamespace.EngineWorkspace);
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe('tokenBucketThrottle', () => {
+    const key = 'test-throttle-key';
+    const maxTokens = 100;
+    const timeWindow = 1000; // 1 second
+
+    it('should allow request when tokens are available (first request)', async () => {
+      cacheStorageService.get.mockResolvedValue(null);
+
+      await service.tokenBucketThrottle(key, 10, maxTokens, timeWindow);
+
+      expect(cacheStorageService.get).toHaveBeenCalledWith(key);
+      expect(cacheStorageService.set).toHaveBeenCalledWith(
+        key,
+        {
+          tokens: 90, // maxTokens - tokensToConsume
+          lastRefillAt: expect.any(Number),
+        },
+        timeWindow * 2,
+      );
+    });
+
+    it('should allow request when sufficient tokens are available', async () => {
+      const now = Date.now();
+
+      cacheStorageService.get.mockResolvedValue({
+        tokens: 50,
+        lastRefillAt: now - 100,
+      });
+
+      await service.tokenBucketThrottle(key, 10, maxTokens, timeWindow);
+
+      expect(cacheStorageService.get).toHaveBeenCalledWith(key);
+      expect(cacheStorageService.set).toHaveBeenCalledWith(
+        key,
+        {
+          tokens: expect.any(Number),
+          lastRefillAt: expect.any(Number),
+        },
+        timeWindow * 2,
+      );
+    });
+
+    it('should throw ThrottlerException when tokens are insufficient', async () => {
+      const now = Date.now();
+
+      cacheStorageService.get.mockResolvedValue({
+        tokens: 5,
+        lastRefillAt: now,
+      });
+
+      await expect(
+        service.tokenBucketThrottle(key, 10, maxTokens, timeWindow),
+      ).rejects.toThrow(ThrottlerException);
+
+      await expect(
+        service.tokenBucketThrottle(key, 10, maxTokens, timeWindow),
+      ).rejects.toThrow('Limit reached');
+
+      expect(cacheStorageService.set).not.toHaveBeenCalled();
+    });
+
+    it('should refill tokens over time', async () => {
+      const now = Date.now();
+      const refillRate = maxTokens / timeWindow; // 100 tokens per 1000ms = 0.1 tokens/ms
+      const timePassed = 500; // 500ms
+      const expectedRefill = Math.floor(timePassed * refillRate); // 50 tokens
+
+      cacheStorageService.get.mockResolvedValue({
+        tokens: 20,
+        lastRefillAt: now - timePassed,
+      });
+
+      jest.spyOn(Date, 'now').mockReturnValue(now);
+
+      await service.tokenBucketThrottle(key, 10, maxTokens, timeWindow);
+
+      expect(cacheStorageService.set).toHaveBeenCalledWith(
+        key,
+        {
+          tokens: 20 + expectedRefill - 10, // 20 + 50 - 10 = 60
+          lastRefillAt: now,
+        },
+        timeWindow * 2,
+      );
+
+      jest.restoreAllMocks();
+    });
+
+    it('should cap tokens at maxTokens', async () => {
+      const now = Date.now();
+      const timePassed = 5000; // Long time passed, would refill beyond max
+
+      cacheStorageService.get.mockResolvedValue({
+        tokens: 80,
+        lastRefillAt: now - timePassed,
+      });
+
+      jest.spyOn(Date, 'now').mockReturnValue(now);
+
+      await service.tokenBucketThrottle(key, 10, maxTokens, timeWindow);
+
+      // Available tokens should be capped at maxTokens (100)
+      expect(cacheStorageService.set).toHaveBeenCalledWith(
+        key,
+        {
+          tokens: 90, // maxTokens (100) - tokensToConsume (10)
+          lastRefillAt: now,
+        },
+        timeWindow * 2,
+      );
+
+      jest.restoreAllMocks();
+    });
+  });
+});
@@ -27,4 +27,39 @@ export class ThrottlerService {

    await this.cacheStorage.set(key, currentCount + 1, ttl);
  }
+
+  async tokenBucketThrottle(
+    key: string,
+    tokensToConsume: number,
+    maxTokens: number,
+    timeWindow: number,
+  ): Promise<void> {
+    const now = Date.now();
+    const refillRate = maxTokens / timeWindow;
+
+    const { tokens, lastRefillAt } = (await this.cacheStorage.get<{
+      tokens: number;
+      lastRefillAt: number;
+    }>(key)) || { tokens: maxTokens, lastRefillAt: now };
+
+    const refillAmount = Math.floor((now - lastRefillAt) * refillRate);
+
+    const availableTokens = Math.min(tokens + refillAmount, maxTokens);
+
+    if (availableTokens < tokensToConsume) {
+      throw new ThrottlerException(
+        `Limit reached (${maxTokens} tokens per ${timeWindow} ms)`,
+        ThrottlerExceptionCode.LIMIT_REACHED,
+      );
+    }
+
+    await this.cacheStorage.set(
+      key,
+      {
+        tokens: availableTokens - tokensToConsume,
+        lastRefillAt: now,
+      },
+      timeWindow * 2,
+    );
+  }
 }
@@ -0,0 +1,19 @@
+import { assertUnreachable } from 'twenty-shared/utils';
+
+import { UserInputError } from 'src/engine/core-modules/graphql/utils/graphql-errors.util';
+import {
+  type ThrottlerException,
+  ThrottlerExceptionCode,
+} from 'src/engine/core-modules/throttler/throttler.exception';
+
+export const throttlerToGraphqlApiExceptionHandler = (
+  error: ThrottlerException,
+) => {
+  switch (error.code) {
+    case ThrottlerExceptionCode.LIMIT_REACHED:
+      throw new UserInputError(error);
+    default: {
+      return assertUnreachable(error.code);
+    }
+  }
+};
@@ -0,0 +1,20 @@
+import { HttpException, HttpStatus } from '@nestjs/common';
+
+import { assertUnreachable } from 'twenty-shared/utils';
+
+import {
+  type ThrottlerException,
+  ThrottlerExceptionCode,
+} from 'src/engine/core-modules/throttler/throttler.exception';
+
+export const throttlerToRestApiExceptionHandler = (
+  error: ThrottlerException,
+): never => {
+  switch (error.code) {
+    case ThrottlerExceptionCode.LIMIT_REACHED:
+      throw new HttpException(error.message, HttpStatus.TOO_MANY_REQUESTS);
+    default: {
+      return assertUnreachable(error.code);
+    }
+  }
+};
@@ -945,20 +945,37 @@ export class ConfigVariables {

  @ConfigVariablesMetadata({
    group: ConfigVariablesGroup.RATE_LIMITING,
-    description: 'Time-to-live for API rate limiting in milliseconds',
+    description: 'Time-to-live for short API rate limiting in milliseconds',
    type: ConfigVariableType.NUMBER,
  })
  @CastToPositiveNumber()
-  API_RATE_LIMITING_TTL = 100;
+  API_RATE_LIMITING_SHORT_TTL_IN_MS = 1000;

  @ConfigVariablesMetadata({
    group: ConfigVariablesGroup.RATE_LIMITING,
    description:
-      'Maximum number of requests allowed in the rate limiting window',
+      'Maximum number of requests allowed in the short rate limiting window',
    type: ConfigVariableType.NUMBER,
  })
  @CastToPositiveNumber()
-  API_RATE_LIMITING_LIMIT = 500;
+  API_RATE_LIMITING_SHORT_LIMIT = 100;
+
+  @ConfigVariablesMetadata({
+    group: ConfigVariablesGroup.RATE_LIMITING,
+    description: 'Time-to-live for long API rate limiting in milliseconds',
+    type: ConfigVariableType.NUMBER,
+  })
+  @CastToPositiveNumber()
+  API_RATE_LIMITING_LONG_TTL_IN_MS = 60000;
+
+  @ConfigVariablesMetadata({
+    group: ConfigVariablesGroup.RATE_LIMITING,
+    description:
+      'Maximum number of requests allowed in the long rate limiting window',
+    type: ConfigVariableType.NUMBER,
+  })
+  @CastToPositiveNumber()
+  API_RATE_LIMITING_LONG_LIMIT = 1000;

  @ConfigVariablesMetadata({
    group: ConfigVariablesGroup.SSL,