feat: periodic ip limit check (#513)

Co-authored-by: Martin Kolárik <martin@kolarik.sk>

Author: alexey-yarmosh

src/lib/server.ts

@@ -1,6 +1,6 @@
 import type { Server } from 'node:http';
 import { initRedisClient } from './redis/client.js';
-import { adoptedProbes, initWsServer } from './ws/server.js';
+import { adoptedProbes, probeIpLimit, initWsServer } from './ws/server.js';
 import { getMetricsAgent } from './metrics.js';
 import { populateMemList as populateMemMalwareList } from './malware/client.js';
 import { populateMemList as populateMemIpRangesList } from './ip-ranges.js';
@@ -33,6 +33,8 @@ export const createServer = async (): Promise<Server> => {
 	await auth.syncTokens();
 	auth.scheduleSync();
 
+	probeIpLimit.scheduleSync();
+
 	reconnectProbes();
 
 	const { getWsServer } = await import('./ws/server.js');

src/lib/ws/helper/probe-ip-limit.ts

@@ -1,19 +1,70 @@
-import { fetchProbes } from '../server.js';
+import _ from 'lodash';
+import config from 'config';
+import type { fetchProbes as serverFetchProbes, fetchRawSockets as serverFetchRawSockets } from '../server.js';
 import { scopedLogger } from '../../logger.js';
 import { ProbeError } from '../../probe-error.js';
 
+const numberOfProcesses = config.get<number>('server.processes');
+
 const logger = scopedLogger('ws:limit');
 
-export const verifyIpLimit = async (ip: string, socketId: string): Promise<void> => {
-	if (process.env['FAKE_PROBE_IP'] || process.env['TEST_MODE'] === 'unit') {
-		return;
+export class ProbeIpLimit {
+	private timer: NodeJS.Timeout | undefined;
+
+	constructor (
+		private readonly fetchProbes: typeof serverFetchProbes,
+		private readonly fetchRawSockets: typeof serverFetchRawSockets,
+	) {}
+
+	scheduleSync () {
+		clearTimeout(this.timer);
+
+		this.timer = setTimeout(() => {
+			this.syncIpLimit()
+				.finally(() => this.scheduleSync())
+				.catch(error => logger.error(error));
+		}, 60_000 * 2 * Math.random() * numberOfProcesses).unref();
+	}
+
+	async syncIpLimit () {
+		const probes = await this.fetchProbes();
+		// Sorting probes by "client" (socket id), so all workers will treat the same probe as "first".
+		const sortedProbes = _.sortBy(probes, [ 'client' ]);
+
+		const uniqIpToSocketId = new Map<string, string>();
+		const socketIdsToDisconnect = new Set<string>();
+
+		for (const probe of sortedProbes) {
+			const prevSocketId = uniqIpToSocketId.get(probe.ipAddress);
+
+			if (prevSocketId && prevSocketId !== probe.client) {
+				logger.warn(`Probe ip duplication occurred (${probe.ipAddress}). Socket id to preserve: ${prevSocketId}, socket id to disconnect: ${probe.client}`);
+				socketIdsToDisconnect.add(probe.client);
+			} else {
+				uniqIpToSocketId.set(probe.ipAddress, probe.client);
+			}
+		}
+
+		if (socketIdsToDisconnect.size > 0) {
+			const sockets = await this.fetchRawSockets();
+			sockets
+				.filter(socket => socketIdsToDisconnect.has(socket.id))
+				.forEach(socket => socket.disconnect());
+		}
 	}
 
-	const probes = await fetchProbes({ allowStale: false });
-	const previousSocket = probes.find(p => p.ipAddress === ip && p.client !== socketId);
+	async verifyIpLimit (ip: string, socketId: string): Promise<void> {
+		if (process.env['FAKE_PROBE_IP'] || process.env['TEST_MODE'] === 'unit') {
+			return;
+		}
+
+		const probes = await this.fetchProbes({ allowStale: false });
+		const previousProbe = probes.find(p => p.ipAddress === ip && p.client !== socketId);
 
-	if (previousSocket) {
-		logger.info(`ws client ${socketId} has reached the concurrent IP limit.`, { message: previousSocket.ipAddress });
-		throw new ProbeError('ip limit');
+		if (previousProbe) {
+			logger.info(`ws client ${socketId} has reached the concurrent IP limit.`, { message: previousProbe.ipAddress });
+			throw new ProbeError('ip limit');
+		}
 	}
-};
+}
+

src/lib/ws/server.ts

@@ -7,6 +7,7 @@ import { getRedisClient } from '../redis/client.js';
 import { SyncedProbeList } from './synced-probe-list.js';
 import { client } from '../sql/client.js';
 import { AdoptedProbes } from '../adopted-probes.js';
+import { ProbeIpLimit } from './helper/probe-ip-limit.js';
 
 export type SocketData = {
 	probe: Probe;
@@ -90,3 +91,5 @@ export const fetchRawProbes = async (): Promise<Probe[]> => {
 };
 
 export const adoptedProbes = new AdoptedProbes(client, fetchRawProbes);
+
+export const probeIpLimit = new ProbeIpLimit(fetchProbes, fetchRawSockets);

src/probe/builder.ts

@@ -17,7 +17,7 @@ import type GeoipClient from '../lib/geoip/client.js';
 import getProbeIp from '../lib/get-probe-ip.js';
 import { getRegion } from '../lib/ip-ranges.js';
 import type { Probe, ProbeLocation, Tag } from './types.js';
-import { verifyIpLimit } from '../lib/ws/helper/probe-ip-limit.js';
+import { probeIpLimit } from '../lib/ws/server.js';
 import { fakeLookup } from '../lib/geoip/fake-client.js';
 
 let geoipClient: GeoipClient;
@@ -62,7 +62,7 @@ export const buildProbe = async (socket: Socket): Promise<Probe> => {
 		throw new Error(`couldn't detect probe location for ip ${ip}`);
 	}
 
-	await verifyIpLimit(ip, socket.id);
+	await probeIpLimit.verifyIpLimit(ip, socket.id);
 
 	const location = getLocation(ipInfo);
 

test/tests/unit/ws/probe-ip-limit.test.ts

@@ -0,0 +1,70 @@
+import * as sinon from 'sinon';
+import { expect } from 'chai';
+
+import { ProbeIpLimit } from '../../../../src/lib/ws/helper/probe-ip-limit.js';
+
+describe('ProbeIpLimit', () => {
+	const sandbox = sinon.createSandbox();
+	const fetchProbes = sandbox.stub();
+	const fetchRawSockets = sandbox.stub();
+
+	const getSocket = (id: string, ip: string) => ({
+		id,
+		data: { probe: { client: id, ipAddress: ip } },
+		disconnect: sandbox.stub(),
+	});
+
+	it('syncIpLimit should disconnect duplicates', async () => {
+		const socket1 = getSocket('a', '1.1.1.1');
+		const socket2 = getSocket('b', '2.2.2.2');
+		const duplicate = getSocket('c', '2.2.2.2');
+
+		fetchProbes.resolves([
+			socket1.data.probe,
+			socket2.data.probe,
+			duplicate.data.probe,
+		]);
+
+		fetchRawSockets.resolves([
+			socket1,
+			socket2,
+			duplicate,
+		]);
+
+		const probeIpLimit = new ProbeIpLimit(fetchProbes, fetchRawSockets);
+		await probeIpLimit.syncIpLimit();
+
+		expect(socket1.disconnect.callCount).to.equal(0);
+		expect(socket2.disconnect.callCount).to.equal(0);
+		expect(duplicate.disconnect.callCount).to.equal(1);
+	});
+
+	it('syncIpLimit should preserve socket with the earliest id', async () => {
+		const socket1 = getSocket('a', '1.1.1.1');
+		const socket2 = getSocket('b', '2.2.2.2');
+		const duplicate1 = getSocket('c', '2.2.2.2');
+		const duplicate2 = getSocket('d', '2.2.2.2');
+
+		fetchProbes.resolves([
+			socket1.data.probe,
+			duplicate1.data.probe,
+			socket2.data.probe,
+			duplicate2.data.probe,
+		]);
+
+		fetchRawSockets.resolves([
+			socket1,
+			duplicate1,
+			socket2,
+			duplicate2,
+		]);
+
+		const probeIpLimit = new ProbeIpLimit(fetchProbes, fetchRawSockets);
+		await probeIpLimit.syncIpLimit();
+
+		expect(socket1.disconnect.callCount).to.equal(0);
+		expect(socket2.disconnect.callCount).to.equal(0);
+		expect(duplicate1.disconnect.callCount).to.equal(1);
+		expect(duplicate2.disconnect.callCount).to.equal(1);
+	});
+});