SponsorBlockServer/src/middleware/requestRateLimit.ts

import {getIP} from "../utils/getIP";
import {getHash} from "../utils/getHash";
import rateLimit from "express-rate-limit";
import {RateLimitConfig} from "../types/config.model";
import {Request} from "express";
import { isUserVIP } from "../utils/isUserVIP";
import { UserID } from "../types/user.model";

export function rateLimitMiddleware(limitConfig: RateLimitConfig, getUserID?: (req: Request) => UserID): rateLimit.RateLimit {
    return rateLimit({
        windowMs: limitConfig.windowMs,
        max: limitConfig.max,
        message: limitConfig.message,
        statusCode: limitConfig.statusCode,
        headers: false,
        keyGenerator: (req) => {
            return getHash(getIP(req), 1);
        },
        handler: async (req, res, next) => {
            if (getUserID === undefined || !await isUserVIP(getHash(getUserID(req)))) {
                return res.status(limitConfig.statusCode).send(limitConfig.message);
            } else {
                return next();
            }
        }
    });
}