From d21fb75592c691a884869df56884765b4420d9fc Mon Sep 17 00:00:00 2001
From: Laura Hausmann <laura@hausmann.dev>
Date: Wed, 20 Nov 2024 04:21:51 +0100
Subject: [PATCH] Apply rate limits to proxyServer and fileServer

This resolves a DoS / DDoS / request amplification attack vector that is being actively exploited.

Signed-off-by: limepotato <limepot@protonmail.ch>
---
 .../src/server/file/send-drive-file.ts        | 30 ++++++++++++++++-
 .../backend/src/server/proxy/proxy-media.ts   | 32 ++++++++++++++++++-
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/packages/backend/src/server/file/send-drive-file.ts b/packages/backend/src/server/file/send-drive-file.ts
index 2482f0ce6..8d6dc0f03 100644
--- a/packages/backend/src/server/file/send-drive-file.ts
+++ b/packages/backend/src/server/file/send-drive-file.ts
@@ -14,7 +14,10 @@ import { detectType } from "@/misc/get-file-info.js";
 import { convertToWebp } from "@/services/drive/image-processor.js";
 import { GenerateVideoThumbnail } from "@/services/drive/generate-video-thumbnail.js";
 import { StatusError } from "@/misc/fetch.js";
-import { FILE_TYPE_BROWSERSAFE } from "@/const.js";
+import { FILE_TYPE_BROWSERSAFE, MINUTE } from "@/const.js";
+import { IEndpointMeta } from "@/server/api/endpoints.js";
+import { getIpHash } from "@/misc/get-ip-hash.js";
+import { limiter } from "@/server/api/limiter.js";
 
 const _filename = fileURLToPath(import.meta.url);
 const _dirname = dirname(_filename);
@@ -31,6 +34,31 @@ const commonReadableHandlerGenerator =
 export default async function (ctx: Koa.Context) {
 	const key = ctx.params.key;
 
+	// koa will automatically load the `X-Forwarded-For` header if `proxy: true` is configured in the app.
+	let limitActor: string;
+	limitActor = getIpHash(ctx.ip);
+
+	const limit: IEndpointMeta["limit"] = {
+		key: `drive-file:${key}`,
+		duration: MINUTE * 10,
+		max: 10
+	}
+
+	// Rate limit
+	await limiter(
+		limit as IEndpointMeta["limit"] & { key: NonNullable<string> },
+		limitActor,
+	).catch((e) => {
+		const remainingTime = e.remainingTime
+			? `Please try again in ${e.remainingTime}.`
+			: "Please try again later.";
+
+		ctx.status = 429;
+		ctx.body = "Rate limit exceeded. " + remainingTime;
+	});
+
+	if (ctx.status == 429) return;
+
 	// Fetch drive file
 	const file = await DriveFiles.createQueryBuilder("file")
 		.where("file.accessKey = :accessKey", { accessKey: key })
diff --git a/packages/backend/src/server/proxy/proxy-media.ts b/packages/backend/src/server/proxy/proxy-media.ts
index b3bb03124..af802c939 100644
--- a/packages/backend/src/server/proxy/proxy-media.ts
+++ b/packages/backend/src/server/proxy/proxy-media.ts
@@ -9,9 +9,12 @@ import { createTemp } from "@/misc/create-temp.js";
 import { downloadUrl } from "@/misc/download-url.js";
 import { detectType } from "@/misc/get-file-info.js";
 import { StatusError } from "@/misc/fetch.js";
-import { FILE_TYPE_BROWSERSAFE } from "@/const.js";
+import { FILE_TYPE_BROWSERSAFE, MINUTE } from "@/const.js";
 import { serverLogger } from "../index.js";
 import { isMimeImage } from "@/misc/is-mime-image.js";
+import { getIpHash } from "@/misc/get-ip-hash.js";
+import { limiter } from "@/server/api/limiter.js";
+import { IEndpointMeta } from "@/server/api/endpoints.js";
 
 export async function proxyMedia(ctx: Koa.Context) {
 	const url = "url" in ctx.query ? ctx.query.url : `https://${ctx.params.url}`;
@@ -21,6 +24,33 @@ export async function proxyMedia(ctx: Koa.Context) {
 		return;
 	}
 
+	// koa will automatically load the `X-Forwarded-For` header if `proxy: true` is configured in the app.
+	let limitActor: string;
+	limitActor = getIpHash(ctx.ip);
+
+	const parsedUrl = new URL(url);
+
+	const limit: IEndpointMeta["limit"] = {
+		key: `media-proxy:${parsedUrl.host}:${parsedUrl.pathname}`,
+		duration: MINUTE * 10,
+		max: 10
+	}
+
+	// Rate limit
+	await limiter(
+		limit as IEndpointMeta["limit"] & { key: NonNullable<string> },
+		limitActor,
+	).catch((e) => {
+		const remainingTime = e.remainingTime
+			? `Please try again in ${e.remainingTime}.`
+			: "Please try again later.";
+
+		ctx.status = 429;
+		ctx.body = "Rate limit exceeded. " + remainingTime;
+	});
+
+	if (ctx.status == 429) return;
+
 	const { hostname } = new URL(url);
 	let resolvedIps;
 	try {