From 7e4282b386ad7bf349e6ee5e0dc58a09c0d23dbb Mon Sep 17 00:00:00 2001 From: Laura Hausmann Date: Sun, 28 Jul 2024 22:06:56 +0200 Subject: [PATCH] [backend/drive] Switch to stream processing for remote media This makes sure that files larger than the configured maximum remote media cache size are not loaded into memory (if the size is known), or are only loaded into memory until the configured maximum size before getting discarded (if the size is not known) --- .../Core/Configuration/Config.cs | 16 +++--- .../Core/Extensions/ServiceExtensions.cs | 1 - .../Core/Extensions/StreamExtensions.cs | 31 ++++++++++++ .../Core/Services/CustomHttpClient.cs | 12 ----- .../Core/Services/DriveService.cs | 49 +++++++++++++++++-- 5 files changed, 83 insertions(+), 26 deletions(-) create mode 100644 Iceshrimp.Backend/Core/Extensions/StreamExtensions.cs diff --git a/Iceshrimp.Backend/Core/Configuration/Config.cs b/Iceshrimp.Backend/Core/Configuration/Config.cs index 448e7173..31e984d0 100644 --- a/Iceshrimp.Backend/Core/Configuration/Config.cs +++ b/Iceshrimp.Backend/Core/Configuration/Config.cs @@ -63,8 +63,8 @@ public sealed class Config public sealed class StorageSection { - public readonly int? MaxCacheSizeBytes; - public readonly int? MaxUploadSizeBytes; + public readonly long? MaxCacheSizeBytes; + public readonly long? MaxUploadSizeBytes; public readonly TimeSpan? MediaRetentionTimeSpan; public bool CleanAvatars = false; @@ -134,9 +134,9 @@ public sealed class Config MaxUploadSizeBytes = suffix switch { null => num, - 'k' or 'K' => num * 1024, - 'm' or 'M' => num * 1024 * 1024, - 'g' or 'G' => num * 1024 * 1024 * 1024, + 'k' or 'K' => num * 1024L, + 'm' or 'M' => num * 1024L * 1024, + 'g' or 'G' => num * 1024L * 1024 * 1024, _ => throw new Exception("Unsupported suffix, use one of: [K]ilobytes, [M]egabytes, [G]igabytes") }; @@ -165,9 +165,9 @@ public sealed class Config MaxCacheSizeBytes = suffix switch { null => num, - 'k' or 'K' => num * 1024, - 'm' or 'M' => num * 1024 * 1024, - 'g' or 'G' => num * 1024 * 1024 * 1024, + 'k' or 'K' => num * 1024L, + 'm' or 'M' => num * 1024L * 1024, + 'g' or 'G' => num * 1024L * 1024 * 1024, _ => throw new Exception("Unsupported suffix, use one of: [K]ilobytes, [M]egabytes, [G]igabytes") }; diff --git a/Iceshrimp.Backend/Core/Extensions/ServiceExtensions.cs b/Iceshrimp.Backend/Core/Extensions/ServiceExtensions.cs index 7e18a8be..5d3cbfcb 100644 --- a/Iceshrimp.Backend/Core/Extensions/ServiceExtensions.cs +++ b/Iceshrimp.Backend/Core/Extensions/ServiceExtensions.cs @@ -86,7 +86,6 @@ public static class ServiceExtensions // Singleton = instantiated once across application lifetime services .AddSingleton() - .AddSingleton() .AddSingleton() .AddSingleton() .AddSingleton() diff --git a/Iceshrimp.Backend/Core/Extensions/StreamExtensions.cs b/Iceshrimp.Backend/Core/Extensions/StreamExtensions.cs new file mode 100644 index 00000000..e9551068 --- /dev/null +++ b/Iceshrimp.Backend/Core/Extensions/StreamExtensions.cs @@ -0,0 +1,31 @@ +using System.Buffers; + +namespace Iceshrimp.Backend.Core.Extensions; + +public static class StreamExtensions +{ + public static async Task CopyToAsync( + this Stream source, Stream destination, long? maxLength, CancellationToken cancellationToken + ) + { + var buffer = ArrayPool.Shared.Rent(81920); + try + { + int bytesRead; + var totalBytesRead = 0L; + while ((maxLength == null || totalBytesRead <= maxLength) && (bytesRead = await DoRead()) != 0) + { + totalBytesRead += bytesRead; + await destination.WriteAsync(new ReadOnlyMemory(buffer, 0, bytesRead), cancellationToken); + } + } + finally + { + ArrayPool.Shared.Return(buffer); + } + + return; + + ValueTask DoRead() => source.ReadAsync(new Memory(buffer), cancellationToken); + } +} \ No newline at end of file diff --git a/Iceshrimp.Backend/Core/Services/CustomHttpClient.cs b/Iceshrimp.Backend/Core/Services/CustomHttpClient.cs index 12cb31bd..afd9a7fd 100644 --- a/Iceshrimp.Backend/Core/Services/CustomHttpClient.cs +++ b/Iceshrimp.Backend/Core/Services/CustomHttpClient.cs @@ -365,16 +365,4 @@ public class CustomHttpClient : HttpClient } } } -} - -public class UnrestrictedHttpClient : CustomHttpClient -{ - public UnrestrictedHttpClient( - IOptions options, - IOptionsMonitor security, - ILoggerFactory loggerFactory - ) : base(options, security, loggerFactory) - { - MaxResponseContentBufferSize = int.MaxValue; - } } \ No newline at end of file diff --git a/Iceshrimp.Backend/Core/Services/DriveService.cs b/Iceshrimp.Backend/Core/Services/DriveService.cs index 119f71b8..9b7d61e5 100644 --- a/Iceshrimp.Backend/Core/Services/DriveService.cs +++ b/Iceshrimp.Backend/Core/Services/DriveService.cs @@ -1,3 +1,4 @@ +using System.Buffers; using System.Diagnostics.CodeAnalysis; using Iceshrimp.Backend.Core.Configuration; using Iceshrimp.Backend.Core.Database; @@ -17,7 +18,7 @@ public class DriveService( [SuppressMessage("ReSharper", "SuggestBaseTypeForParameterInConstructor")] IOptionsSnapshot storageConfig, IOptions instanceConfig, - UnrestrictedHttpClient httpClient, + HttpClient httpClient, QueueService queueSvc, ILogger logger, ImageProcessor imageProcessor @@ -89,7 +90,7 @@ public class DriveService( try { - var res = await httpClient.GetAsync(uri); + var res = await httpClient.GetAsync(uri, HttpCompletionOption.ResponseHeadersRead); res.EnsureSuccessStatusCode(); var filename = res.Content.Headers.ContentDisposition?.FileName ?? @@ -104,7 +105,15 @@ public class DriveService( MimeType = CleanMimeType(mimeType ?? res.Content.Headers.ContentType?.MediaType) }; - return await StoreFile(await res.Content.ReadAsStreamAsync(), user, request, skipImageProcessing); + var input = await res.Content.ReadAsStreamAsync(); + var maxLength = user.IsLocalUser + ? storageConfig.Value.MaxUploadSizeBytes + : storageConfig.Value.MediaRetentionTimeSpan != null + ? storageConfig.Value.MaxCacheSizeBytes + : 0; + + var stream = await GetSafeStreamOrNullAsync(input, maxLength, res.Content.Headers.ContentLength); + return await StoreFile(stream, user, request, skipImageProcessing); } catch (Exception e) { @@ -147,8 +156,7 @@ public class DriveService( DriveFile? file; request.Filename = request.Filename.Trim('"'); - - if (user.IsRemoteUser && input.Length > storageConfig.Value.MaxCacheSizeBytes) + if (input == Stream.Null || user.IsRemoteUser && input.Length > storageConfig.Value.MaxCacheSizeBytes) { file = new DriveFile { @@ -415,6 +423,37 @@ public class DriveService( ? "application/octet-stream" : mimeType; } + + /// + /// We can't trust the Content-Length header, and it might be null. + /// This makes sure that we only ever read up to maxLength into memory. + /// + /// The response content stream + /// The maximum length to buffer (null = unlimited) + /// The content length, if known + /// A CancellationToken, if applicable + /// Either a buffered MemoryStream, or Stream.Null + private static async Task GetSafeStreamOrNullAsync( + Stream stream, long? maxLength, long? contentLength, CancellationToken token = default + ) + { + if (maxLength is 0) return Stream.Null; + if (contentLength > maxLength) return Stream.Null; + + MemoryStream buf = new(); + if (contentLength < maxLength) + maxLength = contentLength.Value; + + await stream.CopyToAsync(buf, maxLength, token); + if (maxLength == null || buf.Length <= maxLength) + { + buf.Seek(0, SeekOrigin.Begin); + return buf; + } + + await buf.DisposeAsync(); + return Stream.Null; + } } public class DriveFileCreationRequest