api: add support for storage limits

Notesnook.API/Controllers/S3Controller.cs

@@ -24,6 +24,12 @@ using System.Threading.Tasks;
 using System.Security.Claims;
 using Notesnook.API.Interfaces;
 using System;
+using System.Net.Http;
+using Streetwriters.Common.Extensions;
+using Streetwriters.Common.Models;
+using Notesnook.API.Helpers;
+using Streetwriters.Common;
+using Streetwriters.Common.Interfaces;
 using Notesnook.API.Models;
 
 namespace Notesnook.API.Controllers
@@ -31,27 +37,55 @@ namespace Notesnook.API.Controllers
     [ApiController]
     [Route("s3")]
     [ResponseCache(NoStore = true, Location = ResponseCacheLocation.None)]
+    [Authorize("Sync")]
     public class S3Controller : ControllerBase
     {
+        private ISyncItemsRepositoryAccessor Repositories { get; }
         private IS3Service S3Service { get; set; }
-        public S3Controller(IS3Service s3Service)
+        public S3Controller(IS3Service s3Service, ISyncItemsRepositoryAccessor syncItemsRepositoryAccessor)
         {
             S3Service = s3Service;
+            Repositories = syncItemsRepositoryAccessor;
         }
 
         [HttpPut]
-        [Authorize("Pro")]
+        public async Task<IActionResult> Upload([FromQuery] string name)
-        public IActionResult Upload([FromQuery] string name)
         {
             var userId = this.User.FindFirstValue("sub");
+
+            if (!HttpContext.Request.Headers.ContentLength.HasValue) return BadRequest(new { error = "No Content-Length header found." });
+
+            long fileSize = HttpContext.Request.Headers.ContentLength.Value;
+            var subscriptionService = await WampServers.SubscriptionServer.GetServiceAsync<IUserSubscriptionService>(SubscriptionServerTopics.UserSubscriptionServiceTopic);
+            var subscription = await subscriptionService.GetUserSubscriptionAsync(Clients.Notesnook.Id, userId);
+
+            if (StorageHelper.IsFileSizeExceeded(subscription, fileSize))
+            {
+                return BadRequest(new { error = "Max file size exceeded." });
+            }
+
+            var userSettings = await Repositories.UsersSettings.FindOneAsync((u) => u.UserId == userId);
+            userSettings.StorageLimit ??= new Limit { Value = 0, UpdatedAt = 0 };
+            userSettings.StorageLimit.Value += fileSize;
+            if (StorageHelper.IsStorageLimitReached(subscription, userSettings.StorageLimit))
+                return BadRequest(new { error = "Storage limit exceeded." });
+
             var url = S3Service.GetUploadObjectUrl(userId, name);
-            if (url == null) return BadRequest("Could not create signed url.");
+            if (url == null) return BadRequest(new { error = "Could not create signed url." });
-            return Ok(url);
+
+            var httpClient = new HttpClient();
+            var content = new StreamContent(HttpContext.Request.BodyReader.AsStream());
+            var response = await httpClient.SendRequestAsync<Response>(url, null, HttpMethod.Put, content);
+            if (!response.Success) return BadRequest(await response.Content.ReadAsStringAsync());
+
+            userSettings.StorageLimit.UpdatedAt = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
+            await Repositories.UsersSettings.UpsertAsync(userSettings, (u) => u.UserId == userId);
+
+            return Ok(response);
         }
 
 
         [HttpGet("multipart")]
-        [Authorize("Pro")]
         public async Task<IActionResult> MultipartUpload([FromQuery] string name, [FromQuery] int parts, [FromQuery] string uploadId)
         {
             var userId = this.User.FindFirstValue("sub");
@@ -64,7 +98,6 @@ namespace Notesnook.API.Controllers
         }
 
         [HttpDelete("multipart")]
-        [Authorize("Pro")]
         public async Task<IActionResult> AbortMultipartUpload([FromQuery] string name, [FromQuery] string uploadId)
         {
             var userId = this.User.FindFirstValue("sub");
@@ -77,7 +110,6 @@ namespace Notesnook.API.Controllers
         }
 
         [HttpPost("multipart")]
-        [Authorize("Pro")]
         public async Task<IActionResult> CompleteMultipartUpload([FromBody] CompleteMultipartUploadRequestWrapper uploadRequestWrapper)
         {
             var userId = this.User.FindFirstValue("sub");
@@ -90,7 +122,6 @@ namespace Notesnook.API.Controllers
         }
 
         [HttpGet]
-        [Authorize("Sync")]
         public IActionResult Download([FromQuery] string name)
         {
             var userId = this.User.FindFirstValue("sub");
@@ -100,7 +131,6 @@ namespace Notesnook.API.Controllers
         }
 
         [HttpHead]
-        [Authorize("Sync")]
         public async Task<IActionResult> Info([FromQuery] string name)
         {
             var userId = this.User.FindFirstValue("sub");
@@ -110,7 +140,6 @@ namespace Notesnook.API.Controllers
         }
 
         [HttpDelete]
-        [Authorize("Sync")]
         public async Task<IActionResult> DeleteAsync([FromQuery] string name)
         {
             try
@@ -125,4 +154,4 @@ namespace Notesnook.API.Controllers
             }
         }
     }
-}
+}

Notesnook.API/Helpers/StorageHelper.cs

@@ -0,0 +1,49 @@
+using System.Collections.Generic;
+using Notesnook.API.Models;
+using Streetwriters.Common.Enums;
+using Streetwriters.Common.Models;
+
+namespace Notesnook.API.Helpers
+{
+    class StorageHelper
+    {
+        const long MB = 1024 * 1024;
+        const long GB = 1024 * MB;
+        public readonly static Dictionary<SubscriptionPlan, long> MAX_STORAGE_PER_MONTH = new()
+            {
+                { SubscriptionPlan.FREE, 50L * MB },
+                { SubscriptionPlan.ESSENTIAL, GB },
+                { SubscriptionPlan.PRO, 10L * GB },
+                { SubscriptionPlan.EDUCATION, 10L * GB },
+                { SubscriptionPlan.BELIEVER, 25L * GB },
+                { SubscriptionPlan.LEGACY_PRO, -1 }
+            };
+        public readonly static Dictionary<SubscriptionPlan, long> MAX_FILE_SIZE = new()
+            {
+                { SubscriptionPlan.FREE, 10 * MB },
+                { SubscriptionPlan.ESSENTIAL, 100 * MB },
+                { SubscriptionPlan.PRO, 1L * GB },
+                { SubscriptionPlan.EDUCATION, 1L * GB },
+                { SubscriptionPlan.BELIEVER, 5L * GB },
+                { SubscriptionPlan.LEGACY_PRO, 512 * MB }
+            };
+
+        public static long GetStorageLimitForPlan(Subscription subscription)
+        {
+            return MAX_STORAGE_PER_MONTH[subscription.Plan];
+        }
+
+        public static bool IsStorageLimitReached(Subscription subscription, Limit limit)
+        {
+            var storageLimit = GetStorageLimitForPlan(subscription);
+            if (storageLimit == -1) return false;
+            return limit.Value > storageLimit;
+        }
+
+        public static bool IsFileSizeExceeded(Subscription subscription, long fileSize)
+        {
+            var maxFileSize = MAX_FILE_SIZE[subscription.Plan];
+            return fileSize > maxFileSize;
+        }
+    }
+}

Notesnook.API/Models/Responses/UserResponse.cs

@@ -1,3 +1,4 @@
+using System.Net.Http;
 using System.Text.Json.Serialization;
 using Streetwriters.Common.Interfaces;
 using Streetwriters.Common.Models;
@@ -7,25 +8,30 @@ namespace Notesnook.API.Models.Responses
     public class UserResponse : UserModel, IResponse
     {
         [JsonPropertyName("salt")]
-        public string Salt { get; set; }
+        public string? Salt { get; set; }
 
         [JsonPropertyName("attachmentsKey")]
-        public EncryptedData AttachmentsKey { get; set; }
+        public EncryptedData? AttachmentsKey { get; set; }
 
         [JsonPropertyName("monographPasswordsKey")]
-        public EncryptedData MonographPasswordsKey { get; set; }
+        public EncryptedData? MonographPasswordsKey { get; set; }
 
         [JsonPropertyName("inboxKeys")]
-        public InboxKeys InboxKeys { get; set; }
+        public InboxKeys? InboxKeys { get; set; }
 
         [JsonPropertyName("subscription")]
-        public ISubscription Subscription { get; set; }
+        public ISubscription? Subscription { get; set; }
 
-        [JsonPropertyName("profile")]
+        [JsonPropertyName("storageUsed")]
-        public EncryptedData Profile { get; set; }
+        public long StorageUsed { get; set; }
+
+        [JsonPropertyName("totalStorage")]
+        public long TotalStorage { get; set; }
 
         [JsonIgnore]
         public bool Success { get; set; }
         public int StatusCode { get; set; }
+        [JsonIgnore]
+        public HttpContent? Content { get; set; }
     }
 }

Notesnook.API/Models/UserSettings.cs

@@ -38,10 +38,10 @@ namespace Notesnook.API.Models
         public string UserId { get; set; }
         public long LastSynced { get; set; }
         public string Salt { get; set; }
-        public EncryptedData VaultKey { get; set; }
+        public EncryptedData? VaultKey { get; set; }
-        public EncryptedData AttachmentsKey { get; set; }
+        public EncryptedData? AttachmentsKey { get; set; }
-        public EncryptedData MonographPasswordsKey { get; set; }
+        public EncryptedData? MonographPasswordsKey { get; set; }
-        public InboxKeys InboxKeys { get; set; }
+        public InboxKeys? InboxKeys { get; set; }
         public Limit StorageLimit { get; set; }
 
         [BsonId]

Notesnook.API/Services/S3Service.cs

@@ -28,9 +28,13 @@ using Amazon.Runtime;
 using Amazon.S3;
 using Amazon.S3.Model;
 using Microsoft.Extensions.Options;
+using Notesnook.API.Helpers;
 using Notesnook.API.Interfaces;
 using Notesnook.API.Models;
 using Streetwriters.Common;
+using Streetwriters.Common.Enums;
+using Streetwriters.Common.Interfaces;
+using Streetwriters.Common.Models;
 
 namespace Notesnook.API.Services
 {
@@ -45,6 +49,7 @@ namespace Notesnook.API.Services
         private readonly string BUCKET_NAME = Constants.S3_BUCKET_NAME ?? "";
         private readonly string INTERNAL_BUCKET_NAME = Constants.S3_INTERNAL_BUCKET_NAME ?? "";
         private AmazonS3Client S3Client { get; }
+        private ISyncItemsRepositoryAccessor Repositories { get; }
 
         // When running in a dockerized environment the sync server doesn't have access
         // to the host's S3 Service URL. It can only talk to S3 server via its own internal
@@ -57,8 +62,9 @@ namespace Notesnook.API.Services
         private AmazonS3Client S3InternalClient { get; }
         private HttpClient httpClient = new HttpClient();
 
-        public S3Service()
+        public S3Service(ISyncItemsRepositoryAccessor syncItemsRepositoryAccessor)
         {
+            Repositories = syncItemsRepositoryAccessor;
             var config = new AmazonS3Config
             {
 #if (DEBUG || STAGING)
@@ -145,12 +151,6 @@ namespace Notesnook.API.Services
 
             var request = new HttpRequestMessage(HttpMethod.Head, url);
             var response = await httpClient.SendAsync(request);
-            const long MAX_SIZE = 513 * 1024 * 1024; // 512 MB
-            if (!Constants.IS_SELF_HOSTED && response.Content.Headers.ContentLength >= MAX_SIZE)
-            {
-                await this.DeleteObjectAsync(userId, name);
-                throw new Exception("File size exceeds the maximum allowed size.");
-            }
             return response.Content.Headers.ContentLength ?? 0;
         }
 
@@ -202,15 +202,55 @@ namespace Notesnook.API.Services
             if (!IsSuccessStatusCode(((int)response.HttpStatusCode))) throw new Exception("Failed to abort multipart upload.");
         }
 
+        private async Task<long> GetMultipartUploadSizeAsync(string userId, string key, string uploadId)
+        {
+            var objectName = GetFullObjectName(userId, key);
+            var parts = await GetS3Client(S3ClientMode.INTERNAL).ListPartsAsync(GetBucketName(S3ClientMode.INTERNAL), objectName, uploadId);
+            long totalSize = 0;
+            foreach (var part in parts.Parts)
+            {
+                totalSize += part.Size;
+            }
+            return totalSize;
+        }
+
         public async Task CompleteMultipartUploadAsync(string userId, CompleteMultipartUploadRequest uploadRequest)
         {
             var objectName = GetFullObjectName(userId, uploadRequest.Key);
             if (userId == null || objectName == null) throw new Exception("Could not abort multipart upload.");
 
+            var subscriptionService = await WampServers.SubscriptionServer.GetServiceAsync<IUserSubscriptionService>(SubscriptionServerTopics.UserSubscriptionServiceTopic);
+            var subscription = await subscriptionService.GetUserSubscriptionAsync(Clients.Notesnook.Id, userId);
+
+            long 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.");
+            }
+
+            var userSettings = await Repositories.UsersSettings.FindOneAsync((u) => u.UserId == userId);
+            if (userSettings == null)
+            {
+                await this.AbortMultipartUploadAsync(userId, uploadRequest.Key, uploadRequest.UploadId);
+                throw new Exception("User settings not found.");
+            }
+
+            userSettings.StorageLimit ??= new Limit { Value = 0, UpdatedAt = 0 };
+            userSettings.StorageLimit.Value += fileSize;
+            if (StorageHelper.IsStorageLimitReached(subscription, userSettings.StorageLimit))
+            {
+                await this.AbortMultipartUploadAsync(userId, uploadRequest.Key, uploadRequest.UploadId);
+                throw new Exception("Storage limit reached.");
+            }
+
             uploadRequest.Key = objectName;
             uploadRequest.BucketName = GetBucketName(S3ClientMode.INTERNAL);
             var response = await GetS3Client(S3ClientMode.INTERNAL).CompleteMultipartUploadAsync(uploadRequest);
             if (!IsSuccessStatusCode(((int)response.HttpStatusCode))) throw new Exception("Failed to complete multipart upload.");
+
+            userSettings.StorageLimit.UpdatedAt = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
+            await Repositories.UsersSettings.UpsertAsync(userSettings, (u) => u.UserId == userId);
         }
 
         private string? GetPresignedURL(string userId, string name, HttpVerb httpVerb, S3ClientMode mode = S3ClientMode.EXTERNAL)

Notesnook.API/Services/UserService.cs

@@ -23,6 +23,7 @@ using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
+using Notesnook.API.Helpers;
 using Notesnook.API.Interfaces;
 using Notesnook.API.Models;
 using Notesnook.API.Models.Responses;
@@ -71,6 +72,7 @@ namespace Notesnook.API.Services
             await Repositories.UsersSettings.InsertAsync(new UserSettings
             {
                 UserId = response.UserId,
+                StorageLimit = new Limit { UpdatedAt = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(), Value = 0 },
                 LastSynced = 0,
                 Salt = GetSalt()
             });
@@ -119,6 +121,14 @@ namespace Notesnook.API.Services
             }
 
             var userSettings = await Repositories.UsersSettings.FindOneAsync((u) => u.UserId == user.UserId) ?? throw new Exception("User settings not found.");
+
+            // reset user's attachment limit every month
+            if (userSettings.StorageLimit == null || DateTimeOffset.UtcNow.Month > DateTimeOffset.FromUnixTimeMilliseconds(userSettings.StorageLimit.UpdatedAt).Month)
+            {
+                userSettings.StorageLimit ??= new Limit { UpdatedAt = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(), Value = 0 };
+                await Repositories.UsersSettings.UpsertAsync(userSettings, (u) => u.UserId == user.UserId);
+            }
+
             return new UserResponse
             {
                 UserId = user.UserId,
@@ -132,6 +142,8 @@ namespace Notesnook.API.Services
                 InboxKeys = userSettings.InboxKeys,
                 Salt = userSettings.Salt,
                 Subscription = subscription,
+                StorageUsed = userSettings.StorageLimit.Value,
+                TotalStorage = StorageHelper.GetStorageLimitForPlan(subscription),
                 Success = true,
                 StatusCode = 200
             };
@@ -262,6 +274,7 @@ namespace Notesnook.API.Services
             userSettings.AttachmentsKey = null;
             userSettings.MonographPasswordsKey = null;
             userSettings.VaultKey = null;
+            userSettings.InboxKeys = null;
             userSettings.LastSynced = 0;
 
             await Repositories.UsersSettings.UpsertAsync(userSettings, (s) => s.UserId == userId);