feat: initially implement locations as id

Author: alexey-yarmosh

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

@@ -12,21 +12,21 @@ const setResponseHeaders = (ctx: Context, response: RateLimiterRes) => {
 };
 
 export const rateLimitHandler = () => async (ctx: Context, next: Next) => {
-	const { isAdmin } = ctx;
+	// const { isAdmin } = ctx;
 	const clientIp = requestIp.getClientIp(ctx.req) ?? '';
-	const request = ctx.request.body as MeasurementRequest;
-	const limit = request.locations.some(l => l.limit) ? request.locations.reduce((sum, { limit = 1 }) => sum + limit, 0) : request.limit;
+	// const request = ctx.request.body as MeasurementRequest;
+	// const limit = request.locations.some(l => l.limit) ? request.locations.reduce((sum, { limit = 1 }) => sum + limit, 0) : request.limit;
 
-	if (isAdmin) {
-		return next();
-	}
+	// if (isAdmin) {
+	// 	return next();
+	// }
 
-	const currentState = await rateLimiter.get(clientIp) ?? defaultState as RateLimiterRes;
+	// const currentState = await rateLimiter.get(clientIp) ?? defaultState as RateLimiterRes;
 
-	if (currentState.remainingPoints < limit) {
-		setResponseHeaders(ctx, currentState);
-		throw createHttpError(429, 'API rate limit exceeded.', { type: 'rate_limit_exceeded' });
-	}
+	// if (currentState.remainingPoints < limit) {
+	// 	setResponseHeaders(ctx, currentState);
+	// 	throw createHttpError(429, 'API rate limit exceeded.', { type: 'rate_limit_exceeded' });
+	// }
 
 	await next();
 	const response = ctx.response.body as object & { probesCount?: number };

Diff for: src/measurement/schema/global-schema.ts

@@ -13,7 +13,7 @@ export const schema = Joi.object({
 	type: Joi.string().valid('ping', 'traceroute', 'dns', 'mtr', 'http').insensitive().required(),
 	target: targetSchema,
 	measurementOptions: measurementSchema,
-	locations: locationSchema,
+	locations: Joi.alternatives().try(locationSchema, Joi.string()),
 	limit: Joi.number().min(1).max(measurementConfig.limits.global).default(GLOBAL_DEFAULTS.limit),
 	inProgressUpdates: Joi.bool().default(GLOBAL_DEFAULTS.inProgressUpdates),
 });

Diff for: src/measurement/store.ts

@@ -21,7 +21,7 @@ const substractObjects = (obj1: Record<string, unknown>, obj2: Record<string, un
 		const value1 = obj1[key];
 		const value2 = obj2[key];
 
-		if (_.isPlainObject(value1)) {
+		if (_.isPlainObject(value1) && _.isPlainObject(value2)) {
 			const difference = substractObjects(value1 as Record<string, unknown>, value2 as Record<string, unknown>);
 
 			if (!_.isEmpty(difference)) {
@@ -93,6 +93,12 @@ export class MeasurementStore {
 		]);
 	}
 
+	async getIpsByMeasurementId (id: string): Promise<string[]> {
+		const key = getMeasurementKey(id, 'ips');
+		const ips = await this.redis.json.get(key) as string[] | null;
+		return ips || [];
+	}
+
 	async markFinishedByTimeout (ids: string[]): Promise<void> {
 		if (ids.length === 0) {
 			return;
@@ -151,8 +157,11 @@ export class MeasurementStore {
 
 	removeDefaults (measurement: MeasurementRecord, request: MeasurementRequest): Partial<MeasurementRecord> {
 		const defaults = getDefaults(request);
+
 		// Removes `"limit": 1` from locations. E.g. [{"country": "US", "limit": 1}] => [{"country": "US"}]
-		measurement.locations = measurement.locations.map(location => location.limit === 1 ? _.omit(location, 'limit') : location);
+		if (_.isArray(measurement.locations)) {
+			measurement.locations = measurement.locations.map(location => location.limit === 1 ? _.omit(location, 'limit') : location);
+		}
 
 		return substractObjects(measurement, defaults) as Partial<MeasurementRecord>;
 	}

Diff for: src/measurement/types.ts

@@ -176,7 +176,7 @@ export type MeasurementRequest = {
 	type: 'ping' | 'traceroute' | 'dns' | 'http' | 'mtr';
 	target: string;
 	measurementOptions: MeasurementOptions;
-	locations: LocationWithLimit[];
+	locations: LocationWithLimit[] | string;
 	limit: number;
 	inProgressUpdates: boolean;
 };
@@ -207,7 +207,7 @@ export type MeasurementRecord = {
 	target: string;
 	limit: number;
 	probesCount: number;
-	locations: LocationWithLimit[];
+	locations: LocationWithLimit[] | string;
 	measurementOptions?: MeasurementOptions;
 	results: MeasurementResult[];
 };

Diff for: src/probe/router.ts

@@ -5,21 +5,25 @@ import type { LocationWithLimit } from '../measurement/types.js';
 import type { Location } from '../lib/location/types.js';
 import type { Probe } from './types.js';
 import { SocketsLocationFilter } from './sockets-location-filter.js';
+import { getMeasurementStore } from '../measurement/store.js';
 
 export class ProbeRouter {
-	constructor (
-		private readonly fetchWsSockets: typeof fetchSockets,
-		private readonly socketsFilter = new SocketsLocationFilter(),
-	) {}
+	private readonly socketsFilter = new SocketsLocationFilter();
+
+	private readonly store = getMeasurementStore();
+
+	constructor (private readonly fetchWsSockets: typeof fetchSockets) {}
 
 	public async findMatchingProbes (
-		locations: LocationWithLimit[] = [],
+		locations: LocationWithLimit[] | string = [],
 		globalLimit = 1,
 	): Promise<Probe[]> {
 		const sockets = await this.fetchSockets();
 		let filtered: RemoteProbeSocket[] = [];
 
-		if (locations.some(l => l.limit)) {
+		if (typeof locations === 'string') {
+			filtered = await this.findWithMeasurementId(sockets, locations);
+		} else if (locations.some(l => l.limit)) {
 			filtered = this.findWithLocationLimit(sockets, locations);
 		} else if (locations.length > 0) {
 			filtered = this.findWithGlobalLimit(sockets, locations, globalLimit);
@@ -68,6 +72,11 @@ export class ProbeRouter {
 
 		return [ ...picked ];
 	}
+
+	private async findWithMeasurementId (sockets: RemoteProbeSocket[], measurementId: string): Promise<RemoteProbeSocket[]> {
+		const ips = await this.store.getIpsByMeasurementId(measurementId);
+		return sockets.filter(socket => ips.includes(socket.data.probe.ipAddress));
+	}
 }
 
 // Factory