feat(file-size-limit): limit file size on POST and Replace

Author: apsantiso

migrations/20260428065637-add-max-upload-file-size-limit.js

@@ -0,0 +1,63 @@
+'use strict';
+
+const { v4 } = require('uuid');
+
+const LIMIT_LABEL = 'max-upload-file-size';
+
+const MB = 1024 * 1024;
+const GB = 1024 * MB;
+
+/** @type {import('sequelize-cli').Migration} */
+module.exports = {
+  async up(queryInterface) {
+    await queryInterface.bulkInsert('limits', [
+      {
+        id: v4(),
+        label: LIMIT_LABEL,
+        type: 'counter',
+        value: String(100 * MB),
+        created_at: new Date(),
+        updated_at: new Date(),
+      },
+      {
+        id: v4(),
+        label: LIMIT_LABEL,
+        type: 'counter',
+        value: String(1 * GB),
+        created_at: new Date(),
+        updated_at: new Date(),
+      },
+      {
+        id: v4(),
+        label: LIMIT_LABEL,
+        type: 'counter',
+        value: String(10 * GB),
+        created_at: new Date(),
+        updated_at: new Date(),
+      },
+      {
+        id: v4(),
+        label: LIMIT_LABEL,
+        type: 'counter',
+        value: String(50 * GB),
+        created_at: new Date(),
+        updated_at: new Date(),
+      },
+      {
+        id: v4(),
+        label: LIMIT_LABEL,
+        type: 'counter',
+        value: String(100 * GB),
+        created_at: new Date(),
+        updated_at: new Date(),
+      },
+    ]);
+  },
+
+  async down(queryInterface) {
+    await queryInterface.sequelize.query(
+      `DELETE FROM limits WHERE label = :limitLabel`,
+      { replacements: { limitLabel: LIMIT_LABEL } },
+    );
+  },
+};

src/modules/cache-manager/cache-manager.service.ts

@@ -101,6 +101,18 @@ export class CacheManagerService {
     return this.cacheManager.del(`${this.AVATAR_KEY_PREFIX}${userUuid}`);
   }
 
+  async getTierLimit(tierId: string, label: string) {
+    return this.cacheManager.get<string>(`tier-limit:${tierId}:${label}`);
+  }
+
+  async setTierLimit(tierId: string, label: string, value: string) {
+    await this.cacheManager.set(
+      `tier-limit:${tierId}:${label}`,
+      value,
+      this.TTL_10_MINUTES,
+    );
+  }
+
   async checkHealth(): Promise<void> {
     const key = 'health_check';
     const value = Date.now().toString();

src/modules/feature-limit/exceptions/payment-required.exception.ts

@@ -1,10 +1,19 @@
 import { HttpException, HttpStatus } from '@nestjs/common';
 
+export enum PaymentRequiredErrorCode {
+  FileUploadSizeExceeded = 'FILE_UPLOAD_SIZE_EXCEEDED',
+  FeatureNotAvailable = 'FEATURE_NOT_AVAILABLE',
+}
+
 export class PaymentRequiredException extends HttpException {
-  constructor(message?: string) {
+  constructor(message?: string, code?: PaymentRequiredErrorCode) {
     super(
-      message ??
-        'It seems you reached the limit or feature is not available for your current plan tier',
+      {
+        message:
+          message ??
+          'It seems you reached the limit or feature is not available for your current plan tier',
+        ...(code ? { error: code } : {}),
+      },
       HttpStatus.PAYMENT_REQUIRED,
     );
   }

src/modules/feature-limit/feature-limit.module.ts

@@ -16,6 +16,7 @@ import { PaidPlansModel } from './models/paid-plans.model';
 import { PaymentsService } from '../../externals/payments/payments.service';
 import { FeatureLimitService } from './feature-limit.service';
 import { WorkspacesModule } from '../workspaces/workspaces.module';
+import { CacheManagerModule } from '../cache-manager/cache-manager.module';
 
 @Module({
   imports: [
@@ -30,6 +31,7 @@ import { WorkspacesModule } from '../workspaces/workspaces.module';
     forwardRef(() => SharingModule),
     forwardRef(() => UserModule),
     WorkspacesModule,
+    CacheManagerModule,
   ],
   providers: [
     SequelizeFeatureLimitsRepository,

src/modules/feature-limit/feature-limit.service.spec.ts

@@ -14,12 +14,15 @@ import {
 } from '../../../test/fixtures';
 import { SequelizeWorkspaceRepository } from '../workspaces/repositories/workspaces.repository';
 import { SequelizeUserRepository } from '../user/user.repository';
+import { CacheManagerService } from '../cache-manager/cache-manager.service';
+import { PaymentRequiredException } from './exceptions/payment-required.exception';
 
 describe('FeatureLimitService', () => {
   let service: FeatureLimitService;
   let limitsRepository: DeepMocked<SequelizeFeatureLimitsRepository>;
   let workspaceRepository: DeepMocked<SequelizeWorkspaceRepository>;
   let userRepository: DeepMocked<SequelizeUserRepository>;
+  let cacheManagerService: DeepMocked<CacheManagerService>;
 
   beforeEach(async () => {
     const moduleRef = await Test.createTestingModule({
@@ -33,6 +36,7 @@ describe('FeatureLimitService', () => {
     limitsRepository = moduleRef.get(SequelizeFeatureLimitsRepository);
     workspaceRepository = moduleRef.get(SequelizeWorkspaceRepository);
     userRepository = moduleRef.get(SequelizeUserRepository);
+    cacheManagerService = moduleRef.get(CacheManagerService);
   });
 
   describe('canUserAccessPlatform', () => {
@@ -370,6 +374,171 @@ describe('FeatureLimitService', () => {
     });
   });
 
+  describe('enforceMaxUploadFileSize', () => {
+    const MB = 1024 * 1024;
+    const GB = 1024 * MB;
+
+    it('When user has no limit configured, then it should allow the upload', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(null);
+      cacheManagerService.getTierLimit.mockResolvedValueOnce(null);
+      limitsRepository.findLimitByLabelAndTier.mockResolvedValueOnce(null);
+
+      await expect(
+        service.enforceMaxUploadFileSize(user, BigInt(500 * MB)),
+      ).resolves.not.toThrow();
+    });
+
+    it('When file size is under tier limit, then it should allow the upload', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+      const limit = newFeatureLimit({
+        type: LimitTypes.Counter,
+        label: LimitLabels.MaxUploadFileSize,
+        value: String(100 * MB),
+      });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(null);
+      cacheManagerService.getTierLimit.mockResolvedValueOnce(null);
+      limitsRepository.findLimitByLabelAndTier.mockResolvedValueOnce(limit);
+
+      await expect(
+        service.enforceMaxUploadFileSize(user, BigInt(50 * MB)),
+      ).resolves.not.toThrow();
+    });
+
+    it('When file size equals tier limit, then it should throw PaymentRequiredException', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+      const limit = newFeatureLimit({
+        type: LimitTypes.Counter,
+        label: LimitLabels.MaxUploadFileSize,
+        value: String(100 * MB),
+      });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(null);
+      cacheManagerService.getTierLimit.mockResolvedValueOnce(null);
+      limitsRepository.findLimitByLabelAndTier.mockResolvedValueOnce(limit);
+
+      await expect(
+        service.enforceMaxUploadFileSize(user, BigInt(100 * MB)),
+      ).rejects.toThrow(PaymentRequiredException);
+    });
+
+    it('When file size exceeds tier limit, then it should throw PaymentRequiredException', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+      const limit = newFeatureLimit({
+        type: LimitTypes.Counter,
+        label: LimitLabels.MaxUploadFileSize,
+        value: String(100 * MB),
+      });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(null);
+      cacheManagerService.getTierLimit.mockResolvedValueOnce(null);
+      limitsRepository.findLimitByLabelAndTier.mockResolvedValueOnce(limit);
+
+      await expect(
+        service.enforceMaxUploadFileSize(user, BigInt(200 * MB)),
+      ).rejects.toThrow(PaymentRequiredException);
+    });
+
+    it('When user has an overridden limit, then it should use it instead of tier limit', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+      const override = newFeatureLimit({
+        type: LimitTypes.Counter,
+        label: LimitLabels.MaxUploadFileSize,
+        value: String(10 * GB),
+      });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(override);
+
+      await expect(
+        service.enforceMaxUploadFileSize(user, BigInt(5 * GB)),
+      ).resolves.not.toThrow();
+      expect(cacheManagerService.getTierLimit).not.toHaveBeenCalled();
+      expect(limitsRepository.findLimitByLabelAndTier).not.toHaveBeenCalled();
+    });
+
+    it('When user overridden limit is exceeded, then it should throw', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+      const override = newFeatureLimit({
+        type: LimitTypes.Counter,
+        label: LimitLabels.MaxUploadFileSize,
+        value: String(100 * MB),
+      });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(override);
+
+      await expect(
+        service.enforceMaxUploadFileSize(user, BigInt(200 * MB)),
+      ).rejects.toThrow(PaymentRequiredException);
+    });
+
+    it('When tier limit is cached, then it should not hit the DB', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(null);
+      cacheManagerService.getTierLimit.mockResolvedValueOnce(String(1 * GB));
+
+      await expect(
+        service.enforceMaxUploadFileSize(user, BigInt(500 * MB)),
+      ).resolves.not.toThrow();
+      expect(limitsRepository.findLimitByLabelAndTier).not.toHaveBeenCalled();
+    });
+
+    it('When tier limit is cached and exceeded, then it should throw without hitting DB', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(null);
+      cacheManagerService.getTierLimit.mockResolvedValueOnce(String(100 * MB));
+
+      await expect(
+        service.enforceMaxUploadFileSize(user, BigInt(200 * MB)),
+      ).rejects.toThrow(PaymentRequiredException);
+      expect(limitsRepository.findLimitByLabelAndTier).not.toHaveBeenCalled();
+    });
+
+    it('When cache miss occurs, then it should populate the cache from DB', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+      const limit = newFeatureLimit({
+        type: LimitTypes.Counter,
+        label: LimitLabels.MaxUploadFileSize,
+        value: String(1 * GB),
+      });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(null);
+      cacheManagerService.getTierLimit.mockResolvedValueOnce(null);
+      limitsRepository.findLimitByLabelAndTier.mockResolvedValueOnce(limit);
+
+      await service.enforceMaxUploadFileSize(user, BigInt(500 * MB));
+
+      expect(cacheManagerService.setTierLimit).toHaveBeenCalledWith(
+        user.tierId,
+        LimitLabels.MaxUploadFileSize,
+        limit.value,
+      );
+    });
+
+    it('When cache write fails, then it should still allow the upload', async () => {
+      const user = newUser({ attributes: { tierId: v4() } });
+      const limit = newFeatureLimit({
+        type: LimitTypes.Counter,
+        label: LimitLabels.MaxUploadFileSize,
+        value: String(1 * GB),
+      });
+
+      limitsRepository.findUserOverriddenLimit.mockResolvedValueOnce(null);
+      cacheManagerService.getTierLimit.mockResolvedValueOnce(null);
+      limitsRepository.findLimitByLabelAndTier.mockResolvedValueOnce(limit);
+      cacheManagerService.setTierLimit.mockRejectedValueOnce(
+        new Error('Redis unavailable'),
+      );
+
+      await expect(
+        service.enforceMaxUploadFileSize(user, BigInt(500 * MB)),
+      ).resolves.not.toThrow();
+    });
+  });
+
   describe('getTier', () => {
     it('When tier exists, then it should return the tier', async () => {
       const tierId = v4();

src/modules/feature-limit/feature-limit.service.ts

@@ -1,11 +1,16 @@
 import { Injectable, Logger, NotFoundException } from '@nestjs/common';
+import {
+  PaymentRequiredException,
+  PaymentRequiredErrorCode,
+} from './exceptions/payment-required.exception';
 import { SequelizeFeatureLimitsRepository } from './feature-limit.repository';
 import { LimitLabels } from './limits.enum';
 import { PlatformName } from '../../common/constants';
 import { SequelizeWorkspaceRepository } from '../workspaces/repositories/workspaces.repository';
 import { SequelizeUserRepository } from '../user/user.repository';
 import { type Limit } from './domain/limit.domain';
 import { type User } from '../user/user.domain';
+import { CacheManagerService } from '../cache-manager/cache-manager.service';
 
 @Injectable()
 export class FeatureLimitService {
@@ -15,6 +20,7 @@ export class FeatureLimitService {
     private readonly limitsRepository: SequelizeFeatureLimitsRepository,
     private readonly workspaceRepository: SequelizeWorkspaceRepository,
     private readonly userRepository: SequelizeUserRepository,
+    private readonly cacheManagerService: CacheManagerService,
   ) {}
 
   async canUserAccessPlatform(
@@ -170,4 +176,52 @@ export class FeatureLimitService {
 
     return userOverriddenLimits ?? tierLimits;
   }
+
+  async enforceMaxUploadFileSize(user: User, fileSize: bigint): Promise<void> {
+    const userOverriddenLimit =
+      await this.limitsRepository.findUserOverriddenLimit(
+        user.uuid,
+        LimitLabels.MaxUploadFileSize,
+      );
+
+    let limitValue: string | null = null;
+
+    if (userOverriddenLimit) {
+      limitValue = userOverriddenLimit.value;
+    } else {
+      const cached = await this.cacheManagerService.getTierLimit(
+        user.tierId,
+        LimitLabels.MaxUploadFileSize,
+      );
+
+      if (cached !== null && cached !== undefined) {
+        limitValue = cached;
+      } else {
+        const tierLimit = await this.limitsRepository.findLimitByLabelAndTier(
+          user.tierId,
+          LimitLabels.MaxUploadFileSize,
+        );
+
+        if (!tierLimit) return;
+
+        limitValue = tierLimit.value;
+        await this.cacheManagerService
+          .setTierLimit(user.tierId, LimitLabels.MaxUploadFileSize, limitValue)
+          .catch((err) => {
+            this.logger.error(
+              `Failed to cache tier limit for tierId ${user.tierId}: ${err.message}`,
+            );
+          });
+      }
+    }
+
+    if (limitValue === null) return;
+
+    if (Number(fileSize) >= Number(limitValue)) {
+      throw new PaymentRequiredException(
+        'File size exceeds the maximum allowed by your plan',
+        PaymentRequiredErrorCode.FileUploadSizeExceeded,
+      );
+    }
+  }
 }