Assist_Design/apps/bff/src/modules/auth/infra/rate-limiting/auth-rate-limit.service.ts

import { Inject, Injectable, InternalServerErrorException } from "@nestjs/common";
import { ConfigService } from "@nestjs/config";
import { ThrottlerException } from "@nestjs/throttler";
import { Logger } from "nestjs-pino";
import type { Request } from "express";
import { RateLimiterRedis, RateLimiterRes } from "rate-limiter-flexible";
import { createHash } from "crypto";
import type { Redis } from "ioredis";
import { getErrorMessage } from "@bff/core/utils/error.util";

export interface RateLimitOutcome {
  key: string;
  remainingPoints: number;
  consumedPoints: number;
  msBeforeNext: number;
  needsCaptcha: boolean;
}

@Injectable()
export class AuthRateLimitService {
  private readonly loginLimiter: RateLimiterRedis;

  private readonly refreshLimiter: RateLimiterRedis;

  private readonly signupLimiter: RateLimiterRedis;

  private readonly passwordResetLimiter: RateLimiterRedis;

  private readonly loginCaptchaThreshold: number;

  private readonly captchaAlwaysOn: boolean;

  constructor(
    @Inject("REDIS_CLIENT") private readonly redis: Redis,
    private readonly configService: ConfigService,
    @Inject(Logger) private readonly logger: Logger
  ) {
    const loginLimit = this.configService.get<number>("LOGIN_RATE_LIMIT_LIMIT", 5);
    const loginTtlMs = this.configService.get<number>("LOGIN_RATE_LIMIT_TTL", 900000);

    const signupLimit = this.configService.get<number>("SIGNUP_RATE_LIMIT_LIMIT", 5);
    const signupTtlMs = this.configService.get<number>("SIGNUP_RATE_LIMIT_TTL", 900000);

    const passwordResetLimit = this.configService.get<number>("PASSWORD_RESET_RATE_LIMIT_LIMIT", 5);
    const passwordResetTtlMs = this.configService.get<number>(
      "PASSWORD_RESET_RATE_LIMIT_TTL",
      900000
    );

    const refreshLimit = this.configService.get<number>("AUTH_REFRESH_RATE_LIMIT_LIMIT", 10);
    const refreshTtlMs = this.configService.get<number>("AUTH_REFRESH_RATE_LIMIT_TTL", 300000);

    this.loginCaptchaThreshold = this.configService.get<number>("LOGIN_CAPTCHA_AFTER_ATTEMPTS", 3);
    this.captchaAlwaysOn = this.configService.get("AUTH_CAPTCHA_ALWAYS_ON", "false") === "true";

    this.loginLimiter = this.createLimiter("auth-login", loginLimit, loginTtlMs);
    this.signupLimiter = this.createLimiter("auth-signup", signupLimit, signupTtlMs);
    this.passwordResetLimiter = this.createLimiter(
      "auth-password-reset",
      passwordResetLimit,
      passwordResetTtlMs
    );
    this.refreshLimiter = this.createLimiter("auth-refresh", refreshLimit, refreshTtlMs);
  }

  async consumeLoginAttempt(request: Request): Promise<RateLimitOutcome> {
    const key = this.buildKey("login", request);
    const result = await this.consume(this.loginLimiter, key, "login");
    return {
      ...result,
      needsCaptcha: this.captchaAlwaysOn || this.requiresCaptcha(result.consumedPoints),
    };
  }

  async clearLoginAttempts(request: Request): Promise<void> {
    const key = this.buildKey("login", request);
    await this.deleteKey(this.loginLimiter, key, "login");
  }

  async consumeSignupAttempt(request: Request): Promise<RateLimitOutcome> {
    const key = this.buildKey("signup", request);
    const outcome = await this.consume(this.signupLimiter, key, "signup");
    return { ...outcome, needsCaptcha: false };
  }

  async consumePasswordReset(request: Request): Promise<RateLimitOutcome> {
    const key = this.buildKey("password-reset", request);
    const outcome = await this.consume(this.passwordResetLimiter, key, "password-reset");
    return { ...outcome, needsCaptcha: false };
  }

  async consumeRefreshAttempt(request: Request, refreshToken: string): Promise<RateLimitOutcome> {
    const tokenHash = this.hashToken(refreshToken);
    const key = this.buildKey("refresh", request, tokenHash);
    return this.consume(this.refreshLimiter, key, "refresh");
  }

  getCaptchaHeaderValue(needsCaptcha: boolean): string {
    if (needsCaptcha || this.captchaAlwaysOn) {
      return "required";
    }
    return "optional";
  }

  private requiresCaptcha(consumedPoints: number): boolean {
    if (this.loginCaptchaThreshold <= 0) {
      return false;
    }
    return consumedPoints >= this.loginCaptchaThreshold;
  }

  private buildKey(type: string, request: Request, suffix?: string): string {
    const ip = this.extractIp(request);
    const userAgent = request.headers["user-agent"] || "unknown";
    const uaHash = createHash("sha256").update(String(userAgent)).digest("hex").slice(0, 16);

    return ["auth", type, ip, uaHash, suffix].filter(Boolean).join(":");
  }

  private extractIp(request: Request): string {
    const forwarded = request.headers["x-forwarded-for"];
    const forwardedIp = Array.isArray(forwarded) ? forwarded[0] : forwarded;
    const rawIp =
      (typeof forwardedIp === "string" ? forwardedIp.split(",")[0]?.trim() : undefined) ||
      (request.headers["x-real-ip"] as string | undefined) ||
      request.socket?.remoteAddress ||
      request.ip ||
      "unknown";

    return rawIp.replace(/^::ffff:/, "");
  }

  private createLimiter(prefix: string, limit: number, ttlMs: number): RateLimiterRedis {
    const duration = Math.max(1, Math.ceil(ttlMs / 1000));

    try {
      return new RateLimiterRedis({
        storeClient: this.redis,
        keyPrefix: prefix,
        points: limit,
        duration,
        inMemoryBlockOnConsumed: limit + 1,
        insuranceLimiter: undefined,
      });
    } catch (error: unknown) {
      this.logger.error(
        { prefix, error: getErrorMessage(error) },
        "Failed to initialize rate limiter"
      );
      throw new InternalServerErrorException("Rate limiter initialization failed");
    }
  }

  private async consume(
    limiter: RateLimiterRedis,
    key: string,
    context: string
  ): Promise<RateLimitOutcome> {
    try {
      const res = await limiter.consume(key);
      const consumedPoints = Math.max(0, limiter.points - res.remainingPoints);

      return {
        key,
        remainingPoints: res.remainingPoints,
        consumedPoints,
        msBeforeNext: res.msBeforeNext,
        needsCaptcha: false,
      };
    } catch (error: unknown) {
      if (error instanceof RateLimiterRes) {
        const retryAfterMs = error?.msBeforeNext ?? 0;
        const message = this.buildThrottleMessage(context, retryAfterMs);

        this.logger.warn({ key, context, retryAfterMs }, "Auth rate limit reached");

        throw new ThrottlerException(message);
      }

      this.logger.error({ key, context, error: getErrorMessage(error) }, "Rate limiter failure");
      throw new ThrottlerException("Authentication temporarily unavailable");
    }
  }

  private async deleteKey(limiter: RateLimiterRedis, key: string, context: string): Promise<void> {
    try {
      await limiter.delete(key);
    } catch (error: unknown) {
      this.logger.warn(
        { key, context, error: getErrorMessage(error) },
        "Failed to reset rate limiter key"
      );
    }
  }

  private buildThrottleMessage(context: string, retryAfterMs: number): string {
    const seconds = Math.ceil(retryAfterMs / 1000);
    if (seconds <= 60) {
      return `Too many ${context} attempts. Try again in ${seconds}s.`;
    }
    const minutes = Math.floor(seconds / 60);
    const remainingSeconds = seconds % 60;
    const timeMessage = remainingSeconds ? `${minutes}m ${remainingSeconds}s` : `${minutes}m`;
    return `Too many ${context} attempts. Try again in ${timeMessage}.`;
  }

  private hashToken(token: string): string {
    return createHash("sha256").update(token).digest("hex");
  }
}