From b86a26668f1955444e469b7bb212b356321218f1 Mon Sep 17 00:00:00 2001 From: AI Builder Date: Sat, 25 Jul 2026 14:17:59 +0100 Subject: [PATCH] feat: add TenantLimitGuard to enforce user and storage limits on tenant resources, return 402 when exceeded --- src/auth/auth.controller.ts | 52 ++++++ src/auth/auth.module.ts | 2 + src/cdn/cdn.controller.ts | 4 + src/cdn/cdn.module.ts | 2 + src/tenancy/guards/tenant-limit.guard.spec.ts | 169 ++++++++++++++++++ src/tenancy/guards/tenant-limit.guard.ts | 84 +++++++++ src/tenancy/tenancy.module.ts | 3 + 7 files changed, 316 insertions(+) create mode 100644 src/tenancy/guards/tenant-limit.guard.spec.ts create mode 100644 src/tenancy/guards/tenant-limit.guard.ts diff --git a/src/auth/auth.controller.ts b/src/auth/auth.controller.ts index 1685e582..f723edc2 100644 --- a/src/auth/auth.controller.ts +++ b/src/auth/auth.controller.ts @@ -7,16 +7,20 @@ import { HttpStatus, Req, UseGuards, + ConflictException, } from '@nestjs/common'; import { AuthService } from './auth.service'; import { MfaService } from './mfa/mfa.service'; import { LoginDto } from './dto/login.dto'; +import { RegisterDto } from './dto/register.dto'; import { RefreshTokenDto } from './dto/refresh-token.dto'; import { InjectRepository } from '@nestjs/typeorm'; import { User } from '../users/entities/user.entity'; import { Repository } from 'typeorm'; import * as bcrypt from 'bcrypt'; import { JwtAuthGuard } from './guards/jwt-auth.guard'; +import { TenantLimitGuard, LimitType } from '../tenancy/guards/tenant-limit.guard'; +import { TenancyService } from '../tenancy/tenancy.service'; import { ApiTags, ApiOperation, ApiResponse, ApiBearerAuth } from '@nestjs/swagger'; @ApiTags('Auth') @@ -27,8 +31,56 @@ export class AuthController { @InjectRepository(User) private readonly userRepository: Repository, private readonly mfaService: MfaService, + private readonly tenancyService: TenancyService, ) {} + @Post('register') + @HttpCode(HttpStatus.CREATED) + @LimitType('user') + @UseGuards(TenantLimitGuard) + @ApiOperation({ summary: 'Register a new user account' }) + @ApiResponse({ status: 201, description: 'User registered successfully' }) + @ApiResponse({ status: 402, description: 'Tenant user limit exceeded' }) + @ApiResponse({ status: 409, description: 'Email or username already exists' }) + async register(@Body() registerDto: RegisterDto, @Req() req: any) { + const existingEmail = await this.userRepository.findOne({ + where: { email: registerDto.email }, + }); + if (existingEmail) { + throw new ConflictException('Email already registered'); + } + + const existingUsername = await this.userRepository.findOne({ + where: { username: registerDto.username }, + }); + if (existingUsername) { + throw new ConflictException('Username already taken'); + } + + const rounds = Number(process.env.BCRYPT_ROUNDS || 12); + const salt = await bcrypt.genSalt(rounds); + const passwordHash = await bcrypt.hash(registerDto.password, salt); + + const user = this.userRepository.create({ + email: registerDto.email, + username: registerDto.username, + firstName: registerDto.firstName, + lastName: registerDto.lastName, + displayName: registerDto.displayName || registerDto.username, + profilePicture: registerDto.avatarUrl, + password: passwordHash, + tenantId: req.tenantId, + }); + + const savedUser = await this.userRepository.save(user); + + if (req.tenantId) { + await this.tenancyService.incrementUserCount(req.tenantId); + } + + return this.authService.login(savedUser); + } + @Post('login') @HttpCode(HttpStatus.OK) @ApiOperation({ summary: 'Log in with email and password' }) diff --git a/src/auth/auth.module.ts b/src/auth/auth.module.ts index b8aedfb7..b8519080 100644 --- a/src/auth/auth.module.ts +++ b/src/auth/auth.module.ts @@ -18,6 +18,7 @@ import { AuthTokensService } from './services/auth-tokens.service'; import { MfaService } from './mfa/mfa.service'; import { MfaController } from './mfa/mfa.controller'; import { SecurityModule } from '../security/security.module'; +import { TenancyModule } from '../tenancy/tenancy.module'; import { createJwtOptions } from './config/jwt-config.factory'; /** @@ -38,6 +39,7 @@ import { createJwtOptions } from './config/jwt-config.factory'; }), TypeOrmModule.forFeature([User]), SecurityModule, + TenancyModule, ], controllers: [AuthController, SocialAuthController, MfaController], providers: [ diff --git a/src/cdn/cdn.controller.ts b/src/cdn/cdn.controller.ts index e0180f34..fdbe8902 100644 --- a/src/cdn/cdn.controller.ts +++ b/src/cdn/cdn.controller.ts @@ -6,16 +6,20 @@ import { Body, HttpException, HttpStatus, + UseGuards, } from '@nestjs/common'; import { FileInterceptor } from '@nestjs/platform-express'; import { CdnService } from './cdn.service'; import { UploadContentDto } from './dto/upload-content.dto'; +import { TenantLimitGuard, LimitType } from '../tenancy/guards/tenant-limit.guard'; @Controller('cdn') export class CdnController { constructor(private readonly cdnService: CdnService) {} @Post('upload') + @LimitType('storage') + @UseGuards(TenantLimitGuard) @UseInterceptors( FileInterceptor('file', { limits: { diff --git a/src/cdn/cdn.module.ts b/src/cdn/cdn.module.ts index 0dd36605..950a8508 100644 --- a/src/cdn/cdn.module.ts +++ b/src/cdn/cdn.module.ts @@ -2,8 +2,10 @@ import { Module } from '@nestjs/common'; import { CdnService } from './cdn.service'; import { CdnEventListener } from './cdn-event.listener'; import { CdnController } from './cdn.controller'; +import { TenancyModule } from '../tenancy/tenancy.module'; @Module({ + imports: [TenancyModule], controllers: [CdnController], providers: [CdnService, CdnEventListener], exports: [CdnService], diff --git a/src/tenancy/guards/tenant-limit.guard.spec.ts b/src/tenancy/guards/tenant-limit.guard.spec.ts new file mode 100644 index 00000000..4382adc1 --- /dev/null +++ b/src/tenancy/guards/tenant-limit.guard.spec.ts @@ -0,0 +1,169 @@ +import { HttpException, HttpStatus } from '@nestjs/common'; +import { TenantLimitGuard, LIMIT_TYPE_KEY } from './tenant-limit.guard'; +import { Tenant } from '../entities/tenant.entity'; + +describe('TenantLimitGuard', () => { + let guard: TenantLimitGuard; + let mockTenancyService: any; + let mockTenant: Partial; + + const createMockContext = (overrides: any = {}) => { + const handler = jest.fn(); + const request: any = { + headers: { 'x-tenant-id': 'tenant-1' }, + hostname: 'example.com', + ...overrides.request, + }; + + if (overrides.limitType) { + Reflect.defineMetadata(LIMIT_TYPE_KEY, overrides.limitType, handler); + } + + return { + getHandler: jest.fn(() => handler), + switchToHttp: jest.fn(() => ({ + getRequest: jest.fn(() => request), + })), + getType: jest.fn(() => 'http'), + } as any; + }; + + beforeEach(() => { + mockTenant = { + id: 'tenant-1', + userLimit: 10, + storageLimit: 1024, + currentUserCount: 5, + currentStorageUsage: 512, + }; + + mockTenancyService = { + resolveTenantIdFromRequest: jest.fn().mockResolvedValue('tenant-1'), + findOne: jest.fn().mockResolvedValue(mockTenant), + }; + + guard = new TenantLimitGuard(mockTenancyService); + }); + + describe('no limit type set', () => { + it('passes when no limit_type metadata is present', async () => { + await expect(guard.canActivate(createMockContext())).resolves.toBe(true); + }); + }); + + describe('user limit enforcement', () => { + it('allows user creation when under the limit', async () => { + mockTenant.currentUserCount = 5; + mockTenant.userLimit = 10; + + const ctx = createMockContext({ limitType: 'user' }); + await expect(guard.canActivate(ctx)).resolves.toBe(true); + }); + + it('allows user creation when exactly one below the limit (boundary)', async () => { + mockTenant.currentUserCount = 9; + mockTenant.userLimit = 10; + + const ctx = createMockContext({ limitType: 'user' }); + await expect(guard.canActivate(ctx)).resolves.toBe(true); + }); + + it('returns 402 when current user count equals the limit', async () => { + mockTenant.currentUserCount = 10; + mockTenant.userLimit = 10; + + const ctx = createMockContext({ limitType: 'user' }); + await expect(guard.canActivate(ctx)).rejects.toThrow( + new HttpException( + { message: 'User limit exceeded', error: 'Payment Required', statusCode: 402 }, + HttpStatus.PAYMENT_REQUIRED, + ), + ); + }); + + it('returns 402 when current user count exceeds the limit', async () => { + mockTenant.currentUserCount = 11; + mockTenant.userLimit = 10; + + const ctx = createMockContext({ limitType: 'user' }); + await expect(guard.canActivate(ctx)).rejects.toThrow( + new HttpException( + { message: 'User limit exceeded', error: 'Payment Required', statusCode: 402 }, + HttpStatus.PAYMENT_REQUIRED, + ), + ); + }); + + it('allows unlimited user creation when userLimit is -1', async () => { + mockTenant.userLimit = -1; + mockTenant.currentUserCount = 999; + + const ctx = createMockContext({ limitType: 'user' }); + await expect(guard.canActivate(ctx)).resolves.toBe(true); + }); + }); + + describe('storage limit enforcement', () => { + it('allows upload when under the storage limit', async () => { + mockTenant.currentStorageUsage = 500; + mockTenant.storageLimit = 1024; + + const ctx = createMockContext({ + limitType: 'storage', + request: { + file: { size: 100 * 1024 * 1024 }, + }, + }); + await expect(guard.canActivate(ctx)).resolves.toBe(true); + }); + + it('allows upload when exactly at the storage limit', async () => { + mockTenant.currentStorageUsage = 924; + mockTenant.storageLimit = 1024; + + const ctx = createMockContext({ + limitType: 'storage', + request: { + file: { size: 100 * 1024 * 1024 }, + }, + }); + await expect(guard.canActivate(ctx)).resolves.toBe(true); + }); + + it('returns 402 when upload would exceed the storage limit', async () => { + mockTenant.currentStorageUsage = 925; + mockTenant.storageLimit = 1024; + + const ctx = createMockContext({ + limitType: 'storage', + request: { + file: { size: 100 * 1024 * 1024 }, + }, + }); + await expect(guard.canActivate(ctx)).rejects.toThrow( + new HttpException( + { message: 'Storage limit exceeded', error: 'Payment Required', statusCode: 402 }, + HttpStatus.PAYMENT_REQUIRED, + ), + ); + }); + + it('passes when no file is present (let the handler validate)', async () => { + const ctx = createMockContext({ limitType: 'storage', request: {} }); + await expect(guard.canActivate(ctx)).resolves.toBe(true); + }); + + it('allows unlimited storage when storageLimit is -1', async () => { + mockTenant.storageLimit = -1; + mockTenant.currentStorageUsage = 999999; + + const ctx = createMockContext({ + limitType: 'storage', + request: { + file: { size: 500 * 1024 * 1024 }, + }, + }); + await expect(guard.canActivate(ctx)).resolves.toBe(true); + }); + }); +}); diff --git a/src/tenancy/guards/tenant-limit.guard.ts b/src/tenancy/guards/tenant-limit.guard.ts new file mode 100644 index 00000000..cc4d3ae4 --- /dev/null +++ b/src/tenancy/guards/tenant-limit.guard.ts @@ -0,0 +1,84 @@ +import { + Injectable, + CanActivate, + ExecutionContext, + HttpException, + HttpStatus, +} from '@nestjs/common'; +import { TenancyService } from '../tenancy.service'; + +export const LIMIT_TYPE_KEY = 'limit_type'; + +export function LimitType(type: 'user' | 'storage') { + return function ( + _target: unknown, + _propertyKey: string, + descriptor: PropertyDescriptor, + ) { + Reflect.defineMetadata(LIMIT_TYPE_KEY, type, descriptor.value); + }; +} + +@Injectable() +export class TenantLimitGuard implements CanActivate { + constructor(private readonly tenancyService: TenancyService) {} + + async canActivate(context: ExecutionContext): Promise { + const limitType = Reflect.getMetadata( + LIMIT_TYPE_KEY, + context.getHandler(), + ) as string | undefined; + + if (!limitType) { + return true; + } + + const request = context.switchToHttp().getRequest(); + const tenantId = await this.tenancyService.resolveTenantIdFromRequest(request); + request.tenantId = tenantId; + const tenant = await this.tenancyService.findOne(tenantId); + + if (limitType === 'user') { + if (tenant.userLimit === -1) { + return true; + } + + if (tenant.currentUserCount >= tenant.userLimit) { + throw new HttpException( + { + message: 'User limit exceeded', + error: 'Payment Required', + statusCode: HttpStatus.PAYMENT_REQUIRED, + }, + HttpStatus.PAYMENT_REQUIRED, + ); + } + } + + if (limitType === 'storage') { + if (tenant.storageLimit === -1) { + return true; + } + + const file = request.file as Express.Multer.File | undefined; + if (!file) { + return true; + } + + const uploadMB = Math.ceil(file.size / (1024 * 1024)); + + if (tenant.currentStorageUsage + uploadMB > tenant.storageLimit) { + throw new HttpException( + { + message: 'Storage limit exceeded', + error: 'Payment Required', + statusCode: HttpStatus.PAYMENT_REQUIRED, + }, + HttpStatus.PAYMENT_REQUIRED, + ); + } + } + + return true; + } +} diff --git a/src/tenancy/tenancy.module.ts b/src/tenancy/tenancy.module.ts index d69b880b..a2d0e2d5 100644 --- a/src/tenancy/tenancy.module.ts +++ b/src/tenancy/tenancy.module.ts @@ -11,6 +11,7 @@ import { TenantBillingService } from './billing/tenant-billing.service'; import { CustomizationService } from './customization/customization.service'; import { TenantAdminService } from './admin/tenant-admin.service'; import { TenantGuard } from './guards/tenant.guard'; +import { TenantLimitGuard } from './guards/tenant-limit.guard'; import { TenantMiddleware } from '../middleware/tenant/tenant.middleware'; import { TenantRlsSubscriber } from '../middleware/tenant/tenant-rls.subscriber'; import { TenantAccessValidationGuard } from '../middleware/tenant/tenant-access-validation.guard'; @@ -28,6 +29,7 @@ import { TenantAccessValidationGuard } from '../middleware/tenant/tenant-access- CustomizationService, TenantAdminService, TenantGuard, + TenantLimitGuard, TenantMiddleware, TenantRlsSubscriber, TenantAccessValidationGuard, @@ -39,6 +41,7 @@ import { TenantAccessValidationGuard } from '../middleware/tenant/tenant-access- CustomizationService, TenantAdminService, TenantGuard, + TenantLimitGuard, TenantMiddleware, TenantRlsSubscriber, TenantAccessValidationGuard,