feat: separate limits for authenticated users (#483)

* docs: document auth tokens

* refactor: allow auth only when creating measurements

* feat: add is authenticated mw

* feat: add token validation

* feat: save userId to ctx

* refactor: rename rate limiter

* feat: add auth rate limiter

* feat: update date_last_used

* test: fix tests

* feat: add int tests

* test: add auth tests

* test: add unit tests

* fix: remove args from deleteFakeProbes

* fix: small fixes

* fix: store userId in state

* refactor: use ExtendedContext definitions

* refactor: add userId to state type

* refactor: move AuthenticateState to the middleware

---------

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

config/default.cjs

@@ -47,7 +47,8 @@ module.exports = {
 		fetchSocketsCacheTTL: 1000,
 	},
 	measurement: {
-		rateLimit: 100000,
+		anonymousRateLimit: 100000,
+		authenticatedRateLimit: 250,
 		rateLimitReset: 3600,
 		maxInProgressProbes: 5,
 		// Timeout after which measurement will be marked as finished even if not all probes respond

migrations/create-tables.js.sql

@@ -1,3 +1,8 @@
+CREATE TABLE IF NOT EXISTS directus_users (
+  id CHAR(36),
+  github_username VARCHAR(255)
+);
+
 CREATE TABLE IF NOT EXISTS gp_adopted_probes (
   id INT(10) UNSIGNED AUTO_INCREMENT PRIMARY KEY,
   user_created CHAR(36),
@@ -29,3 +34,13 @@ CREATE TABLE IF NOT EXISTS directus_notifications (
   subject VARCHAR(255),
   message TEXT
 );
+
+CREATE TABLE IF NOT EXISTS gp_tokens (
+  id INT(10) UNSIGNED AUTO_INCREMENT PRIMARY KEY,
+	user_created CHAR(36),
+	name VARCHAR(255),
+	value VARCHAR(255),
+	origins LONGTEXT,
+	expire DATE,
+	date_last_used DATE
+);

package.json

@@ -5,6 +5,7 @@
 	"type": "module",
 	"main": "dist/src/index.js",
 	"dependencies": {
+		"@isaacs/ttlcache": "^1.4.1",
 		"@koa/router": "^12.0.1",
 		"@maxmind/geoip2-node": "^4.2.0",
 		"@redocly/openapi-core": "^1.6.0",

public/v1/spec.yaml

@@ -75,6 +75,9 @@ paths:
           $ref: 'components/responses.yaml#/components/responses/measurements429'
       tags:
         - Measurements
+      security:
+        - {}
+        - bearerAuth: []
   /v1/measurements/{id}:
     parameters:
       - $ref: 'components/parameters.yaml#/components/parameters/measurementId'
@@ -110,3 +113,8 @@ paths:
           $ref: 'components/responses.yaml#/components/responses/probes200'
       tags:
         - Probes
+components:
+  securitySchemes:
+    bearerAuth:
+      type: http
+      scheme: bearer

src/lib/adopted-probes.ts

@@ -142,7 +142,7 @@ export class AdoptedProbes {
 	}
 
 	private async fetchAdoptedProbes () {
-		const rows = await this.sql({ probes: ADOPTED_PROBES_TABLE }).select<Row[]>();
+		const rows = await this.sql(ADOPTED_PROBES_TABLE).select<Row[]>();
 
 
 		const adoptedProbes: AdoptedProbe[] = rows.map(row => ({

src/lib/http/auth.ts

@@ -0,0 +1,137 @@
+import { createHash } from 'node:crypto';
+import type { Knex } from 'knex';
+import TTLCache from '@isaacs/ttlcache';
+import { scopedLogger } from '../logger.js';
+import { client } from '../sql/client.js';
+
+export const GP_TOKENS_TABLE = 'gp_tokens';
+export const USERS_TABLE = 'directus_users';
+
+const logger = scopedLogger('auth');
+
+const TOKEN_TTL = 2 * 60 * 1000;
+
+export type Token = {
+	user_created: string,
+	value: string,
+	expire: Date | null,
+	origins: string[],
+	date_last_used: Date | null
+}
+
+type Row = Omit<Token, 'origins'> & {
+	origins: string | null,
+}
+
+export class Auth {
+	private validTokens = new TTLCache<string, Token>({ ttl: TOKEN_TTL });
+	private invalidTokens = new TTLCache<string, true>({ ttl: TOKEN_TTL });
+	constructor (private readonly sql: Knex) {}
+
+	scheduleSync () {
+		setTimeout(() => {
+			this.syncTokens()
+				.finally(() => this.scheduleSync())
+				.catch(error => logger.error(error));
+		}, 60_000);
+	}
+
+	async syncTokens () {
+		const tokens = await this.fetchTokens();
+		const newValidTokens = new TTLCache<string, Token>({ ttl: TOKEN_TTL });
+		const newInvalidTokens = new TTLCache<string, true>({ ttl: TOKEN_TTL });
+
+		tokens.forEach((token) => {
+			if (token.expire && this.isExpired(token.expire)) {
+				newInvalidTokens.set(token.value, true);
+			} else {
+				newValidTokens.set(token.value, token);
+			}
+		});
+
+		this.validTokens = newValidTokens;
+		this.invalidTokens = newInvalidTokens;
+	}
+
+	async syncSpecificToken (hash: string) {
+		const tokens = await this.fetchTokens({ value: hash });
+
+		if (tokens.length === 0) {
+			this.invalidTokens.set(hash, true);
+			return undefined;
+		}
+
+		const token = tokens[0]!;
+
+		if (token.expire && this.isExpired(token.expire)) {
+			this.invalidTokens.set(hash, true);
+			return undefined;
+		}
+
+		this.validTokens.set(hash, token);
+		return token;
+	}
+
+	async fetchTokens (filter: Partial<Row> = {}) {
+		const rows = await this.sql(GP_TOKENS_TABLE).where(filter)
+			.select<Row[]>([ 'user_created', 'value', 'expire', 'origins', 'date_last_used' ]);
+
+		const tokens: Token[] = rows.map(row => ({
+			...row,
+			origins: (row.origins ? JSON.parse(row.origins) as string[] : []),
+		}));
+
+		return tokens;
+	}
+
+	async validate (tokenString: string, origin: string) {
+		const bytes = Buffer.from(tokenString, 'base64');
+		const hash = createHash('sha256').update(bytes).digest('base64');
+
+		if (this.invalidTokens.get(hash)) {
+			return null;
+		}
+
+		let token = this.validTokens.get(hash);
+
+		if (!token) {
+			token = await this.syncSpecificToken(hash);
+		}
+
+		if (!token) {
+			return null;
+		}
+
+		if (!this.isValidOrigin(origin, token.origins)) {
+			return null;
+		}
+
+		await this.updateLastUsedDate(token);
+		return token.user_created;
+	}
+
+	private async updateLastUsedDate (token: Token) {
+		if (!token.date_last_used || !this.isToday(token.date_last_used)) {
+			const date = new Date();
+			await this.sql(GP_TOKENS_TABLE).where({ value: token.value }).update({ date_last_used: date });
+			token.date_last_used = date;
+		}
+	}
+
+	private isExpired (date: Date) {
+		const currentDate = new Date();
+		currentDate.setHours(0, 0, 0, 0);
+		return date < currentDate;
+	}
+
+	private isValidOrigin (origin: string, validOrigins: string[]) {
+		return validOrigins.length > 0 ? validOrigins.includes(origin) : true;
+	}
+
+	private isToday (date: Date) {
+		const currentDate = new Date();
+		return date.toDateString() === currentDate.toDateString();
+	}
+}
+
+export const auth = new Auth(client);

src/lib/http/middleware/authenticate.ts

@@ -0,0 +1,30 @@
+import { auth } from '../auth.js';
+import type { ExtendedMiddleware } from '../../../types.js';
+
+export const authenticate: ExtendedMiddleware = async (ctx, next) => {
+	const { headers } = ctx.request;
+
+	if (headers && headers.authorization) {
+		const parts = headers.authorization.split(' ');
+
+		if (parts.length !== 2 || parts[0] !== 'Bearer') {
+			ctx.status = 401;
+			return;
+		}
+
+		const token = parts[1]!;
+		const origin = ctx.get('Origin');
+		const userId = await auth.validate(token, origin);
+
+		if (!userId) {
+			ctx.status = 401;
+			return;
+		}
+
+		ctx.state.userId = userId;
+	}
+
+	return next();
+};
+
+export type AuthenticateState = { userId?: string };

src/lib/http/middleware/cors.ts

@@ -2,11 +2,18 @@ import type { Context, Next } from 'koa';
 
 export const corsHandler = () => async (ctx: Context, next: Next) => {
 	ctx.set('Access-Control-Allow-Origin', '*');
-	ctx.set('Access-Control-Allow-Headers', '*, Authorization');
+	ctx.set('Access-Control-Allow-Headers', '*');
 	ctx.set('Access-Control-Expose-Headers', '*');
+	ctx.set('Access-Control-Max-Age', '600');
 	ctx.set('Cross-Origin-Resource-Policy', 'cross-origin');
 	ctx.set('Timing-Allow-Origin', '*');
 	ctx.set('Vary', 'Accept-Encoding');
 
 	return next();
 };
+
+export const corsAuthHandler = () => async (ctx: Context, next: Next) => {
+	ctx.set('Access-Control-Allow-Headers', '*, Authorization');
+
+	return next();
+};

src/lib/rate-limiter.ts

@@ -0,0 +1,59 @@
+import config from 'config';
+import type { Context } from 'koa';
+import { RateLimiterRedis, RateLimiterRes } from 'rate-limiter-flexible';
+import requestIp from 'request-ip';
+import { createPersistentRedisClient } from './redis/persistent-client.js';
+import createHttpError from 'http-errors';
+import type { ExtendedContext } from '../types.js';
+
+const redisClient = await createPersistentRedisClient({ legacyMode: true });
+
+export const anonymousRateLimiter = new RateLimiterRedis({
+	storeClient: redisClient,
+	keyPrefix: 'rate:anon',
+	points: config.get<number>('measurement.anonymousRateLimit'),
+	duration: config.get<number>('measurement.rateLimitReset'),
+});
+
+export const authenticatedRateLimiter = new RateLimiterRedis({
+	storeClient: redisClient,
+	keyPrefix: 'rate:auth',
+	points: config.get<number>('measurement.authenticatedRateLimit'),
+	duration: config.get<number>('measurement.rateLimitReset'),
+});
+
+export const rateLimit = async (ctx: ExtendedContext, numberOfProbes: number) => {
+	if (ctx['isAdmin']) {
+		return;
+	}
+
+	let rateLimiter: RateLimiterRedis;
+	let id: string;
+
+	if (ctx.state.userId) {
+		rateLimiter = authenticatedRateLimiter;
+		id = ctx.state.userId;
+	} else {
+		rateLimiter = anonymousRateLimiter;
+		id = requestIp.getClientIp(ctx.req) ?? '';
+	}
+
+	try {
+		const result = await rateLimiter.consume(id, numberOfProbes);
+		setRateLimitHeaders(ctx, result, rateLimiter);
+	} catch (error) {
+		if (error instanceof RateLimiterRes) {
+			const result = await rateLimiter.reward(id, numberOfProbes);
+			setRateLimitHeaders(ctx, result, rateLimiter);
+			throw createHttpError(429, 'Too Many Probes Requested', { type: 'too_many_probes' });
+		}
+
+		throw createHttpError(500);
+	}
+};
+
+const setRateLimitHeaders = (ctx: Context, result: RateLimiterRes, rateLimiter: RateLimiterRedis) => {
+	ctx.set('X-RateLimit-Reset', `${Math.round(result.msBeforeNext / 1000)}`);
+	ctx.set('X-RateLimit-Limit', `${rateLimiter.points}`);
+	ctx.set('X-RateLimit-Remaining', `${result.remainingPoints}`);
+};