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

@@ -5,22 +5,24 @@ import { PassportModule } from '@nestjs/passport';
 import { TypeOrmModule } from '@nestjs/typeorm';
 import { AuditLogModule } from '../audit-log/audit-log.module';
 import { RbacModule } from '../rbac/rbac.module';
+import { SecurityModule } from '../security/security.module';
+import { TenancyModule } from '../tenancy/tenancy.module';
 import { User } from '../users/entities/user.entity';
+
 import { createJwtOptions } from './config/jwt-config.factory';
 import { AuthController } from './auth.controller';
 import { AuthService } from './auth.service';
 import { JwtStrategy } from './jwt.strategy';
 import { PermissionsGuard } from './guards/permissions.guard';
 import { RolesGuard } from './guards/roles.guard';
-import { AuthTokensService } from './services/auth-tokens.service';
+import { MfaController } from './mfa/mfa.controller';
+import { MfaService } from './mfa/mfa.service';
 import { SocialAuthController } from './controllers/social-auth.controller';
+import { AuthTokensService } from './services/auth-tokens.service';
 import { SocialAuthService } from './services/social-auth.service';
 import { TokenBlacklistService } from './services/token-blacklist.service';
-import { GoogleStrategy } from './strategies/google.strategy';
 import { GitHubStrategy } from './strategies/github.strategy';
-import { MfaController } from './mfa/mfa.controller';
-import { MfaService } from './mfa/mfa.service';
-import { SecurityModule } from '../security/security.module';
+import { GoogleStrategy } from './strategies/google.strategy';
 
 /**
  * Registers the authentication module with Passport and JWT support.
@@ -36,10 +38,12 @@ import { SecurityModule } from '../security/security.module';
     JwtModule.registerAsync({
       imports: [ConfigModule],
       inject: [ConfigService],
-      useFactory: (configService: ConfigService) => createJwtOptions(configService),
+      useFactory: (configService: ConfigService) =>
+        createJwtOptions(configService),
     }),
     TypeOrmModule.forFeature([User]),
     SecurityModule,
+    TenancyModule,
     RbacModule,
     AuditLogModule,
   ],
@@ -66,4 +70,4 @@ import { SecurityModule } from '../security/security.module';
     PermissionsGuard,
   ],
 })
-export class AuthModule {}
+export class AuthModule {}

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