feat: credits support (#489)

* feat: add credits consuming

* test: fix existing tests

* test: add int tests

* test: add unit tests

* feat: support credits headers

* fix: fix tests

* test: add more tests

* refactor: rate limiter

* docs: update openapi

* fix: header tests

* fix: required defaults to false

* docs: update descriptions

* fix: stricter check

* fix: simplify and improve requiredPoints calculation

* feat: add credits-required headers

* docs: update headers docs

* feat: separate limits for authenticated and anon users

* fix: location limits sum <= global limit

---------

Co-authored-by: Martin Kolárik <[email protected]>

config/default.cjs

@@ -56,8 +56,10 @@ module.exports = {
 		// measurement result TTL in redis
 		resultTTL: 7 * 24 * 60 * 60, // 7 days
 		limits: {
-			global: 500,
-			location: 200,
+			anonymousTestsPerLocation: 200,
+			anonymousTestsPerMeasurement: 500,
+			authenticatedTestsPerLocation: 500,
+			authenticatedTestsPerMeasurement: 500,
 		},
 		globalDistribution: {
 			AF: 5,

migrations/create-tables.js.sql

@@ -44,3 +44,12 @@ CREATE TABLE IF NOT EXISTS gp_tokens (
 	expire DATE,
 	date_last_used DATE
 );
+
+CREATE TABLE IF NOT EXISTS gp_credits (
+  id INT(10) UNSIGNED AUTO_INCREMENT PRIMARY KEY,
+  date_created TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+	AMOUNT INT,
+	user_id VARCHAR(36) NOT NULL,
+	CONSTRAINT gp_credits_user_id_unique UNIQUE (user_id),
+	CONSTRAINT gp_credits_amount_positive CHECK (`amount` >= 0)
+);

public/v1/components/headers.yaml

@@ -1,19 +1,38 @@
 components:
   headers:
+    CreditsConsumed:
+      description: The number of credits consumed by the request. Returned only when credits were successfully consumed.
+      required: false
+      schema:
+        type: integer
+    CreditsRequired:
+      description: The number of credits required to process the request. Returned only when the credits in your account were not sufficient, and the request was rejected.
+      required: false
+      schema:
+        type: integer
+    CreditsRemaining:
+      description: The number of credits remaining. Returned only when an attempt to use credits was made (requests with a valid token exceeding the hourly rate limit).
+      required: false
+      schema:
+        type: integer
     MeasurementLocation:
       description: A link to the newly created measurement.
+      required: true
       schema:
         type: string
         format: uri
     RateLimitLimit:
       description: The number of requests available in a given time window.
+      required: true
       schema:
         type: integer
     RateLimitRemaining:
       description: The number of requests remaining in the current time window.
+      required: true
       schema:
         type: integer
     RateLimitReset:
       description: The number of seconds until the limit is reset.
+      required: true
       schema:
         type: integer

public/v1/components/responses.yaml

@@ -117,6 +117,12 @@ components:
           $ref: 'headers.yaml#/components/headers/RateLimitRemaining'
         X-RateLimit-Reset:
           $ref: 'headers.yaml#/components/headers/RateLimitReset'
+        X-Credits-Consumed:
+          $ref: 'headers.yaml#/components/headers/CreditsConsumed'
+        X-Credits-Required:
+          $ref: 'headers.yaml#/components/headers/CreditsRequired'
+        X-Credits-Remaining:
+          $ref: 'headers.yaml#/components/headers/CreditsRemaining'
       content:
         application/json:
           schema:

src/lib/credits.ts

@@ -0,0 +1,45 @@
+import type { Knex } from 'knex';
+import { client } from './sql/client.js';
+
+export const CREDITS_TABLE = 'gp_credits';
+const ER_CONSTRAINT_FAILED_CODE = 4025;
+
+export class Credits {
+	constructor (private readonly sql: Knex) {}
+
+	async consume (userId: string, credits: number): Promise<{ isConsumed: boolean, remainingCredits: number }> {
+		let numberOfUpdates = null;
+
+		try {
+			numberOfUpdates = await this.sql(CREDITS_TABLE).where({ user_id: userId }).update({ amount: this.sql.raw('amount - ?', [ credits ]) });
+		} catch (error) {
+			if (error && (error as Error & {errno?: number}).errno === ER_CONSTRAINT_FAILED_CODE) {
+				const remainingCredits = await this.getRemainingCredits(userId);
+				return { isConsumed: false, remainingCredits };
+			}
+
+			throw error;
+		}
+
+		if (numberOfUpdates === 0) {
+			return { isConsumed: false, remainingCredits: 0 };
+		}
+
+		const remainingCredits = await this.getRemainingCredits(userId);
+		return { isConsumed: true, remainingCredits };
+	}
+
+	async getRemainingCredits (userId: string): Promise<number> {
+		const result = await this.sql(CREDITS_TABLE).where({ user_id: userId }).select<[{ amount: number }]>('amount');
+
+		const remainingCredits = result[0]?.amount;
+
+		if (remainingCredits || remainingCredits === 0) {
+			return remainingCredits;
+		}
+
+		throw new Error('Credits data for the user not found.');
+	}
+}
+
+export const credits = new Credits(client);

src/lib/http/middleware/validate.ts

@@ -2,7 +2,7 @@ import type { Schema } from 'joi';
 import type { ExtendedMiddleware } from '../../../types.js';
 
 export const validate = (schema: Schema): ExtendedMiddleware => async (ctx, next) => {
-	const valid = schema.validate(ctx.request.body, { convert: true });
+	const valid = schema.validate(ctx.request.body, { convert: true, context: ctx.state });
 
 	if (valid.error) {
 		const fields = valid.error.details.map(field => [ field.path.join('.'), String(field?.message) ]);

src/lib/malware/client.ts

@@ -42,14 +42,8 @@ export const joiValidate = (value: string, helpers?: CustomHelpers): string | Er
 	return String(value);
 };
 
-export const joiErrorMessage = (field = 'address'): string => `Provided ${field} is blacklisted.`;
-
-export const joiSchemaErrorMessage = (field?: string): Record<string, string> => {
-	const message = joiErrorMessage(field);
-
-	return {
-		'ip.blacklisted': message,
-		'domain.blacklisted': message,
-		'any.blacklisted': message,
-	};
+export const joiSchemaErrorMessages = {
+	'ip.blacklisted': `Provided address is blacklisted.`,
+	'domain.blacklisted': `Provided address is blacklisted.`,
+	'any.blacklisted': `Provided address is blacklisted.`,
 };

src/lib/rate-limiter.ts

@@ -1,10 +1,10 @@
 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';
+import { credits } from './credits.js';
 
 const redisClient = await createPersistentRedisClient({ legacyMode: true });
 
@@ -43,6 +43,19 @@ export const rateLimit = async (ctx: ExtendedContext, numberOfProbes: number) =>
 		setRateLimitHeaders(ctx, result, rateLimiter);
 	} catch (error) {
 		if (error instanceof RateLimiterRes) {
+			if (ctx.state.userId) {
+				const { isConsumed, requiredCredits, remainingCredits } = await consumeCredits(ctx.state.userId, error, numberOfProbes);
+
+				if (isConsumed) {
+					const result = await rateLimiter.reward(id, requiredCredits);
+					setConsumedHeaders(ctx, requiredCredits, remainingCredits);
+					setRateLimitHeaders(ctx, result, rateLimiter);
+					return;
+				}
+
+				setRequiredHeaders(ctx, requiredCredits, remainingCredits);
+			}
+
 			const result = await rateLimiter.reward(id, numberOfProbes);
 			setRateLimitHeaders(ctx, result, rateLimiter);
 			throw createHttpError(429, 'Too Many Probes Requested', { type: 'too_many_probes' });
@@ -52,8 +65,30 @@ export const rateLimit = async (ctx: ExtendedContext, numberOfProbes: number) =>
 	}
 };
 
-const setRateLimitHeaders = (ctx: Context, result: RateLimiterRes, rateLimiter: RateLimiterRedis) => {
+const consumeCredits = async (userId: string, rateLimiterRes: RateLimiterRes, numberOfProbes: number) => {
+	const freeCredits = config.get<number>('measurement.authenticatedRateLimit');
+	const requiredCredits = Math.min(rateLimiterRes.consumedPoints - freeCredits, numberOfProbes);
+	const { isConsumed, remainingCredits } = await credits.consume(userId, requiredCredits);
+
+	return {
+		isConsumed,
+		requiredCredits,
+		remainingCredits,
+	};
+};
+
+const setRateLimitHeaders = (ctx: ExtendedContext, 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}`);
 };
+
+const setConsumedHeaders = (ctx: ExtendedContext, consumedCredits: number, remainingCredits: number) => {
+	ctx.set('X-Credits-Consumed', `${consumedCredits}`);
+	ctx.set('X-Credits-Remaining', `${remainingCredits}`);
+};
+
+const setRequiredHeaders = (ctx: ExtendedContext, requiredCredits: number, remainingCredits: number) => {
+	ctx.set('X-Credits-Required', `${requiredCredits}`);
+	ctx.set('X-Credits-Remaining', `${remainingCredits}`);
+};

src/measurement/schema/command-schema.ts

@@ -1,6 +1,6 @@
 import Joi from 'joi';
 import {
-	joiSchemaErrorMessage as joiMalwareSchemaErrorMessage,
+	joiSchemaErrorMessages as joiMalwareSchemaErrorMessages,
 } from '../../lib/malware/client.js';
 import {
 	joiValidateDomain,
@@ -11,7 +11,7 @@ import {
 } from './utils.js';
 
 export const schemaErrorMessages = {
-	...joiMalwareSchemaErrorMessage(),
+	...joiMalwareSchemaErrorMessages,
 	'ip.private': 'Private hostnames are not allowed.',
 	'domain.invalid': 'Provided target is not a valid domain name',
 };

src/measurement/schema/global-schema.ts

@@ -7,13 +7,18 @@ import {
 import { schema as locationSchema } from './location-schema.js';
 import { GLOBAL_DEFAULTS } from './utils.js';
 
-const measurementConfig = config.get<{limits: {global: number; location: number}}>('measurement');
+const authenticatedTestsPerMeasurement = config.get<number>('measurement.limits.authenticatedTestsPerMeasurement');
+const anonymousTestsPerMeasurement = config.get<number>('measurement.limits.anonymousTestsPerMeasurement');
 
 export const schema = Joi.object({
 	type: Joi.string().valid('ping', 'traceroute', 'dns', 'mtr', 'http').insensitive().required(),
 	target: targetSchema,
 	measurementOptions: measurementSchema,
 	locations: locationSchema,
-	limit: Joi.number().min(1).max(measurementConfig.limits.global).default(GLOBAL_DEFAULTS.limit),
+	limit: Joi.number().min(1).when('$userId', {
+		is: Joi.exist(),
+		then: Joi.number().max(authenticatedTestsPerMeasurement),
+		otherwise: Joi.number().max(anonymousTestsPerMeasurement),
+	}).default(GLOBAL_DEFAULTS.limit),
 	inProgressUpdates: Joi.bool().default(GLOBAL_DEFAULTS.inProgressUpdates),
 });

src/measurement/schema/location-schema.ts

@@ -1,17 +1,31 @@
 
 import anyAscii from 'any-ascii';
-import Joi from 'joi';
+import Joi, { CustomHelpers, ErrorReport } from 'joi';
 import config from 'config';
 
 import { continents, countries } from 'countries-list';
 import { states } from '../../lib/location/states.js';
 import { regionNames } from '../../lib/location/regions.js';
 import { GLOBAL_DEFAULTS } from './utils.js';
+import type { LocationWithLimit } from '../types.js';
 
-const measurementConfig = config.get<{limits: {global: number; location: number}}>('measurement');
+const authenticatedTestsPerMeasurement = config.get<number>('measurement.limits.authenticatedTestsPerMeasurement');
+const anonymousTestsPerMeasurement = config.get<number>('measurement.limits.anonymousTestsPerMeasurement');
+const authenticatedTestsPerLocation = config.get<number>('measurement.limits.authenticatedTestsPerLocation');
+const anonymousTestsPerLocation = config.get<number>('measurement.limits.anonymousTestsPerLocation');
 
 const normalizeValue = (value: string): string => anyAscii(value);
 
+export const sumOfLocationsLimits = (code: string, max: number) => (value: LocationWithLimit[], helpers: CustomHelpers): LocationWithLimit[] | ErrorReport => {
+	const sum = value.reduce((sum, location) => sum + (location.limit || 1), 0);
+
+	if (sum > max) {
+		return helpers.error(code);
+	}
+
+	return value;
+};
+
 export const schema = Joi.alternatives().try(
 	Joi.string(),
 	Joi.array().items(Joi.object().keys({
@@ -27,10 +41,22 @@ export const schema = Joi.alternatives().try(
 		asn: Joi.number().integer().positive(),
 		magic: Joi.string().min(1).custom(normalizeValue),
 		tags: Joi.array().items(Joi.string().min(1).max(128).lowercase().custom(normalizeValue)),
-		limit: Joi.number().min(1).max(measurementConfig.limits.location).when(Joi.ref('/limit'), {
+		limit: Joi.number().min(1).when('$userId', {
+			is: Joi.exist(),
+			then: Joi.number().max(authenticatedTestsPerLocation),
+			otherwise: Joi.number().max(anonymousTestsPerLocation),
+		}).when(Joi.ref('/limit'), {
 			is: Joi.exist(),
 			then: Joi.forbidden().messages({ 'any.unknown': 'limit per location is not allowed when a global limit is set' }),
 			otherwise: Joi.number().default(1),
 		}),
-	}).or('continent', 'region', 'country', 'state', 'city', 'network', 'asn', 'magic', 'tags')),
+	}).or('continent', 'region', 'country', 'state', 'city', 'network', 'asn', 'magic', 'tags'))
+		.when('$userId', {
+			is: Joi.exist(),
+			then: Joi.custom(sumOfLocationsLimits('limits.sum.auth', authenticatedTestsPerMeasurement)),
+			otherwise: Joi.custom(sumOfLocationsLimits('limits.sum.anon', anonymousTestsPerMeasurement)),
+		}).messages({
+			'limits.sum.auth': `Sum of limits must be less than or equal to ${authenticatedTestsPerMeasurement}`,
+			'limits.sum.anon': `Sum of limits must be less than or equal to ${anonymousTestsPerMeasurement}`,
+		}),
 ).default(GLOBAL_DEFAULTS.locations);

test/plugins/oas/index.js

@@ -130,6 +130,12 @@ export default async ({ specPath, ajvBodyOptions = {}, ajvHeadersOptions = {} })
 					return;
 				}
 
+				const headerOptional = dereference(responseHeadersSpec, header, 'required') !== true;
+
+				if (response.headers[header.toLowerCase()] === undefined && headerOptional) {
+					return;
+				}
+
 				new chai.Assertion(response.headers, `expected the response to have a header ${header}`).to.have.property(header.toLowerCase());
 
 				let responseHeaderValue = response.headers[header.toLowerCase()];