05a4f050bb
* chore: Improve graceful server shutdown * Replace node timers with custom promise timeout
92 lines
2.5 KiB
TypeScript
92 lines
2.5 KiB
TypeScript
import { Context, Next } from "koa";
|
||
import { defaults } from "lodash";
|
||
import env from "@server/env";
|
||
import { RateLimitExceededError } from "@server/errors";
|
||
import Metrics from "@server/logging/Metrics";
|
||
import Redis from "@server/redis";
|
||
import RateLimiter from "@server/utils/RateLimiter";
|
||
|
||
/**
|
||
* Middleware that limits the number of requests that are allowed within a given
|
||
* window. Should only be applied once to a server – do not use on individual
|
||
* routes.
|
||
*
|
||
* @returns The middleware function.
|
||
*/
|
||
export function defaultRateLimiter() {
|
||
return async function rateLimiterMiddleware(ctx: Context, next: Next) {
|
||
if (!env.RATE_LIMITER_ENABLED) {
|
||
return next();
|
||
}
|
||
|
||
const key = RateLimiter.hasRateLimiter(ctx.path)
|
||
? `${ctx.path}:${ctx.ip}`
|
||
: `${ctx.ip}`;
|
||
const limiter = RateLimiter.getRateLimiter(ctx.path);
|
||
|
||
try {
|
||
await limiter.consume(key);
|
||
} catch (rateLimiterRes) {
|
||
ctx.set("Retry-After", `${rateLimiterRes.msBeforeNext / 1000}`);
|
||
ctx.set("RateLimit-Limit", `${limiter.points}`);
|
||
ctx.set("RateLimit-Remaining", `${rateLimiterRes.remainingPoints}`);
|
||
ctx.set(
|
||
"RateLimit-Reset",
|
||
`${new Date(Date.now() + rateLimiterRes.msBeforeNext)}`
|
||
);
|
||
|
||
Metrics.increment("rate_limit.exceeded", {
|
||
path: ctx.path,
|
||
});
|
||
|
||
throw RateLimitExceededError();
|
||
}
|
||
|
||
return next();
|
||
};
|
||
}
|
||
|
||
type RateLimiterConfig = {
|
||
/** The window for which this rate limiter is considered (defaults to 60s) */
|
||
duration?: number;
|
||
/** The number of requests per IP address that are allowed within the window */
|
||
requests: number;
|
||
};
|
||
|
||
/**
|
||
* Middleware that limits the number of requests per IP address that are allowed
|
||
* within a window, overrides default middleware when used on a route.
|
||
*
|
||
* @returns The middleware function.
|
||
*/
|
||
export function rateLimiter(config: RateLimiterConfig) {
|
||
return async function registerRateLimiterMiddleware(
|
||
ctx: Context,
|
||
next: Next
|
||
) {
|
||
if (!env.RATE_LIMITER_ENABLED) {
|
||
return next();
|
||
}
|
||
|
||
if (!RateLimiter.hasRateLimiter(ctx.path)) {
|
||
RateLimiter.setRateLimiter(
|
||
ctx.path,
|
||
defaults(
|
||
{
|
||
...config,
|
||
points: config.requests,
|
||
},
|
||
{
|
||
duration: 60,
|
||
points: env.RATE_LIMITER_REQUESTS,
|
||
keyPrefix: RateLimiter.RATE_LIMITER_REDIS_KEY_PREFIX,
|
||
storeClient: Redis.defaultClient,
|
||
}
|
||
)
|
||
);
|
||
}
|
||
|
||
return next();
|
||
};
|
||
}
|