From b50121d0a1409ac5b8d821526cf5c5de22032f7c Mon Sep 17 00:00:00 2001
From: Laura Hausmann <laura@hausmann.dev>
Date: Thu, 5 Sep 2024 22:08:50 +0200
Subject: [PATCH] [backend/api-shared] Fix upload of files larger than 28MB.

---
 .../Controllers/Mastodon/MediaController.cs   |  1 +
 .../MaxRequestSizeIsMaxUploadSize.cs          | 24 +++++++++++++++++++
 .../Controllers/Web/DriveController.cs        |  1 +
 .../Core/Middleware/ErrorHandlerMiddleware.cs |  3 +++
 4 files changed, 29 insertions(+)
 create mode 100644 Iceshrimp.Backend/Controllers/Shared/Attributes/MaxRequestSizeIsMaxUploadSize.cs

diff --git a/Iceshrimp.Backend/Controllers/Mastodon/MediaController.cs b/Iceshrimp.Backend/Controllers/Mastodon/MediaController.cs
index 424596eb..36823afd 100644
--- a/Iceshrimp.Backend/Controllers/Mastodon/MediaController.cs
+++ b/Iceshrimp.Backend/Controllers/Mastodon/MediaController.cs
@@ -25,6 +25,7 @@ namespace Iceshrimp.Backend.Controllers.Mastodon;
 [Produces(MediaTypeNames.Application.Json)]
 public class MediaController(DriveService driveSvc, DatabaseContext db) : ControllerBase
 {
+	[MaxRequestSizeIsMaxUploadSize]
 	[HttpPost("/api/v1/media")]
 	[HttpPost("/api/v2/media")]
 	[ProducesResults(HttpStatusCode.OK)]
diff --git a/Iceshrimp.Backend/Controllers/Shared/Attributes/MaxRequestSizeIsMaxUploadSize.cs b/Iceshrimp.Backend/Controllers/Shared/Attributes/MaxRequestSizeIsMaxUploadSize.cs
new file mode 100644
index 00000000..359f74d6
--- /dev/null
+++ b/Iceshrimp.Backend/Controllers/Shared/Attributes/MaxRequestSizeIsMaxUploadSize.cs
@@ -0,0 +1,24 @@
+using Iceshrimp.Backend.Core.Configuration;
+using Iceshrimp.Backend.Core.Middleware;
+using Microsoft.AspNetCore.Http.Features;
+using Microsoft.AspNetCore.Mvc.Filters;
+using Microsoft.Extensions.Options;
+
+namespace Iceshrimp.Backend.Controllers.Shared.Attributes;
+
+public class MaxRequestSizeIsMaxUploadSize : Attribute, IResourceFilter
+{
+	public void OnResourceExecuting(ResourceExecutingContext context)
+	{
+		var feature = context.HttpContext.Features.Get<IHttpMaxRequestBodySizeFeature>() ??
+		              throw new Exception("Failed to get IHttpMaxRequestBodySizeFeature");
+		var config = context.HttpContext.RequestServices.GetRequiredService<IOptionsSnapshot<Config.StorageSection>>();
+		feature.MaxRequestBodySize = config.Value.MaxUploadSizeBytes;
+
+		if (context.HttpContext.Request.ContentLength > config.Value.MaxUploadSizeBytes)
+			throw GracefulException.RequestEntityTooLarge("Attachment is too large.",
+			                                              $"The media upload size limit is set to {config.Value.MaxUploadSizeBytes} bytes.");
+	}
+
+	public void OnResourceExecuted(ResourceExecutedContext context) { }
+}
\ No newline at end of file
diff --git a/Iceshrimp.Backend/Controllers/Web/DriveController.cs b/Iceshrimp.Backend/Controllers/Web/DriveController.cs
index a74ef194..7503e577 100644
--- a/Iceshrimp.Backend/Controllers/Web/DriveController.cs
+++ b/Iceshrimp.Backend/Controllers/Web/DriveController.cs
@@ -82,6 +82,7 @@ public class DriveController(
 	[Authorize]
 	[Produces(MediaTypeNames.Application.Json)]
 	[ProducesResults(HttpStatusCode.OK)]
+	[MaxRequestSizeIsMaxUploadSize]
 	public async Task<DriveFileResponse> UploadFile(IFormFile file)
 	{
 		var user = HttpContext.GetUserOrFail();
diff --git a/Iceshrimp.Backend/Core/Middleware/ErrorHandlerMiddleware.cs b/Iceshrimp.Backend/Core/Middleware/ErrorHandlerMiddleware.cs
index ec7494f6..f100f7c6 100644
--- a/Iceshrimp.Backend/Core/Middleware/ErrorHandlerMiddleware.cs
+++ b/Iceshrimp.Backend/Core/Middleware/ErrorHandlerMiddleware.cs
@@ -231,6 +231,9 @@ public class GracefulException(
 
 	public static GracefulException Conflict(string message, string? details = null) =>
 		new(HttpStatusCode.Conflict, message, details);
+	
+	public static GracefulException RequestEntityTooLarge(string message, string? details = null) =>
+		new(HttpStatusCode.RequestEntityTooLarge, message, details);
 
 	public static GracefulException RecordNotFound() => new(HttpStatusCode.NotFound, "Record not found");