From 461ede13669890750c974c90af90f59e191f7dca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Kol=C3=A1rik?= <martin@kolarik.sk>
Date: Sun, 11 Feb 2024 13:51:21 +0100
Subject: [PATCH] fix: simplify and improve requiredPoints calculation

---
 src/lib/rate-limiter.ts                  |  4 +---
 test/tests/integration/ratelimit.test.ts | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/lib/rate-limiter.ts b/src/lib/rate-limiter.ts
index 300d99c3..fa6df39a 100644
--- a/src/lib/rate-limiter.ts
+++ b/src/lib/rate-limiter.ts
@@ -65,9 +65,7 @@ export const rateLimit = async (ctx: ExtendedContext, numberOfProbes: number) =>
 
 const consumeCredits = async (userId: string, rateLimiterRes: RateLimiterRes, numberOfProbes: number) => {
 	const freePoints = config.get<number>('measurement.authenticatedRateLimit');
-	const alreadyUsedPoints = rateLimiterRes.consumedPoints - numberOfProbes;
-	const remainingFreePoints = freePoints - alreadyUsedPoints;
-	const requiredPoints = numberOfProbes - remainingFreePoints;
+	const requiredPoints = Math.min(rateLimiterRes.consumedPoints - freePoints, numberOfProbes);
 	const { isConsumed, remainingCredits } = await credits.consume(userId, requiredPoints);
 
 	if (isConsumed) {
diff --git a/test/tests/integration/ratelimit.test.ts b/test/tests/integration/ratelimit.test.ts
index 93f9b3e2..e72b50bd 100644
--- a/test/tests/integration/ratelimit.test.ts
+++ b/test/tests/integration/ratelimit.test.ts
@@ -302,6 +302,24 @@ describe('rate limiter', () => {
 			expect(amount).to.equal(1);
 		});
 
+		it('should not consume more paid credits than the cost of the full request', async () => {
+			await authenticatedRateLimiter.set('89da69bd-a236-4ab7-9c5d-b5f52ce09959', 255, 0);
+
+			const response = await requestAgent.post('/v1/measurements')
+				.set('Authorization', 'Bearer v2lUHEVLtVSskaRKDBabpyp4AkzdMnob')
+				.send({
+					type: 'ping',
+					target: 'jsdelivr.com',
+					limit: 2,
+				}).expect(202) as Response;
+
+			expect(response.headers['x-ratelimit-remaining']).to.equal('0');
+			expect(response.headers['x-credits-cost']).to.equal('2');
+			expect((response.headers['x-credits-remaining'])).to.equal('8');
+			const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
+			expect(amount).to.equal(8);
+		});
+
 		it('should not consume free credits if there are not enough to satisfy the request', async () => {
 			await authenticatedRateLimiter.set('89da69bd-a236-4ab7-9c5d-b5f52ce09959', 249, 0);