From 3fa193ec569caaa639f0dfd4371d6d1659e3abc3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=A9my=20Rialland?= <wiinxt@gmail.com>
Date: Wed, 31 Aug 2022 10:55:04 +0200
Subject: [PATCH] Add rate limit information in header

---
 api/source/rate-limiter.ts | 32 ++++++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/api/source/rate-limiter.ts b/api/source/rate-limiter.ts
index 15ae5a589..fbb1ca676 100644
--- a/api/source/rate-limiter.ts
+++ b/api/source/rate-limiter.ts
@@ -1,6 +1,10 @@
-import { BaseContext, Next } from 'koa'
-import { RateLimiterMemory, RateLimiterRedis } from 'rate-limiter-flexible'
 import IORedis from 'ioredis'
+import { BaseContext, Next } from 'koa'
+import {
+	RateLimiterMemory,
+	RateLimiterRedis,
+	RateLimiterRes,
+} from 'rate-limiter-flexible'
 
 const Redis = IORedis.default
 
@@ -26,8 +30,32 @@ export const rateLimiterMiddleware = async (ctx: BaseContext, next: Next) => {
 		ctx.status = 429
 		ctx.body = 'Too Many Requests'
 
+		if (isRateLimiterRes(rejRes)) {
+			ctx.set({
+				'Retry-After': (rejRes.msBeforeNext / 1000).toString(),
+				'X-RateLimit-Limit': rateLimiter.points.toString(),
+				'X-RateLimit-Remaining': rejRes.remainingPoints.toString(),
+				'X-RateLimit-Reset': new Date(
+					Date.now() + rejRes.msBeforeNext
+				).toString(),
+			})
+		}
+
 		return
 	}
 
 	await next()
 }
+
+const isRateLimiterRes = (val: unknown): val is RateLimiterRes => {
+	return !!(
+		val &&
+		typeof val === 'object' &&
+		[
+			'msBeforeNext',
+			'remainingPoints',
+			'consumedPoints',
+			'isFirstInDuration',
+		].every((s) => s in val)
+	)
+}