feat: add TenantLimitGuard to enforce user and storage limits on tena…

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<User>,
     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' })

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: [

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: {

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],

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<Tenant>;
+
+  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);
+    });
+  });
+});

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<boolean> {
+    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;
+  }
+}