From 67b1c97d8921d894e61a5747af56913205680160 Mon Sep 17 00:00:00 2001
From: Calvin Montgomery <calzoneman@gmail.com>
Date: Wed, 25 Jul 2018 21:07:07 -0700
Subject: [PATCH] Add io.throttle-in-rate-limit for socket event rate

---
 package.json             |  2 +-
 src/config.js            |  6 ++++++
 src/io/ioserver.js       | 23 +++++++++++++++++++++++
 src/util/token-bucket.js | 17 ++++++++++++++---
 4 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/package.json b/package.json
index c1231cdb..bd6b492c 100644
--- a/package.json
+++ b/package.json
@@ -2,7 +2,7 @@
   "author": "Calvin Montgomery",
   "name": "CyTube",
   "description": "Online media synchronizer and chat",
-  "version": "3.56.4",
+  "version": "3.56.5",
   "repository": {
     "url": "http://github.com/calzoneman/sync"
   },
diff --git a/src/config.js b/src/config.js
index 8d08801f..d54f378d 100644
--- a/src/config.js
+++ b/src/config.js
@@ -386,6 +386,12 @@ function preprocessConfig(cfg) {
         return contact.name !== 'calzoneman';
     });
 
+    if (!cfg.io.throttle) {
+        cfg.io.throttle = {
+            'in-rate-limit': Infinity
+        };
+    }
+
     return cfg;
 }
 
diff --git a/src/io/ioserver.js b/src/io/ioserver.js
index d7227394..e69c14d8 100644
--- a/src/io/ioserver.js
+++ b/src/io/ioserver.js
@@ -223,6 +223,8 @@ class IOServer {
             return;
         }
 
+        this.setRateLimiter(socket);
+
         emitMetrics(socket);
 
         LOGGER.info('Accepted socket from %s', socket.context.ipAddress);
@@ -240,6 +242,25 @@ class IOServer {
         }
     }
 
+    setRateLimiter(socket) {
+        const thunk = () => Config.get('io.throttle.in-rate-limit');
+
+        socket._inRateLimit = new TokenBucket(thunk, thunk);
+
+        socket.on('cytube:count-event', () => {
+            if (socket._inRateLimit.throttle()) {
+                LOGGER.warn(
+                    'Kicking client %s: exceeded in-rate-limit of %d',
+                    socket.context.ipAddress,
+                    thunk()
+                );
+
+                socket.emit('kick', { reason: 'Rate limit exceeded' });
+                socket.disconnect();
+            }
+        });
+    }
+
     initSocketIO() {
         patchSocketMetrics();
         patchTypecheckedFunctions();
@@ -277,10 +298,12 @@ const outgoingPacketCount = new Counter({
 function patchSocketMetrics() {
     const onevent = Socket.prototype.onevent;
     const packet = Socket.prototype.packet;
+    const emit = require('events').EventEmitter.prototype.emit;
 
     Socket.prototype.onevent = function patchedOnevent() {
         onevent.apply(this, arguments);
         incomingEventCount.inc(1);
+        emit.call(this, 'cytube:count-event');
     };
 
     Socket.prototype.packet = function patchedPacket() {
diff --git a/src/util/token-bucket.js b/src/util/token-bucket.js
index c71e97a5..1fb84025 100644
--- a/src/util/token-bucket.js
+++ b/src/util/token-bucket.js
@@ -1,16 +1,27 @@
 class TokenBucket {
     constructor(capacity, refillRate) {
+        if (typeof refillRate !== 'function') {
+            const _refillRate = refillRate;
+            refillRate = () => _refillRate;
+        }
+        if (typeof capacity !== 'function') {
+            const _capacity = capacity;
+            capacity = () => _capacity;
+        }
+
         this.capacity = capacity;
         this.refillRate = refillRate;
-        this.count = capacity;
+        this.count = capacity();
         this.lastRefill = Date.now();
     }
 
     throttle() {
         const now = Date.now();
-        const delta = Math.floor((now - this.lastRefill) / 1000 * this.refillRate);
+        const delta = Math.floor(
+            (now - this.lastRefill) / 1000 * this.refillRate()
+        );
         if (delta > 0) {
-            this.count = Math.min(this.capacity, this.count + delta);
+            this.count = Math.min(this.capacity(), this.count + delta);
             this.lastRefill = now;
         }