@trifrost/core/esm/modules/RateLimit/_RateLimit.js

Blazingly fast, runtime-agnostic server framework for modern edge and node environments

import { isIntGt } from '@valkyriestudios/utils/number';
import { Sym_TriFrostDescription, Sym_TriFrostFingerPrint, Sym_TriFrostName } from '../../types/constants';
import { Sliding } from './strategies/Sliding';
import { Fixed } from './strategies/Fixed';
/* Specific symbol attached to limit mware to identify them by */
export const Sym_TriFrostMiddlewareRateLimit = Symbol('TriFrost.Middleware.RateLimit');
/**
 * Prebuilt Key Gen Registry
 */
export const RateLimitKeyGeneratorRegistry = {
    ip: ctx => ctx.ip || 'unknown',
    ip_name: ctx => (ctx.ip || 'unknown') + ':' + ctx.name,
    ip_method: ctx => (ctx.ip || 'unknown') + ':' + ctx.method,
    ip_name_method: ctx => (ctx.ip || 'unknown') + ':' + ctx.name + ':' + ctx.method,
};
export class TriFrostRateLimit {
    #keygen;
    #exceeded;
    #store;
    #headers;
    #strategy;
    #window;
    constructor(opts) {
        if (typeof opts?.store?.get !== 'function' || typeof opts?.store?.set !== 'function')
            throw new Error('TriFrostRateLimit: Expected a store initializer');
        /* Define keygen or fallback to ip_name_method */
        this.#keygen = (typeof opts?.keygen === 'function'
            ? opts.keygen
            : typeof opts?.keygen === 'string'
                ? RateLimitKeyGeneratorRegistry[opts.keygen] // eslint-disable-line prettier/prettier
                : RateLimitKeyGeneratorRegistry.ip_name_method // eslint-disable-line prettier/prettier
        ); // eslint-disable-line prettier/prettier
        /* Define exceeded behavior */
        this.#exceeded = typeof opts?.exceeded === 'function' ? opts.exceeded : (ctx) => ctx.status(429);
        /* Whether or not rate limit headers should be set (Defaults to true) */
        this.#headers = opts?.headers !== false;
        /* Set strategy */
        this.#strategy = opts?.strategy === 'sliding' ? 'sliding' : 'fixed';
        /* Set window */
        this.#window = isIntGt(opts?.window, 0) ? opts?.window : 60;
        /* Create lazy store */
        switch (this.#strategy) {
            case 'sliding':
                this.#store = new Sliding(this.#window, opts.store);
                break;
            default:
                this.#store = new Fixed(this.#window, opts.store);
                break;
        }
    }
    /**
     * Configured keygen handler
     */
    get keygen() {
        return this.#keygen;
    }
    /**
     * Configured exceeded behavior
     */
    get exceeded() {
        return this.#exceeded;
    }
    /**
     * Configured store
     */
    get store() {
        return this.#store;
    }
    /**
     * Configured headers (default=true)
     */
    get headers() {
        return this.#headers;
    }
    /**
     * The configured strategy type (default=fixed)
     */
    get strategy() {
        return this.#strategy;
    }
    /**
     * The configured window (default=60) in seconds
     */
    get window() {
        return this.#window;
    }
    /**
     * This function is meant specifically to call a 'stop' function on implementing stores.
     */
    async stop() {
        await this.#store.stop();
    }
}
/**
 * Creates a reusable "rate limit" middleware with a given limit.
 *
 * @param {Lazy<TriFrostRateLimit>} rateLimiter - Lazy rate limit instance resolver
 * @param {number|TriFrostRateLimitLimitFunction} limit - The limit to use, either a number or a rate limit function
 */
export function limitMiddleware(limiter, limit) {
    const limit_fn = (typeof limit === 'function' ? limit : () => limit);
    const mware = async function TriFrostRateLimitedMiddleware(ctx) {
        if (ctx.kind !== 'std')
            return;
        const instance = limiter.resolve(ctx);
        /* Get limit for context */
        const n_limit = limit_fn(ctx);
        if (!isIntGt(n_limit, 0))
            return ctx.status(500);
        /* Consume */
        const key = instance.keygen(ctx);
        const usage = await instance.store.consume(typeof key === 'string' && key.length ? key : 'unknown', n_limit);
        if (usage.amt > n_limit) {
            if (instance.headers) {
                ctx.setHeaders({
                    'retry-after': usage.reset - Math.floor(Date.now() / 1000),
                    'x-ratelimit-limit': n_limit,
                    'x-ratelimit-remaining': '0',
                    'x-ratelimit-reset': usage.reset,
                });
            }
            else {
                ctx.setHeader('retry-after', usage.reset - Math.floor(Date.now() / 1000));
            }
            return instance.exceeded(ctx);
        }
        if (instance.headers) {
            ctx.setHeaders({
                'x-ratelimit-limit': n_limit,
                'x-ratelimit-remaining': n_limit - usage.amt,
                'x-ratelimit-reset': usage.reset,
            });
        }
    };
    /* Add symbols for introspection/use further down the line */
    Reflect.set(mware, Sym_TriFrostName, 'TriFrostRateLimit');
    Reflect.set(mware, Sym_TriFrostDescription, 'Middleware for rate limitting contexts passing through it');
    Reflect.set(mware, Sym_TriFrostFingerPrint, Sym_TriFrostMiddlewareRateLimit);
    return mware;
}