Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 63 additions & 0 deletions migrations/20260428065637-add-max-upload-file-size-limit.js
Original file line number Diff line number Diff line change
@@ -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 } },
);
},
};
12 changes: 12 additions & 0 deletions src/modules/cache-manager/cache-manager.service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
@@ -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,
);
}
Expand Down
2 changes: 2 additions & 0 deletions src/modules/feature-limit/feature-limit.module.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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: [
Expand All @@ -30,6 +31,7 @@ import { WorkspacesModule } from '../workspaces/workspaces.module';
forwardRef(() => SharingModule),
forwardRef(() => UserModule),
WorkspacesModule,
CacheManagerModule,
],
providers: [
SequelizeFeatureLimitsRepository,
Expand Down
152 changes: 152 additions & 0 deletions src/modules/feature-limit/feature-limit.service.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand All @@ -33,6 +36,7 @@ describe('FeatureLimitService', () => {
limitsRepository = moduleRef.get(SequelizeFeatureLimitsRepository);
workspaceRepository = moduleRef.get(SequelizeWorkspaceRepository);
userRepository = moduleRef.get(SequelizeUserRepository);
cacheManagerService = moduleRef.get(CacheManagerService);
});

describe('canUserAccessPlatform', () => {
Expand Down Expand Up @@ -370,6 +374,154 @@ 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 exceeds tier limit, then it should throw', 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();
Expand Down
54 changes: 54 additions & 0 deletions src/modules/feature-limit/feature-limit.service.ts
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -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(
Expand Down Expand Up @@ -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,
);
}
}
}
1 change: 1 addition & 0 deletions src/modules/feature-limit/limits.enum.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ export enum LimitLabels {
RcloneAccess = 'rclone-access',
TrashRetentionDays = 'trash-retention-days',
ReferralAccess = 'referral-access',
MaxUploadFileSize = 'max-upload-file-size',
}

export enum LimitTypes {
Expand Down
Loading
Loading