s3: re-enable uploads with rate limiting

Notesnook.API/Controllers/S3Controller.cs

@@ -20,10 +20,12 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 using System;
 using System.Net.Http;
 using System.Security.Claims;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Authorization;
 using Microsoft.AspNetCore.Http.Extensions;
 using Microsoft.AspNetCore.Mvc;
+using Microsoft.AspNetCore.RateLimiting;
 using Microsoft.Extensions.Logging;
 using MongoDB.Driver;
 using Notesnook.API.Helpers;
@@ -32,7 +34,6 @@ using Notesnook.API.Models;
 using Streetwriters.Common;
 using Streetwriters.Common.Accessors;
 using Streetwriters.Common.Extensions;
-using Streetwriters.Common.Models;
 
 namespace Notesnook.API.Controllers
 {
@@ -40,39 +41,45 @@ namespace Notesnook.API.Controllers
     [Route("s3")]
     [ResponseCache(NoStore = true, Location = ResponseCacheLocation.None)]
     [Authorize("Sync")]
-    public class S3Controller(IS3Service s3Service, ISyncItemsRepositoryAccessor repositories, WampServiceAccessor serviceAccessor, ILogger<S3Controller> logger) : ControllerBase
+    public class S3Controller(IS3Service s3Service, ISyncItemsRepositoryAccessor repositories, WampServiceAccessor serviceAccessor, IHttpClientFactory httpClientFactory, ILogger<S3Controller> logger) : ControllerBase
     {
         [HttpPut]
+        [EnableRateLimiting("s3-direct")]
         public async Task<IActionResult> Upload([FromQuery] string name)
         {
-            return BadRequest(new { error = "Attachment storage is temporarily unavailable. Please try again later." });
-            // try
-            // {
-            //     var userId = this.User.GetUserId();
+            try
+            {
+                var userId = this.User.GetUserId();
 
-            //     var fileSize = HttpContext.Request.ContentLength ?? 0;
-            //     bool hasBody = fileSize > 0;
+                var fileSize = HttpContext.Request.ContentLength ?? 0;
+                bool hasBody = fileSize > 0;
 
-            //     if (!hasBody)
-            //     {
-            //         return Ok(Request.GetEncodedUrl() + "&access_token=" + Request.Headers.Authorization.ToString().Replace("Bearer ", ""));
-            //     }
+                if (!hasBody)
+                {
+                    return Ok(Request.GetEncodedUrl() + "&access_token=" + Request.Headers.Authorization.ToString().Replace("Bearer ", ""));
+                }
 
-            //     if (Constants.IS_SELF_HOSTED) await UploadFileAsync(userId, name, fileSize);
-            //     else await UploadFileWithChecksAsync(userId, name, fileSize);
+                if (Constants.IS_SELF_HOSTED) await UploadFileAsync(userId, name, fileSize);
+                else await UploadFileWithChecksAsync(userId, name, fileSize);
 
-            //     return Ok();
-            // }
-            // catch (Exception ex)
-            // {
-            //     logger.LogError(ex, "Error uploading attachment for user.");
-            //     return BadRequest(new { error = "Failed to upload attachment." });
-            // }
+                return Ok();
+            }
+            catch (Exception ex) when (S3TransientErrorClassifier.IsTransient(ex))
+            {
+                return StorageUnavailable(ex);
+            }
+            catch (Exception ex)
+            {
+                logger.LogError(ex, "Error uploading attachment for user.");
+                return BadRequest(new { error = "Failed to upload attachment." });
+            }
         }
 
         private async Task UploadFileWithChecksAsync(string userId, string name, long fileSize)
         {
-            var userSettings = await repositories.UsersSettings.FindOneAsync((u) => u.UserId == userId);
+            var userSettings = await repositories.UsersSettings.FindOneAsync((u) => u.UserId == userId)
+                ?? throw new Exception("User settings not found.");
+            var observedStorageLimit = userSettings.StorageLimit;
 
             var subscription = await serviceAccessor.UserSubscriptionService.GetUserSubscriptionAsync(Clients.Notesnook.Id, userId) ?? throw new Exception("User subscription not found.");
 
@@ -85,49 +92,68 @@ namespace Notesnook.API.Controllers
 
             var uploadedFileSize = await UploadFileAsync(userId, name, fileSize);
 
-            userSettings.StorageLimit.Value += uploadedFileSize;
-            await repositories.UsersSettings.Collection.UpdateOneAsync(
-                Builders<UserSettings>.Filter.Eq(u => u.UserId, userId),
-                Builders<UserSettings>.Update.Set(u => u.StorageLimit, userSettings.StorageLimit)
-            );
-
-            // extra check in case user sets wrong ContentLength in the HTTP header
-            if (uploadedFileSize != fileSize && StorageHelper.IsStorageLimitReached(subscription, userSettings.StorageLimit.Value))
+            try
             {
-                await s3Service.DeleteObjectAsync(userId, name);
-                throw new Exception("Storage limit exceeded.");
+                await s3Service.IncrementStorageUsageAsync(userId, uploadedFileSize, observedStorageLimit);
             }
+            catch (Exception ex)
+            {
+                logger.LogError(ex, "Failed to account for attachment usage for user {UserId}.", userId);
+            }
+
         }
 
         private async Task<long> UploadFileAsync(string userId, string name, long fileSize)
         {
             var url = await s3Service.GetInternalUploadObjectUrlAsync(userId, name) ?? throw new Exception("Could not create signed url.");
 
-            var httpClient = new HttpClient();
+            var httpClient = httpClientFactory.CreateClient("S3Upload");
             var content = new StreamContent(HttpContext.Request.BodyReader.AsStream());
             content.Headers.ContentLength = fileSize;
-            var response = await httpClient.SendRequestAsync<Response>(url, null, HttpMethod.Put, content);
-            if (!response.Success) throw new Exception(response.Content != null ? await response.Content.ReadAsStringAsync() : "Could not upload file.");
+            using var request = new HttpRequestMessage(HttpMethod.Put, url)
+            {
+                Content = content
+            };
+            using var uploadTimeout = CancellationTokenSource.CreateLinkedTokenSource(HttpContext.RequestAborted);
+            uploadTimeout.CancelAfter(TimeSpan.FromMinutes(15 + (15 * fileSize / (1024d * 1024d * 1024d))));
+            using var response = await httpClient.SendAsync(
+                request,
+                HttpCompletionOption.ResponseHeadersRead,
+                uploadTimeout.Token);
+            if (!response.IsSuccessStatusCode)
+            {
+                var statusCode = (int)response.StatusCode;
+                if (S3TransientErrorClassifier.IsTransient(statusCode))
+                    throw new S3StorageUnavailableException("Attachment storage is temporarily unavailable.");
 
-            return await s3Service.GetObjectSizeAsync(userId, name);
+                throw new Exception(await response.Content.ReadAsStringAsync(uploadTimeout.Token));
+            }
+
+            // The PUT response confirms this request. A follow-up HEAD can lag or
+            // be throttled, so direct accounting uses the accepted request bytes.
+            return fileSize;
         }
 
 
         [HttpGet("multipart")]
+        [EnableRateLimiting("s3-multipart-control")]
         public async Task<IActionResult> MultipartUpload([FromQuery] string name, [FromQuery] int parts, [FromQuery] string? uploadId)
         {
-            return BadRequest(new { error = "Attachment storage is temporarily unavailable. Please try again later." });
-            // var userId = this.User.GetUserId();
-            // try
-            // {
-            //     var meta = await s3Service.StartMultipartUploadAsync(userId, name, parts, uploadId);
-            //     return Ok(meta);
-            // }
-            // catch (Exception ex)
-            // {
-            //     logger.LogError(ex, "Error starting multipart upload for user.");
-            //     return BadRequest(new { error = "Failed to start multipart upload." });
-            // }
+            var userId = this.User.GetUserId();
+            try
+            {
+                var meta = await s3Service.StartMultipartUploadAsync(userId, name, parts, uploadId);
+                return Ok(meta);
+            }
+            catch (Exception ex) when (S3TransientErrorClassifier.IsTransient(ex))
+            {
+                return StorageUnavailable(ex);
+            }
+            catch (Exception ex)
+            {
+                logger.LogError(ex, "Error starting multipart upload for user.");
+                return BadRequest(new { error = "Failed to start multipart upload." });
+            }
         }
 
         [HttpDelete("multipart")]
@@ -139,6 +165,10 @@ namespace Notesnook.API.Controllers
                 await s3Service.AbortMultipartUploadAsync(userId, name, uploadId);
                 return Ok();
             }
+            catch (Exception ex) when (S3TransientErrorClassifier.IsTransient(ex))
+            {
+                return StorageUnavailable(ex);
+            }
             catch (Exception ex)
             {
                 logger.LogError(ex, "Error aborting multipart upload for user.");
@@ -147,6 +177,7 @@ namespace Notesnook.API.Controllers
         }
 
         [HttpPost("multipart")]
+        [EnableRateLimiting("s3-multipart-control")]
         public async Task<IActionResult> CompleteMultipartUpload([FromBody] CompleteMultipartUploadRequestWrapper uploadRequestWrapper)
         {
             var userId = this.User.GetUserId();
@@ -155,6 +186,10 @@ namespace Notesnook.API.Controllers
                 await s3Service.CompleteMultipartUploadAsync(userId, uploadRequestWrapper.ToRequest());
                 return Ok();
             }
+            catch (Exception ex) when (S3TransientErrorClassifier.IsTransient(ex))
+            {
+                return StorageUnavailable(ex);
+            }
             catch (Exception ex)
             {
                 logger.LogError(ex, "Error completing multipart upload for user.");
@@ -162,6 +197,12 @@ namespace Notesnook.API.Controllers
             }
         }
 
+        private IActionResult StorageUnavailable(Exception exception)
+        {
+            logger.LogWarning(exception, "Attachment storage is temporarily unavailable.");
+            return StatusCode(503, new { error = "Attachment storage is temporarily unavailable. Please try again later." });
+        }
+
         [HttpGet]
         public async Task<IActionResult> Download([FromQuery] string name)
         {

Notesnook.API/Helpers/S3FailoverHelper.cs

@@ -20,6 +20,8 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Net.Http;
+using System.Net.Sockets;
 using System.Threading.Tasks;
 using Amazon.S3;
 using Amazon.S3.Model;
@@ -47,12 +49,6 @@ namespace Notesnook.API.Helpers
         /// </summary>
         public bool UseExponentialBackoff { get; set; } = true;
 
-        /// <summary>
-        /// Whether to allow failover for write operations (PUT, POST, DELETE).
-        /// Default is false to prevent data consistency issues.
-        /// </summary>
-        public bool AllowWriteFailover { get; set; } = false;
-
         /// <summary>
         /// List of exception types that should trigger failover
         /// </summary>
@@ -77,6 +73,39 @@ namespace Notesnook.API.Helpers
         };
     }
 
+    public sealed class S3StorageUnavailableException : Exception
+    {
+        public S3StorageUnavailableException(string message, Exception? innerException = null)
+            : base(message, innerException)
+        {
+        }
+    }
+
+    public static class S3TransientErrorClassifier
+    {
+        public static bool IsTransient(int statusCode)
+        {
+            return statusCode == 408 || statusCode == 429 || statusCode >= 500;
+        }
+
+        public static bool IsTransient(Exception exception)
+        {
+            if (exception is S3StorageUnavailableException)
+                return true;
+
+            if (exception is AmazonS3Exception s3Exception)
+            {
+                return IsTransient((int)s3Exception.StatusCode)
+                    || s3Exception.ErrorCode is "SlowDown" or "ServiceUnavailable" or "InternalError" or "RequestTimeout";
+            }
+
+            if (exception is HttpRequestException or SocketException or TimeoutException or TaskCanceledException)
+                return true;
+
+            return exception.InnerException != null && IsTransient(exception.InnerException);
+        }
+    }
+
     /// <summary>
     /// Result of a failover operation
     /// </summary>
@@ -161,10 +190,11 @@ namespace Notesnook.API.Helpers
             var result = new S3FailoverResult<T>();
             Exception? lastException = null;
 
-            // Determine max clients to try based on write operation flag
-            var maxClientsToTry = (isWriteOperation && !config.AllowWriteFailover) ? 1 : clients.Count;
+            // Writes stay on the provider that owns the object or multipart upload.
+            // They get one application attempt and never fail over.
+            var maxClientsToTry = isWriteOperation ? 1 : clients.Count;
 
-            if (isWriteOperation && !config.AllowWriteFailover && clients.Count > 1)
+            if (isWriteOperation && clients.Count > 1)
             {
                 logger?.LogDebug(
                     "Write operation {Operation} will only use primary endpoint. Failover is disabled for write operations.",
@@ -185,7 +215,12 @@ namespace Notesnook.API.Helpers
                         operationName, clientName, i + 1, maxClientsToTry);
                 }
 
-                var (success, value, exception, attempts) = await TryExecuteAsync(client, operation, operationName, clientName);
+                var (success, value, exception, attempts) = await TryExecuteAsync(
+                    client,
+                    operation,
+                    operationName,
+                    clientName,
+                    isWriteOperation ? 0 : config.MaxRetries);
                 result.AttemptsUsed += attempts;
 
                 if (success && value != null)
@@ -222,19 +257,22 @@ namespace Notesnook.API.Helpers
                 operationName, maxClientsToTry, result.AttemptsUsed);
 
             return result;
-        }        /// <summary>
-                 /// Try to execute an operation with retries
-                 /// </summary>
+        }
+
+        /// <summary>
+        /// Try to execute an operation with retries
+        /// </summary>
         private async Task<(bool success, T? value, Exception? exception, int attempts)> TryExecuteAsync<T>(
             AmazonS3Client client,
             Func<AmazonS3Client, Task<T>> operation,
             string operationName,
-            string endpointName)
+            string endpointName,
+            int maxRetries)
         {
             Exception? lastException = null;
             int attempts = 0;
 
-            for (int retry = 0; retry <= config.MaxRetries; retry++)
+            for (int retry = 0; retry <= maxRetries; retry++)
             {
                 attempts++;
                 try
@@ -246,12 +284,12 @@ namespace Notesnook.API.Helpers
                 {
                     lastException = ex;
 
-                    if (retry < config.MaxRetries && ShouldRetry(ex))
+                    if (retry < maxRetries && ShouldRetry(ex))
                     {
                         var delay = CalculateRetryDelay(retry);
                         logger?.LogWarning(ex,
                             "Attempt {Attempt}/{MaxAttempts} failed for {Operation} on {Endpoint}. Retrying in {Delay}ms",
-                            retry + 1, config.MaxRetries + 1, operationName, endpointName, delay);
+                            retry + 1, maxRetries + 1, operationName, endpointName, delay);
 
                         await Task.Delay(delay);
                     }
@@ -322,10 +360,10 @@ namespace Notesnook.API.Helpers
             string operationName = "S3Operation",
             bool isWriteOperation = false)
         {
-            await ExecuteWithFailoverAsync<object?>(async (client) =>
+            await ExecuteWithFailoverAsync(async (client) =>
             {
                 await operation(client);
-                return null;
+                return true;
             }, operationName, isWriteOperation);
         }
     }
@@ -384,6 +422,7 @@ namespace Notesnook.API.Helpers
                         ServiceURL = url,
                         AuthenticationRegion = region,
                         ForcePathStyle = forcePathStyle,
+                        MaxErrorRetry = 0,
                         SignatureMethod = Amazon.Runtime.SigningAlgorithm.HmacSHA256,
                         SignatureVersion = "4"
                     };

Notesnook.API/Interfaces/IS3Service.cs

@@ -35,5 +35,6 @@ namespace Notesnook.API.Interfaces
         Task<MultipartUploadMeta> StartMultipartUploadAsync(string userId, string name, int parts, string? uploadId = null);
         Task AbortMultipartUploadAsync(string userId, string name, string uploadId);
         Task CompleteMultipartUploadAsync(string userId, CompleteMultipartUploadRequest uploadRequest);
+        Task IncrementStorageUsageAsync(string userId, long fileSize, Limit? observedStorageLimit);
     }
 }

Notesnook.API/Services/S3Service.cs

@@ -60,9 +60,11 @@ namespace Notesnook.API.Services
         // That is why we create 2 separate S3 clients. One for internal traffic and one for external.
         private readonly S3FailoverHelper S3InternalClient;
         private readonly HttpClient httpClient = new();
+        private readonly ILogger<S3Service> logger;
 
         public S3Service(ISyncItemsRepositoryAccessor syncItemsRepositoryAccessor, WampServiceAccessor wampServiceAccessor, ILogger<S3Service> logger)
         {
+            this.logger = logger;
             Repositories = syncItemsRepositoryAccessor;
             ServiceAccessor = wampServiceAccessor;
             S3Client = new S3FailoverHelper(
@@ -277,12 +279,32 @@ namespace Notesnook.API.Services
 
         private async Task<long> GetMultipartUploadSizeAsync(string userId, string key, string uploadId)
         {
-            var objectName = GetFullObjectName(userId, key);
-            var parts = await S3InternalClient.ExecuteWithFailoverAsync((client) => client.ListPartsAsync(INTERNAL_BUCKET_NAME, objectName, uploadId), operationName: "ListParts");
+            var objectName = GetFullObjectName(userId, key) ?? throw new Exception("Invalid object name.");
             long totalSize = 0;
-            foreach (var part in parts.Parts)
+            var marker = 0;
+            while (true)
             {
-                totalSize += part.Size;
+                var request = new ListPartsRequest
+                {
+                    BucketName = INTERNAL_BUCKET_NAME,
+                    Key = objectName,
+                    UploadId = uploadId,
+                    PartNumberMarker = marker.ToString()
+                };
+                var parts = await S3InternalClient.ExecuteWithFailoverAsync(
+                    (client) => client.ListPartsAsync(request),
+                    operationName: "ListParts",
+                    isWriteOperation: true);
+                foreach (var part in parts.Parts)
+                {
+                    totalSize = checked(totalSize + part.Size);
+                }
+
+                if (!parts.IsTruncated)
+                    break;
+                if (parts.NextPartNumberMarker <= marker)
+                    throw new Exception("S3 returned an invalid parts page.");
+                marker = parts.NextPartNumberMarker;
             }
             return totalSize;
         }
@@ -298,21 +320,22 @@ namespace Notesnook.API.Services
                 await this.AbortMultipartUploadAsync(userId, uploadRequest.Key, uploadRequest.UploadId);
                 throw new Exception("User settings not found.");
             }
-            userSettings.StorageLimit ??= StorageHelper.RolloverStorageLimit(userSettings.StorageLimit);
+            var observedStorageLimit = userSettings.StorageLimit;
+            userSettings.StorageLimit = StorageHelper.RolloverStorageLimit(userSettings.StorageLimit);
+            long fileSize = 0;
 
             if (!Constants.IS_SELF_HOSTED)
             {
                 var subscription = await ServiceAccessor.UserSubscriptionService.GetUserSubscriptionAsync(Clients.Notesnook.Id, userId) ?? throw new Exception("User subscription not found.");
 
-                long fileSize = await GetMultipartUploadSizeAsync(userId, uploadRequest.Key, uploadRequest.UploadId);
+                fileSize = await GetMultipartUploadSizeAsync(userId, uploadRequest.Key, uploadRequest.UploadId);
                 if (StorageHelper.IsFileSizeExceeded(subscription, fileSize))
                 {
                     await this.AbortMultipartUploadAsync(userId, uploadRequest.Key, uploadRequest.UploadId);
                     throw new Exception("Max file size exceeded.");
                 }
 
-                userSettings.StorageLimit.Value += fileSize;
-                if (StorageHelper.IsStorageLimitReached(subscription, userSettings.StorageLimit.Value))
+                if (StorageHelper.IsStorageLimitReached(subscription, userSettings.StorageLimit.Value + fileSize))
                 {
                     await this.AbortMultipartUploadAsync(userId, uploadRequest.Key, uploadRequest.UploadId);
                     throw new Exception("Storage limit reached.");
@@ -326,14 +349,75 @@ namespace Notesnook.API.Services
 
             if (!Constants.IS_SELF_HOSTED)
             {
-                await Repositories.UsersSettings.UpsertAsync(userSettings, (u) => u.UserId == userId);
-                await Repositories.UsersSettings.Collection.UpdateOneAsync(
-                    Builders<UserSettings>.Filter.Eq(u => u.UserId, userId),
-                    Builders<UserSettings>.Update.Set(u => u.StorageLimit, userSettings.StorageLimit)
-                );
+                try
+                {
+                    await IncrementStorageUsageAsync(userId, fileSize, observedStorageLimit);
+                }
+                catch (Exception ex)
+                {
+                    // The object is already committed. Accounting is best effort so
+                    // a Mongo failure cannot turn a successful upload into a retry.
+                    logger.LogError(ex, "Failed to account for multipart attachment usage for user {UserId}.", userId);
+                }
             }
         }
 
+        public async Task IncrementStorageUsageAsync(string userId, long fileSize, Limit? observedStorageLimit)
+        {
+            var collection = Repositories.UsersSettings.Collection;
+            var userFilter = Builders<UserSettings>.Filter.Eq(u => u.UserId, userId);
+
+            for (var attempt = 0; attempt < 3; attempt++)
+            {
+                var now = DateTimeOffset.UtcNow;
+                var monthStart = new DateTimeOffset(now.Year, now.Month, 1, 0, 0, 0, TimeSpan.Zero);
+                var nextMonthStart = monthStart.AddMonths(1);
+                var nowMilliseconds = now.ToUnixTimeMilliseconds();
+                var monthStartMilliseconds = monthStart.ToUnixTimeMilliseconds();
+                var nextMonthStartMilliseconds = nextMonthStart.ToUnixTimeMilliseconds();
+                var currentStorageLimit = observedStorageLimit;
+
+                if (attempt > 0)
+                {
+                    currentStorageLimit = (await collection.Find(userFilter).FirstOrDefaultAsync())?.StorageLimit;
+                }
+
+                var isCurrentMonth = currentStorageLimit != null
+                    && currentStorageLimit.UpdatedAt >= monthStartMilliseconds
+                    && currentStorageLimit.UpdatedAt < nextMonthStartMilliseconds;
+
+                if (!isCurrentMonth)
+                {
+                    var observedFilter = currentStorageLimit == null
+                        ? Builders<UserSettings>.Filter.Eq(u => u.StorageLimit, null)
+                        : Builders<UserSettings>.Filter.Eq("StorageLimit.UpdatedAt", currentStorageLimit.UpdatedAt);
+
+                    var resetResult = await collection.UpdateOneAsync(
+                        Builders<UserSettings>.Filter.And(userFilter, observedFilter),
+                        Builders<UserSettings>.Update.Set(
+                            "StorageLimit",
+                            new Limit { Value = 0, UpdatedAt = nowMilliseconds }));
+
+                    if (resetResult.MatchedCount == 0)
+                        continue;
+                }
+
+                var result = await collection.UpdateOneAsync(
+                    Builders<UserSettings>.Filter.And(
+                        userFilter,
+                        Builders<UserSettings>.Filter.Gte("StorageLimit.UpdatedAt", monthStartMilliseconds),
+                        Builders<UserSettings>.Filter.Lt("StorageLimit.UpdatedAt", nextMonthStartMilliseconds)),
+                    Builders<UserSettings>.Update.Combine(
+                        Builders<UserSettings>.Update.Inc("StorageLimit.Value", fileSize),
+                        Builders<UserSettings>.Update.Set("StorageLimit.UpdatedAt", nowMilliseconds)));
+
+                if (result.MatchedCount > 0)
+                    return;
+            }
+
+            throw new Exception("Storage usage counter was not available after rollover retries.");
+        }
+
         private async Task<string?> GetPresignedURLAsync(string userId, string name, HttpVerb httpVerb, S3ClientMode mode = S3ClientMode.EXTERNAL)
         {
             var objectName = GetFullObjectName(userId, name);
@@ -393,4 +477,4 @@ namespace Notesnook.API.Services
             return ((int)statusCode >= 200) && ((int)statusCode <= 299);
         }
     }
-}
+}

Notesnook.API/Startup.cs

@@ -23,6 +23,7 @@ using System.IO.Compression;
 using System.Security.Claims;
 using System.Text;
 using System.Text.Encodings.Web;
+using System.Threading;
 using System.Threading.Tasks;
 using Amazon.Runtime;
 using StackExchange.Redis;
@@ -35,6 +36,7 @@ using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Connections;
 using Microsoft.AspNetCore.HttpOverrides;
+using Microsoft.AspNetCore.RateLimiting;
 using Microsoft.AspNetCore.ResponseCompression;
 using Microsoft.Extensions.Caching.Distributed;
 using Microsoft.Extensions.Caching.Memory;
@@ -68,6 +70,7 @@ using Streetwriters.Data;
 using Streetwriters.Data.DbContexts;
 using Streetwriters.Data.Interfaces;
 using Streetwriters.Data.Repositories;
+using System.Threading.RateLimiting;
 
 namespace Notesnook.API
 {
@@ -125,6 +128,27 @@ namespace Notesnook.API
                 options.DefaultPolicy = options.GetPolicy("Notesnook") ?? throw new Exception("Notesnook policy not found");
             }).AddSingleton<IAuthorizationMiddlewareResultHandler, AuthorizationResultTransformer>();
 
+            services.AddRateLimiter(options =>
+            {
+                options.RejectionStatusCode = StatusCodes.Status503ServiceUnavailable;
+                options.AddPolicy("s3-direct", _ => RateLimitPartition.GetConcurrencyLimiter(
+                    "s3-direct",
+                    _ => new ConcurrencyLimiterOptions
+                    {
+                        PermitLimit = ReadRateLimit("S3_DIRECT_UPLOAD_CONCURRENCY", 128),
+                        QueueLimit = 0,
+                        QueueProcessingOrder = QueueProcessingOrder.OldestFirst
+                    }));
+                options.AddPolicy("s3-multipart-control", _ => RateLimitPartition.GetConcurrencyLimiter(
+                    "s3-multipart-control",
+                    _ => new ConcurrencyLimiterOptions
+                    {
+                        PermitLimit = ReadRateLimit("S3_MULTIPART_CONTROL_CONCURRENCY", 32),
+                        QueueLimit = 0,
+                        QueueProcessingOrder = QueueProcessingOrder.OldestFirst
+                    }));
+            });
+
             services.AddAuthentication(JwtBearerDefaults.AuthenticationScheme)
             .AddOAuth2Introspection("introspection", options =>
             {
@@ -204,6 +228,10 @@ namespace Notesnook.API
             services.AddScoped<SyncDeviceService>();
             services.AddScoped<IUserService, UserService>();
             services.AddScoped<IS3Service, S3Service>();
+            services.AddHttpClient("S3Upload", client =>
+            {
+                client.Timeout = Timeout.InfiniteTimeSpan;
+            });
             services.AddScoped<IURLAnalyzer, URLAnalyzer>();
 
             services.AddWampServiceAccessor(Servers.NotesnookAPI);
@@ -314,6 +342,7 @@ namespace Notesnook.API
 
             app.UseAuthentication();
             app.UseAuthorization();
+            app.UseRateLimiter();
 
             app.UseEndpoints(endpoints =>
             {
@@ -326,6 +355,13 @@ namespace Notesnook.API
                 });
             });
         }
+
+        private int ReadRateLimit(string name, int defaultValue)
+        {
+            return int.TryParse(Configuration[name], out var value) && value > 0
+                ? value
+                : defaultValue;
+        }
     }
 
     public static class ServiceCollectionMongoCollectionExtensions