From 7cac813a24a39f18b675aaafd6dd9bcae8be6fdb Mon Sep 17 00:00:00 2001
From: Alice Gaudon <alice@gaudon.pro>
Date: Fri, 24 Apr 2020 10:54:21 +0200
Subject: [PATCH] Add throttling utils

See https://gitlab.com/ArisuOngaku/wms/issues/35
---
 package.json     |  2 +-
 src/HttpError.ts | 15 +++++++++
 src/Throttler.ts | 82 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 1 deletion(-)
 create mode 100644 src/Throttler.ts

diff --git a/package.json b/package.json
index a3bb867..6253763 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "wms-core",
-  "version": "0.2.7",
+  "version": "0.2.8",
   "description": "Node web framework",
   "repository": "git@gitlab.com:ArisuOngaku/wms-core.git",
   "author": "Alice Gaudon <alice@gaudon.pro>",
diff --git a/src/HttpError.ts b/src/HttpError.ts
index de03928..0b8c232 100644
--- a/src/HttpError.ts
+++ b/src/HttpError.ts
@@ -58,6 +58,21 @@ export class NotFoundHttpError extends BadRequestError {
     }
 }
 
+export class TooManyRequestsHttpError extends BadRequestError {
+    constructor(cause?: Error) {
+        super(
+            `You're making too many requests!`,
+            `We need some rest.`,
+            '',
+            cause
+        );
+    }
+
+    get errorCode(): number {
+        return 404;
+    }
+}
+
 export class ServerError extends HttpError {
     constructor(message: string, cause?: Error) {
         super(message, `Maybe you should contact us; see instructions below.`, cause);
diff --git a/src/Throttler.ts b/src/Throttler.ts
new file mode 100644
index 0000000..0619faf
--- /dev/null
+++ b/src/Throttler.ts
@@ -0,0 +1,82 @@
+import {TooManyRequestsHttpError} from "./HttpError";
+
+export default class Throttler {
+    private static readonly throttles: { [throttleName: string]: Throttle } = {};
+
+    /**
+     * Throttle function; will throw a TooManyRequestsHttpError when the threshold is reached.
+     *
+     * This throttle is adaptive: it will slowly decrease (linear) until it reaches 0 after {@param resetPeriod} ms.
+     * Threshold will hold for {@param holdPeriod} ms.
+     *
+     * @param action a unique action name (can be used multiple times, but it'll account for a single action).
+     * @param max how many times this action can be triggered per id.
+     * @param resetPeriod after how much time in ms the throttle will reach 0.
+     * @param id an identifier of who triggered the action.
+     * @param holdPeriod time in ms after each call before the threshold begins to decrease.
+     * @param jailPeriod time in ms for which the throttle will throw when it is triggered.
+     */
+    public static throttle(action: string, max: number, resetPeriod: number, id: string, holdPeriod: number = 100, jailPeriod: number = 30 * 1000) {
+        let throttle = this.throttles[action];
+        if (!throttle) throttle = this.throttles[action] = new Throttle(max, resetPeriod, holdPeriod, jailPeriod);
+        throttle.trigger(id);
+    }
+
+    private constructor() {
+    }
+}
+
+class Throttle {
+    private readonly max: number;
+    private readonly resetPeriod: number;
+    private readonly holdPeriod: number;
+    private readonly jailPeriod: number;
+    private readonly triggers: {
+        [id: string]: {
+            count: number,
+            lastTrigger?: number,
+            jailed?: number;
+        }
+    } = {};
+
+    constructor(max: number, resetPeriod: number, holdPeriod: number, jailPeriod: number) {
+        this.max = max;
+        this.resetPeriod = resetPeriod;
+        this.holdPeriod = holdPeriod;
+        this.jailPeriod = jailPeriod;
+    }
+
+
+    public trigger(id: string) {
+        let trigger = this.triggers[id];
+        if (!trigger) trigger = this.triggers[id] = {count: 0};
+
+        let currentDate = new Date().getTime();
+
+        if (trigger.jailed && currentDate - trigger.jailed > 0) {
+            this.throw();
+            return;
+        }
+
+        if (trigger.lastTrigger) {
+            let timeDiff = currentDate - trigger.lastTrigger;
+            if (timeDiff > this.holdPeriod) {
+                timeDiff -= this.holdPeriod;
+                trigger.count = Math.floor(Math.min(trigger.count, this.max * (1 - timeDiff / this.resetPeriod)));
+            }
+        }
+
+        trigger.count++;
+        trigger.lastTrigger = currentDate;
+
+        if (trigger.count > this.max) {
+            trigger.jailed = currentDate;
+            this.throw();
+            return;
+        }
+    }
+
+    private throw() {
+        throw new TooManyRequestsHttpError();
+    }
+}
\ No newline at end of file