[PM 29610]Update Account Storage Endpoint (#6750)

* update account storage endpoint

* Fix the failing test

* Added flag and refactor base on pr comments

* fix the lint error

* Resolve the pr comments

* Fix the failing test

* Fix the failing test

* Return none

* Resolve the lint error

* Fix the failing test

* Add the missing test

* Formatting issues fixed

src/Api/Billing/Controllers/VNext/AccountBillingVNextController.cs

@@ -1,6 +1,7 @@
 using Bit.Api.Billing.Attributes;
 using Bit.Api.Billing.Models.Requests.Payment;
 using Bit.Api.Billing.Models.Requests.Premium;
+using Bit.Api.Billing.Models.Requests.Storage;
 using Bit.Core;
 using Bit.Core.Billing.Licenses.Queries;
 using Bit.Core.Billing.Payment.Commands;
@@ -23,7 +24,8 @@ public class AccountBillingVNextController(
     IGetCreditQuery getCreditQuery,
     IGetPaymentMethodQuery getPaymentMethodQuery,
     IGetUserLicenseQuery getUserLicenseQuery,
-    IUpdatePaymentMethodCommand updatePaymentMethodCommand) : BaseBillingController
+    IUpdatePaymentMethodCommand updatePaymentMethodCommand,
+    IUpdatePremiumStorageCommand updatePremiumStorageCommand) : BaseBillingController
 {
     [HttpGet("credit")]
     [InjectUser]
@@ -88,4 +90,15 @@ public class AccountBillingVNextController(
         var response = await getUserLicenseQuery.Run(user);
         return TypedResults.Ok(response);
     }
+
+    [HttpPut("storage")]
+    [RequireFeature(FeatureFlagKeys.PM29594_UpdateIndividualSubscriptionPage)]
+    [InjectUser]
+    public async Task<IResult> UpdateStorageAsync(
+        [BindNever] User user,
+        [FromBody] StorageUpdateRequest request)
+    {
+        var result = await updatePremiumStorageCommand.Run(user, request.AdditionalStorageGb);
+        return Handle(result);
+    }
 }

src/Api/Billing/Models/Requests/Storage/StorageUpdateRequest.cs

@@ -0,0 +1,35 @@
+using System.ComponentModel.DataAnnotations;
+
+namespace Bit.Api.Billing.Models.Requests.Storage;
+
+/// <summary>
+/// Request model for updating storage allocation on a user's premium subscription.
+/// Allows for both increasing and decreasing storage in an idempotent manner.
+/// </summary>
+public class StorageUpdateRequest : IValidatableObject
+{
+    /// <summary>
+    /// The additional storage in GB beyond the base storage.
+    /// Must be between 0 and the maximum allowed (minus base storage).
+    /// </summary>
+    [Required]
+    [Range(0, 99)]
+    public short AdditionalStorageGb { get; set; }
+
+    public IEnumerable<ValidationResult> Validate(ValidationContext validationContext)
+    {
+        if (AdditionalStorageGb < 0)
+        {
+            yield return new ValidationResult(
+                "Additional storage cannot be negative.",
+                new[] { nameof(AdditionalStorageGb) });
+        }
+
+        if (AdditionalStorageGb > 99)
+        {
+            yield return new ValidationResult(
+                "Maximum additional storage is 99 GB.",
+                new[] { nameof(AdditionalStorageGb) });
+        }
+    }
+}

src/Core/Billing/Extensions/ServiceCollectionExtensions.cs

@@ -53,6 +53,7 @@ public static class ServiceCollectionExtensions
         services.AddScoped<ICreatePremiumCloudHostedSubscriptionCommand, CreatePremiumCloudHostedSubscriptionCommand>();
         services.AddScoped<ICreatePremiumSelfHostedSubscriptionCommand, CreatePremiumSelfHostedSubscriptionCommand>();
         services.AddTransient<IPreviewPremiumTaxCommand, PreviewPremiumTaxCommand>();
+        services.AddScoped<IUpdatePremiumStorageCommand, UpdatePremiumStorageCommand>();
     }
 
     private static void AddPremiumQueries(this IServiceCollection services)

src/Core/Billing/Premium/Commands/UpdatePremiumStorageCommand.cs

@@ -0,0 +1,144 @@
+using Bit.Core.Billing.Commands;
+using Bit.Core.Billing.Pricing;
+using Bit.Core.Billing.Services;
+using Bit.Core.Entities;
+using Bit.Core.Services;
+using Bit.Core.Utilities;
+using Microsoft.Extensions.Logging;
+using OneOf.Types;
+using Stripe;
+
+namespace Bit.Core.Billing.Premium.Commands;
+
+/// <summary>
+/// Updates the storage allocation for a premium user's subscription.
+/// Handles both increases and decreases in storage in an idempotent manner.
+/// </summary>
+public interface IUpdatePremiumStorageCommand
+{
+    /// <summary>
+    /// Updates the user's storage by the specified additional amount.
+    /// </summary>
+    /// <param name="user">The premium user whose storage should be updated.</param>
+    /// <param name="additionalStorageGb">The additional storage amount in GB beyond base storage.</param>
+    /// <returns>A billing command result indicating success or failure.</returns>
+    Task<BillingCommandResult<None>> Run(User user, short additionalStorageGb);
+}
+
+public class UpdatePremiumStorageCommand(
+    IStripeAdapter stripeAdapter,
+    IUserService userService,
+    IPricingClient pricingClient,
+    ILogger<UpdatePremiumStorageCommand> logger)
+    : BaseBillingCommand<UpdatePremiumStorageCommand>(logger), IUpdatePremiumStorageCommand
+{
+    public Task<BillingCommandResult<None>> Run(User user, short additionalStorageGb) => HandleAsync<None>(async () =>
+    {
+        if (!user.Premium)
+        {
+            return new BadRequest("User does not have a premium subscription.");
+        }
+
+        if (!user.MaxStorageGb.HasValue)
+        {
+            return new BadRequest("No access to storage.");
+        }
+
+        // Fetch all premium plans and the user's subscription to find which plan they're on
+        var premiumPlans = await pricingClient.ListPremiumPlans();
+        var subscription = await stripeAdapter.GetSubscriptionAsync(user.GatewaySubscriptionId);
+
+        // Find the password manager subscription item (seat, not storage) and match it to a plan
+        var passwordManagerItem = subscription.Items.Data.FirstOrDefault(i =>
+            premiumPlans.Any(p => p.Seat.StripePriceId == i.Price.Id));
+
+        if (passwordManagerItem == null)
+        {
+            return new BadRequest("Premium subscription item not found.");
+        }
+
+        var premiumPlan = premiumPlans.First(p => p.Seat.StripePriceId == passwordManagerItem.Price.Id);
+
+        var baseStorageGb = (short)premiumPlan.Storage.Provided;
+
+        if (additionalStorageGb < 0)
+        {
+            return new BadRequest("Additional storage cannot be negative.");
+        }
+
+        var newTotalStorageGb = (short)(baseStorageGb + additionalStorageGb);
+
+        if (newTotalStorageGb > 100)
+        {
+            return new BadRequest("Maximum storage is 100 GB.");
+        }
+
+        // Idempotency check: if user already has the requested storage, return success
+        if (user.MaxStorageGb == newTotalStorageGb)
+        {
+            return new None();
+        }
+
+        var remainingStorage = user.StorageBytesRemaining(newTotalStorageGb);
+        if (remainingStorage < 0)
+        {
+            return new BadRequest(
+                $"You are currently using {CoreHelpers.ReadableBytesSize(user.Storage.GetValueOrDefault(0))} of storage. " +
+                "Delete some stored data first.");
+        }
+
+        // Find the storage line item in the subscription
+        var storageItem = subscription.Items.Data.FirstOrDefault(i => i.Price.Id == premiumPlan.Storage.StripePriceId);
+
+        var subscriptionItemOptions = new List<SubscriptionItemOptions>();
+
+        if (additionalStorageGb > 0)
+        {
+            if (storageItem != null)
+            {
+                // Update existing storage item
+                subscriptionItemOptions.Add(new SubscriptionItemOptions
+                {
+                    Id = storageItem.Id,
+                    Price = premiumPlan.Storage.StripePriceId,
+                    Quantity = additionalStorageGb
+                });
+            }
+            else
+            {
+                // Add new storage item
+                subscriptionItemOptions.Add(new SubscriptionItemOptions
+                {
+                    Price = premiumPlan.Storage.StripePriceId,
+                    Quantity = additionalStorageGb
+                });
+            }
+        }
+        else if (storageItem != null)
+        {
+            // Remove storage item if setting to 0
+            subscriptionItemOptions.Add(new SubscriptionItemOptions
+            {
+                Id = storageItem.Id,
+                Deleted = true
+            });
+        }
+
+        // Update subscription with prorations
+        // Storage is billed annually, so we create prorations and invoice immediately
+        var subscriptionUpdateOptions = new SubscriptionUpdateOptions
+        {
+            Items = subscriptionItemOptions,
+            ProrationBehavior = Core.Constants.CreateProrations
+        };
+
+        await stripeAdapter.UpdateSubscriptionAsync(subscription.Id, subscriptionUpdateOptions);
+
+        // Update the user's max storage
+        user.MaxStorageGb = newTotalStorageGb;
+        await userService.SaveUserAsync(user);
+
+        // No payment intent needed - the subscription update will automatically create and finalize the invoice
+        return new None();
+    });
+}

src/Core/Constants.cs

@@ -196,6 +196,7 @@ public static class FeatureFlagKeys
     public const string PM26462_Milestone_3 = "pm-26462-milestone-3";
     public const string PM28265_EnableReconcileAdditionalStorageJob = "pm-28265-enable-reconcile-additional-storage-job";
     public const string PM28265_ReconcileAdditionalStorageJobEnableLiveMode = "pm-28265-reconcile-additional-storage-job-enable-live-mode";
+    public const string PM29594_UpdateIndividualSubscriptionPage = "pm-29594-update-individual-subscription-page";
 
     /* Key Management Team */
     public const string PrivateKeyRegeneration = "pm-12241-private-key-regeneration";