fix: rate limits

Author: alexey-yarmosh

Diff for: src/lib/http/middleware/ratelimit.ts

@@ -1,21 +1,39 @@
 import config from 'config';
+import type { Context } from 'koa';
 import { RateLimiterRedis } from 'rate-limiter-flexible';
+import requestIp from 'request-ip';
+import type { RateLimiterRes } from 'rate-limiter-flexible';
 import { createRedisClient } from './redis/client.js';
+import createHttpError from 'http-errors';
 
 const redisClient = await createRedisClient({ legacyMode: true });
 
-export const defaultState = {
-	remainingPoints: config.get<number>('measurement.rateLimit'),
-	msBeforeNext: config.get<number>('measurement.rateLimitReset') * 1000,
-	consumedPoints: 0,
-	isFirstInDuration: true,
-};
-
 const rateLimiter = new RateLimiterRedis({
 	storeClient: redisClient,
 	keyPrefix: 'rate',
 	points: config.get<number>('measurement.rateLimit'),
 	duration: config.get<number>('measurement.rateLimitReset'),
 });
 
+const setRateLimitHeaders = (ctx: Context, result: RateLimiterRes) => {
+	ctx.set('X-RateLimit-Reset', `${Math.round(result.msBeforeNext / 1000)}`);
+	ctx.set('X-RateLimit-Limit', `${rateLimiter.points}`);
+	ctx.set('X-RateLimit-Remaining', `${result.remainingPoints}`);
+};
+
+export const checkRateLimits = async (ctx: Context, numberOfProbes: number) => {
+	if (ctx['isAdmin']) {
+		return;
+	}
+
+	try {
+		const clientIp = requestIp.getClientIp(ctx.req) ?? '';
+		const result = await rateLimiter.consume(clientIp, numberOfProbes);
+		setRateLimitHeaders(ctx, result);
+	} catch (error) {
+		setRateLimitHeaders(ctx, error as RateLimiterRes);
+		throw createHttpError(429, 'Too Many Probes Requested', { type: 'too_many_probes' });
+	}
+};
+
 export default rateLimiter;

Diff for: src/lib/ratelimiter.ts

@@ -2,18 +2,15 @@ import config from 'config';
 import type { Context } from 'koa';
 import type Router from '@koa/router';
 import { getMeasurementRunner } from '../runner.js';
-import type { MeasurementRequest } from '../types.js';
 import { bodyParser } from '../../lib/http/middleware/body-parser.js';
 import { validate } from '../../lib/http/middleware/validate.js';
 import { schema } from '../schema/global-schema.js';
-import { rateLimitHandler } from '../../lib/http/middleware/ratelimit.js';
 
 const hostConfig = config.get<string>('server.host');
 const runner = getMeasurementRunner();
 
 const handle = async (ctx: Context): Promise<void> => {
-	const request = ctx.request.body as MeasurementRequest;
-	const { measurementId, probesCount } = await runner.run(request);
+	const { measurementId, probesCount } = await runner.run(ctx);
 
 	ctx.status = 202;
 	ctx.set('Location', `${hostConfig}/v1/measurements/${measurementId}`);
@@ -25,5 +22,5 @@ const handle = async (ctx: Context): Promise<void> => {
 };
 
 export const registerCreateMeasurementRoute = (router: Router): void => {
-	router.post('/measurements', '/measurements', bodyParser(), validate(schema), rateLimitHandler(), handle);
+	router.post('/measurements', '/measurements', bodyParser(), validate(schema), handle);
 };

Diff for: src/measurement/route/create-measurement.ts

@@ -1,3 +1,4 @@
+import type { Context } from 'koa';
 import config from 'config';
 import type { Server } from 'socket.io';
 import createHttpError from 'http-errors';
@@ -9,11 +10,8 @@ import type { Probe } from '../probe/types.js';
 import { getMetricsAgent, type MetricsAgent } from '../lib/metrics.js';
 import type { MeasurementStore } from './store.js';
 import { getMeasurementStore } from './store.js';
-import type {
-	MeasurementRequest,
-	MeasurementResultMessage,
-	MeasurementProgressMessage,
-} from './types.js';
+import type { MeasurementRequest, MeasurementResultMessage, MeasurementProgressMessage } from './types.js';
+import { checkRateLimits } from '../lib/ratelimiter.js';
 
 export class MeasurementRunner {
 	constructor (
@@ -24,13 +22,16 @@ export class MeasurementRunner {
 		private readonly metrics: MetricsAgent,
 	) {}
 
-	async run (request: MeasurementRequest): Promise<{measurementId: string; probesCount: number;}> {
+	async run (ctx: Context): Promise<{measurementId: string; probesCount: number;}> {
+		const request = ctx.request.body as MeasurementRequest;
 		const probes = await this.router.findMatchingProbes(request.locations, request.limit);
 
 		if (probes.length === 0) {
 			throw createHttpError(422, 'No suitable probes found.', { type: 'no_probes_found' });
 		}
 
+		await checkRateLimits(ctx, probes.filter(probe => probe.status !== 'offline').length);
+
 		const measurementId = await this.store.createMeasurement(request, probes);
 
 		this.sendToProbes(measurementId, probes, request);