From 1df2ad13dbe133ea21734822a16c26d4913a8e2b Mon Sep 17 00:00:00 2001
From: Yonle <yonle@lecturify.net>
Date: Sun, 19 May 2024 23:09:49 +0700
Subject: [PATCH] config: max_conn_per_ip for limiting maximum incomming
 connections per IP.

Signed-off-by: Yonle <yonle@lecturify.net>
---
 bouncer.js        | 4 +++-
 config.js.example | 4 ++++
 http.js           | 7 ++++++-
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/bouncer.js b/bouncer.js
index a9eaca4..6cd15d0 100644
--- a/bouncer.js
+++ b/bouncer.js
@@ -43,7 +43,7 @@ let zeroStats = {
 }
 let stats = {};
 
-function handleConnection(ws, req) {
+function handleConnection(ws, req, onClose) {
   let query = querystring.parse(req.url.slice(2));
   let authKey = null;
   let authorized = true;
@@ -147,6 +147,8 @@ function handleConnection(ws, req) {
 
   ws.on('error', console.error);
   ws.on('close', _ => {
+    onClose();
+
     delete idents[ws.ident];
 
     console.log(process.pid, "---", `${ws.ip} disconnected`);
diff --git a/config.js.example b/config.js.example
index 255d6a7..0bb92ff 100644
--- a/config.js.example
+++ b/config.js.example
@@ -53,6 +53,10 @@ module.exports = {
   // Setting as 0 will disable ratelimit handling.
   upstream_ratelimit_expiration: 10000,
 
+  // Maximum incomming connections per IP.
+  // By default, This is Infinity. Change the value as Integer (number) to override.
+  max_conn_per_ip: Infinity,
+
   // Maximum subscriptions that client could open.
   // Setting as -1 will disable max subscription limit.
   max_client_subs: -1,
diff --git a/http.js b/http.js
index e1faa8d..991a545 100644
--- a/http.js
+++ b/http.js
@@ -18,6 +18,8 @@ const log = _ => console.log(process.pid, curD(), "-", _);
 let server = null;
 let config = require(process.env.BOSTR_CONFIG_PATH || "./config");
 
+let connectedHosts = [];
+
 let wslinkregex = /(?:^- )(wss?:\/\/.*)(?: \(.*\))/gm;
 let loadbalancerUpstreamLinks = [];
 
@@ -147,8 +149,11 @@ server.on('upgrade', (req, sock, head) => {
   const ip = req.headers["x-forwarded-for"]?.split(",")[0] || sock.address()?.address;
 
   if (config.blocked_hosts && config.blocked_hosts.includes(ip)) return sock.destroy();
+  if (connectedHosts.filter(i => i === ip).length >= (config.max_conn_per_ip || Infinity)) return sock.destroy();
 
-  wss.handleUpgrade(req, sock, head, _ => bouncer.handleConnection(_, req));
+  connectedHosts.push(ip);
+
+  wss.handleUpgrade(req, sock, head, _ => bouncer.handleConnection(_, req, _ => delete connectedHosts[connectedHosts.indexOf(ip)]));
 });
 
 const listened = server.listen(process.env.PORT || config.port, config.address || "0.0.0.0", _ => {