FrenoCorp/apps/api/src/middleware/spam-rate-limit.middleware.ts

import { FastifyInstance, FastifyRequest, FastifyReply } from 'fastify';
import { redis } from '../config/redis';
import { spamRateLimits } from '../services/spamshield/spamshield.config';

const REDIS_PREFIX = 'spamshield:ratelimit';

class RedisRateLimiter {
  async checkLimit(
    key: string,
    windowSeconds: number,
    maxRequests: number
  ): Promise<{
    remaining: number;
    resetTime: number;
    retryAfter?: number;
  }> {
    const redisKey = `${REDIS_PREFIX}:${key}`;
    const now = Date.now();

    const current = await redis.get(redisKey);
    const windowStart = now - (now % (windowSeconds * 1000));
    const resetTime = windowStart + windowSeconds * 1000;

    if (!current) {
      const expirySeconds = Math.ceil((resetTime - now) / 1000);
      await redis.set(redisKey, '1', 'EX', expirySeconds);

      return {
        remaining: maxRequests - 1,
        resetTime,
      };
    }

    const count = parseInt(current, 10) + 1;
    await redis.set(redisKey, String(count), 'EX', Math.ceil((resetTime - now) / 1000));

    const remaining = maxRequests - count;

    if (count > maxRequests) {
      return {
        remaining: 0,
        resetTime,
        retryAfter: resetTime - now,
      };
    }

    return {
      remaining,
      resetTime,
    };
  }

  async checkDailyLimit(
    key: string,
    maxPerDay: number
  ): Promise<{
    remaining: number;
    retryAfter?: number;
  }> {
    const redisKey = `${REDIS_PREFIX}:daily:${key}`;
    const now = Date.now();
    const dayStart = new Date(now);
    dayStart.setHours(0, 0, 0, 0);
    const dayEnd = new Date(dayStart);
    dayEnd.setDate(dayEnd.getDate() + 1);
    const resetTime = dayEnd.getTime();

    const current = await redis.get(redisKey);
    const expirySeconds = Math.ceil((resetTime - now) / 1000);

    if (!current) {
      await redis.set(redisKey, '1', 'EX', expirySeconds);

      return {
        remaining: maxPerDay - 1,
      };
    }

    const count = parseInt(current, 10) + 1;
    await redis.set(redisKey, String(count), 'EX', expirySeconds);

    const remaining = maxPerDay - count;

    if (count > maxPerDay) {
      return {
        remaining: 0,
        retryAfter: resetTime - now,
      };
    }

    return {
      remaining,
    };
  }

  reset(key: string) {
    const redisKey = `${REDIS_PREFIX}:${key}`;
    return redis.del(redisKey);
  }
}

export const spamRateLimiter = new RedisRateLimiter();

export async function spamRateLimitMiddleware(fastify: FastifyInstance) {
  fastify.addHook('preHandler', async (request: FastifyRequest, reply: FastifyReply) => {
    const url = request.url || '';

    if (!url.startsWith('/spamshield')) {
      return;
    }

    const clientIp = request.ip || (request.headers['x-forwarded-for'] as string) || 'unknown';
    const apiKey = request.headers['x-api-key'] as string | undefined;
    const key = apiKey ? `api:${apiKey}` : `ip:${clientIp}`;

    let tier = 'basic';
    if (apiKey) {
      if (apiKey.startsWith('premium_')) {
        tier = 'premium';
      } else if (apiKey.startsWith('plus_')) {
        tier = 'plus';
      }
    }

    const config = spamRateLimits[tier as keyof typeof spamRateLimits];

    const minuteResult = await spamRateLimiter.checkLimit(
      key,
      60,
      config.analysesPerMinute
    );

    const dailyResult = await spamRateLimiter.checkDailyLimit(
      key,
      config.analysesPerDay
    );

    reply.header('X-RateLimit-Limit', config.analysesPerMinute);
    reply.header('X-RateLimit-Remaining', minuteResult.remaining);
    reply.header('X-RateLimit-Reset', Math.ceil(minuteResult.resetTime / 1000));
    reply.header('X-RateLimit-Daily-Limit', config.analysesPerDay);
    reply.header('X-RateLimit-Daily-Remaining', dailyResult.remaining);

    const retryAfter = minuteResult.retryAfter || dailyResult.retryAfter;

    if (retryAfter) {
      reply.header('Retry-After', Math.ceil(retryAfter / 1000));
      reply.code(429);

      return {
        error: 'Too Many Requests',
        message: `Spam analysis rate limit exceeded. Try again in ${Math.ceil(retryAfter / 1000)}s`,
        tier,
        limit: config.analysesPerMinute,
        dailyLimit: config.analysesPerDay,
        reset: new Date(minuteResult.resetTime).toISOString(),
      };
    }

    (request as any).spamRateLimitTier = tier;
  });
}

export { RedisRateLimiter };