update Rate Limit with new event transmission

2 files changed, 47 insertions, 46 deletions

diff --git a/api/src/Server.ts b/api/src/Server.ts
index 332c28de..ef72df5f 100644
--- a/api/src/Server.ts
+++ b/api/src/Server.ts
@@ -64,7 +64,7 @@ export class FosscordServer extends Server {
 		this.app = api;
 
 		api.use(Authentication);
-		initRateLimits(api);
+		await initRateLimits(api);
 		await initTranslation(api);
 
 		this.routes = await this.registerRoutes(path.join(__dirname, "routes", "/"));
diff --git a/api/src/middlewares/RateLimit.ts b/api/src/middlewares/RateLimit.ts
index a1bb0c44..37c54969 100644
--- a/api/src/middlewares/RateLimit.ts
+++ b/api/src/middlewares/RateLimit.ts
@@ -1,30 +1,16 @@
 // @ts-nocheck
-import { db, Bucket, Config } from "@fosscord/util";
+import { db, Bucket, Config, listenEvent, emitEvent } from "@fosscord/util";
 import { NextFunction, Request, Response, Router } from "express";
 import { getIpAdress } from "../util/ipAddress";
 import { API_PREFIX_TRAILING_SLASH } from "./Authentication";
 
-// const Cache = new MongooseCache(
-// 	db.collection("ratelimits"),
-// 	[
-// 		// TODO: uncomment $match and fix error: not receiving change events
-// 		// { $match: { blocked: true } }
-// 	],
-// 	{ onlyEvents: false, array: true }
-// );
-
 // Docs: https://discord.com/developers/docs/topics/rate-limits
 
 /*
 ? bucket limit? Max actions/sec per bucket?
 
-TODO: ip rate limit
-TODO: user rate limit
-TODO: different rate limit for bots/user/oauth/webhook
 TODO: delay database requests to include multiple queries
 TODO: different for methods (GET/POST)
-TODO: bucket major parameters (channel_id, guild_id, webhook_id)
-TODO: use config values
 
 > IP addresses that make too many invalid HTTP requests are automatically and temporarily restricted from accessing the Discord API. Currently, this limit is 10,000 per 10 minutes. An invalid request is one that results in 401, 403, or 429 statuses.
 
@@ -32,6 +18,9 @@ TODO: use config values
 
 */
 
+var Cache = new Map<string, Bucket>();
+const EventRateLimit = "ratelimit";
+
 // TODO: FIX with new event handling
 export default function RateLimit(opts: {
 	bucket?: string;
@@ -46,9 +35,6 @@ export default function RateLimit(opts: {
 	success?: boolean;
 	onlyIp?: boolean;
 }): any {
-	return (req, res, next) => next();
-	Cache.init(); // will only initalize it once
-
 	return async (req: Request, res: Response, next: NextFunction): Promise<any> => {
 		const bucket_id = opts.bucket || req.originalUrl.replace(API_PREFIX_TRAILING_SLASH, "");
 		var user_id = getIpAdress(req);
@@ -59,7 +45,7 @@ export default function RateLimit(opts: {
 		if (opts.GET && ["GET", "OPTIONS", "HEAD"].includes(req.method)) max_hits = opts.GET;
 		else if (opts.MODIFY && ["POST", "DELETE", "PATCH", "PUT"].includes(req.method)) max_hits = opts.MODIFY;
 
-		const offender = Cache.data?.find((x: Bucket) => x.user_id == user_id && x.id === bucket_id) as Bucket | null;
+		const offender = Cache.get(user_id + bucket_id) as Bucket | null;
 
 		if (offender && offender.blocked) {
 			const reset = offender.expires_at.getTime();
@@ -88,6 +74,7 @@ export default function RateLimit(opts: {
 				offender.blocked = false;
 				// mongodb ttl didn't update yet -> manually update/delete
 				db.collection("ratelimits").updateOne({ id: bucket_id, user_id }, { $set: offender });
+				Cache.delete(user_id + bucket_id);
 			}
 		}
 		next();
@@ -108,8 +95,18 @@ export default function RateLimit(opts: {
 	};
 }
 
-export function initRateLimits(app: Router) {
+export async function initRateLimits(app: Router) {
 	const { routes, global, ip, error } = Config.get().limits.rate;
+	await listenEvent(EventRateLimit, (event) => {
+		Cache.set(event.channel_id, event.data);
+		event.acknowledge?.();
+	});
+
+	setInterval(() => {
+		Cache.forEach((x, key) => {
+			if (Date.now() > x.expires_at) Cache.delete(key);
+		});
+	}, 1000 * 60 * 10);
 
 	app.use(
 		RateLimit({
@@ -134,32 +131,36 @@ export function initRateLimits(app: Router) {
 	app.use("/auth/register", RateLimit({ onlyIp: true, success: true, ...routes.auth.register }));
 }
 
-function hitRoute(opts: { user_id: string; bucket_id: string; max_hits: number; window: number }) {
-	return db.collection("ratelimits").updateOne(
-		{ id: opts.bucket_id, user_id: opts.user_id },
-		[
-			{
-				$replaceRoot: {
-					newRoot: {
-						// similar to $setOnInsert
-						$mergeObjects: [
-							{
-								id: opts.bucket_id,
-								user_id: opts.user_id,
-								expires_at: new Date(Date.now() + opts.window * 1000)
-							},
-							"$$ROOT"
-						]
-					}
-				}
+async function hitRoute(opts: { user_id: string; bucket_id: string; max_hits: number; window: number }) {
+	const filter = { id: opts.bucket_id, user_id: opts.user_id };
+	const { value } = await db.collection("ratelimits").findOneAndUpdate(
+		filter,
+		{
+			$setOnInsert: {
+				id: opts.bucket_id,
+				user_id: opts.user_id,
+				expires_at: new Date(Date.now() + opts.window * 1000)
 			},
-			{
-				$set: {
-					hits: { $sum: [{ $ifNull: ["$hits", 0] }, 1] },
-					blocked: { $gte: ["$hits", opts.max_hits] }
-				}
+			$inc: {
+				hits: 1
 			}
-		],
-		{ upsert: true }
+			// Conditionally update blocked doesn't work
+		},
+		{ upsert: true, returnDocument: "before" }
 	);
+	if (!value) return;
+	const updateBlock = !value.blocked && value.hits >= opts.max_hits;
+
+	if (updateBlock) {
+		value.blocked = true;
+		Cache.set(opts.user_id + opts.bucket_id, value);
+		await emitEvent({
+			channel_id: EventRateLimit,
+			event: EventRateLimit,
+			data: value
+		});
+		await db.collection("ratelimits").updateOne(filter, { $set: { blocked: true } });
+	} else {
+		Cache.delete(opts.user_id);
+	}
 }