@trifrost/core
Version:
Blazingly fast, runtime-agnostic server framework for modern edge and node environments
146 lines (145 loc) • 5.49 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.TriFrostRateLimit = exports.RateLimitKeyGeneratorRegistry = exports.Sym_TriFrostMiddlewareRateLimit = void 0;
exports.limitMiddleware = limitMiddleware;
const number_1 = require("@valkyriestudios/utils/number");
const constants_1 = require("../../types/constants");
const Sliding_1 = require("./strategies/Sliding");
const Fixed_1 = require("./strategies/Fixed");
/* Specific symbol attached to limit mware to identify them by */
exports.Sym_TriFrostMiddlewareRateLimit = Symbol('TriFrost.Middleware.RateLimit');
/**
* Prebuilt Key Gen Registry
*/
exports.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,
};
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'
? exports.RateLimitKeyGeneratorRegistry[opts.keygen] // eslint-disable-line prettier/prettier
: exports.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 = (0, number_1.isIntGt)(opts?.window, 0) ? opts?.window : 60;
/* Create lazy store */
switch (this.#strategy) {
case 'sliding':
this.#store = new Sliding_1.Sliding(this.#window, opts.store);
break;
default:
this.#store = new Fixed_1.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();
}
}
exports.TriFrostRateLimit = TriFrostRateLimit;
/**
* 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
*/
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 (!(0, number_1.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, constants_1.Sym_TriFrostName, 'TriFrostRateLimit');
Reflect.set(mware, constants_1.Sym_TriFrostDescription, 'Middleware for rate limitting contexts passing through it');
Reflect.set(mware, constants_1.Sym_TriFrostFingerPrint, exports.Sym_TriFrostMiddlewareRateLimit);
return mware;
}