[PM-25088] - refactor premium purchase endpoint (#6262)

* [PM-25088] add feature flag for new premium subscription flow

* [PM-25088] refactor premium endpoint

* forgot the punctuation change in the test

* [PM-25088] - pr feedback

* [PM-25088] - pr feedback round two

src/Core/AdminConsole/Services/OrganizationFactory.cs

@@ -23,7 +23,7 @@ public static class OrganizationFactory
             PlanType = claimsPrincipal.GetValue<PlanType>(OrganizationLicenseConstants.PlanType),
             Seats = claimsPrincipal.GetValue<int?>(OrganizationLicenseConstants.Seats),
             MaxCollections = claimsPrincipal.GetValue<short?>(OrganizationLicenseConstants.MaxCollections),
-            MaxStorageGb = 10240,
+            MaxStorageGb = Constants.SelfHostedMaxStorageGb,
             UsePolicies = claimsPrincipal.GetValue<bool>(OrganizationLicenseConstants.UsePolicies),
             UseSso = claimsPrincipal.GetValue<bool>(OrganizationLicenseConstants.UseSso),
             UseKeyConnector = claimsPrincipal.GetValue<bool>(OrganizationLicenseConstants.UseKeyConnector),
@@ -75,7 +75,7 @@ public static class OrganizationFactory
             PlanType = license.PlanType,
             Seats = license.Seats,
             MaxCollections = license.MaxCollections,
-            MaxStorageGb = 10240,
+            MaxStorageGb = Constants.SelfHostedMaxStorageGb,
             UsePolicies = license.UsePolicies,
             UseSso = license.UseSso,
             UseKeyConnector = license.UseKeyConnector,

src/Core/Billing/Constants/StripeConstants.cs

@@ -79,6 +79,7 @@ public static class StripeConstants
     public static class Prices
     {
         public const string StoragePlanPersonal = "personal-storage-gb-annually";
+        public const string PremiumAnnually = "premium-annually";
     }
 
     public static class ProrationBehavior

src/Core/Billing/Extensions/ServiceCollectionExtensions.cs

@@ -5,6 +5,7 @@ using Bit.Core.Billing.Organizations.Commands;
 using Bit.Core.Billing.Organizations.Queries;
 using Bit.Core.Billing.Organizations.Services;
 using Bit.Core.Billing.Payment;
+using Bit.Core.Billing.Premium.Commands;
 using Bit.Core.Billing.Pricing;
 using Bit.Core.Billing.Services;
 using Bit.Core.Billing.Services.Implementations;
@@ -30,6 +31,7 @@ public static class ServiceCollectionExtensions
         services.AddTransient<IPreviewTaxAmountCommand, PreviewTaxAmountCommand>();
         services.AddPaymentOperations();
         services.AddOrganizationLicenseCommandsQueries();
+        services.AddPremiumCommands();
         services.AddTransient<IGetOrganizationWarningsQuery, GetOrganizationWarningsQuery>();
     }
 
@@ -39,4 +41,10 @@ public static class ServiceCollectionExtensions
         services.AddScoped<IGetSelfHostedOrganizationLicenseQuery, GetSelfHostedOrganizationLicenseQuery>();
         services.AddScoped<IUpdateOrganizationLicenseCommand, UpdateOrganizationLicenseCommand>();
     }
+
+    private static void AddPremiumCommands(this IServiceCollection services)
+    {
+        services.AddScoped<ICreatePremiumCloudHostedSubscriptionCommand, CreatePremiumCloudHostedSubscriptionCommand>();
+        services.AddScoped<ICreatePremiumSelfHostedSubscriptionCommand, CreatePremiumSelfHostedSubscriptionCommand>();
+    }
 }

src/Core/Billing/Payment/Models/TokenizablePaymentMethodType.cs

@@ -6,3 +6,17 @@ public enum TokenizablePaymentMethodType
     Card,
     PayPal
 }
+
+public static class TokenizablePaymentMethodTypeExtensions
+{
+    public static TokenizablePaymentMethodType From(string type)
+    {
+        return type switch
+        {
+            "bankAccount" => TokenizablePaymentMethodType.BankAccount,
+            "card" => TokenizablePaymentMethodType.Card,
+            "payPal" => TokenizablePaymentMethodType.PayPal,
+            _ => throw new InvalidOperationException($"Invalid value for {nameof(TokenizedPaymentMethod)}.{nameof(TokenizedPaymentMethod.Type)}")
+        };
+    }
+}

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

@@ -0,0 +1,308 @@
+using Bit.Core.Billing.Caches;
+using Bit.Core.Billing.Commands;
+using Bit.Core.Billing.Constants;
+using Bit.Core.Billing.Payment.Models;
+using Bit.Core.Billing.Services;
+using Bit.Core.Entities;
+using Bit.Core.Enums;
+using Bit.Core.Platform.Push;
+using Bit.Core.Services;
+using Bit.Core.Settings;
+using Bit.Core.Utilities;
+using Braintree;
+using Microsoft.Extensions.Logging;
+using OneOf.Types;
+using Stripe;
+using Customer = Stripe.Customer;
+using Subscription = Stripe.Subscription;
+
+namespace Bit.Core.Billing.Premium.Commands;
+
+using static Utilities;
+
+/// <summary>
+/// Creates a premium subscription for a cloud-hosted user with Stripe payment processing.
+/// Handles customer creation, payment method setup, and subscription creation.
+/// </summary>
+public interface ICreatePremiumCloudHostedSubscriptionCommand
+{
+    /// <summary>
+    /// Creates a premium cloud-hosted subscription for the specified user.
+    /// </summary>
+    /// <param name="user">The user to create the premium subscription for. Must not already be a premium user.</param>
+    /// <param name="paymentMethod">The tokenized payment method containing the payment type and token for billing.</param>
+    /// <param name="billingAddress">The billing address information required for tax calculation and customer creation.</param>
+    /// <param name="additionalStorageGb">Additional storage in GB beyond the base 1GB included with premium (must be >= 0).</param>
+    /// <returns>A billing command result indicating success or failure with appropriate error details.</returns>
+    Task<BillingCommandResult<None>> Run(
+        User user,
+        TokenizedPaymentMethod paymentMethod,
+        BillingAddress billingAddress,
+        short additionalStorageGb);
+}
+
+public class CreatePremiumCloudHostedSubscriptionCommand(
+    IBraintreeGateway braintreeGateway,
+    IGlobalSettings globalSettings,
+    ISetupIntentCache setupIntentCache,
+    IStripeAdapter stripeAdapter,
+    ISubscriberService subscriberService,
+    IUserService userService,
+    IPushNotificationService pushNotificationService,
+    ILogger<CreatePremiumCloudHostedSubscriptionCommand> logger)
+    : BaseBillingCommand<CreatePremiumCloudHostedSubscriptionCommand>(logger), ICreatePremiumCloudHostedSubscriptionCommand
+{
+    private static readonly List<string> _expand = ["tax"];
+    private readonly ILogger<CreatePremiumCloudHostedSubscriptionCommand> _logger = logger;
+
+    public Task<BillingCommandResult<None>> Run(
+        User user,
+        TokenizedPaymentMethod paymentMethod,
+        BillingAddress billingAddress,
+        short additionalStorageGb) => HandleAsync<None>(async () =>
+    {
+        if (user.Premium)
+        {
+            return new BadRequest("Already a premium user.");
+        }
+
+        if (additionalStorageGb < 0)
+        {
+            return new BadRequest("Additional storage must be greater than 0.");
+        }
+
+        var customer = string.IsNullOrEmpty(user.GatewayCustomerId)
+            ? await CreateCustomerAsync(user, paymentMethod, billingAddress)
+            : await subscriberService.GetCustomerOrThrow(user, new CustomerGetOptions { Expand = _expand });
+
+        customer = await ReconcileBillingLocationAsync(customer, billingAddress);
+
+        var subscription = await CreateSubscriptionAsync(user.Id, customer, additionalStorageGb > 0 ? additionalStorageGb : null);
+
+        switch (paymentMethod)
+        {
+            case { Type: TokenizablePaymentMethodType.PayPal }
+                when subscription.Status == StripeConstants.SubscriptionStatus.Incomplete:
+            case { Type: not TokenizablePaymentMethodType.PayPal }
+                when subscription.Status == StripeConstants.SubscriptionStatus.Active:
+                {
+                    user.Premium = true;
+                    user.PremiumExpirationDate = subscription.CurrentPeriodEnd;
+                    break;
+                }
+        }
+
+        user.Gateway = GatewayType.Stripe;
+        user.GatewayCustomerId = customer.Id;
+        user.GatewaySubscriptionId = subscription.Id;
+        user.MaxStorageGb = (short)(1 + additionalStorageGb);
+        user.LicenseKey = CoreHelpers.SecureRandomString(20);
+        user.RevisionDate = DateTime.UtcNow;
+
+        await userService.SaveUserAsync(user);
+        await pushNotificationService.PushSyncVaultAsync(user.Id);
+
+        return new None();
+    });
+
+    private async Task<Customer> CreateCustomerAsync(User user,
+        TokenizedPaymentMethod paymentMethod,
+        BillingAddress billingAddress)
+    {
+        var subscriberName = user.SubscriberName();
+        var customerCreateOptions = new CustomerCreateOptions
+        {
+            Address = new AddressOptions
+            {
+                Line1 = billingAddress.Line1,
+                Line2 = billingAddress.Line2,
+                City = billingAddress.City,
+                PostalCode = billingAddress.PostalCode,
+                State = billingAddress.State,
+                Country = billingAddress.Country
+            },
+            Description = user.Name,
+            Email = user.Email,
+            Expand = _expand,
+            InvoiceSettings = new CustomerInvoiceSettingsOptions
+            {
+                CustomFields =
+                [
+                    new CustomerInvoiceSettingsCustomFieldOptions
+                    {
+                        Name = user.SubscriberType(),
+                        Value = subscriberName.Length <= 30
+                            ? subscriberName
+                            : subscriberName[..30]
+                    }
+                ]
+            },
+            Metadata = new Dictionary<string, string>
+            {
+                [StripeConstants.MetadataKeys.Region] = globalSettings.BaseServiceUri.CloudRegion,
+                [StripeConstants.MetadataKeys.UserId] = user.Id.ToString()
+            },
+            Tax = new CustomerTaxOptions
+            {
+                ValidateLocation = StripeConstants.ValidateTaxLocationTiming.Immediately
+            }
+        };
+
+        var braintreeCustomerId = "";
+
+        // ReSharper disable once SwitchStatementHandlesSomeKnownEnumValuesWithDefault
+        switch (paymentMethod.Type)
+        {
+            case TokenizablePaymentMethodType.BankAccount:
+                {
+                    var setupIntent =
+                        (await stripeAdapter.SetupIntentList(new SetupIntentListOptions { PaymentMethod = paymentMethod.Token }))
+                        .FirstOrDefault();
+
+                    if (setupIntent == null)
+                    {
+                        _logger.LogError("Cannot create customer for user ({UserID}) without a setup intent for their bank account", user.Id);
+                        throw new BillingException();
+                    }
+
+                    await setupIntentCache.Set(user.Id, setupIntent.Id);
+                    break;
+                }
+            case TokenizablePaymentMethodType.Card:
+                {
+                    customerCreateOptions.PaymentMethod = paymentMethod.Token;
+                    customerCreateOptions.InvoiceSettings.DefaultPaymentMethod = paymentMethod.Token;
+                    break;
+                }
+            case TokenizablePaymentMethodType.PayPal:
+                {
+                    braintreeCustomerId = await subscriberService.CreateBraintreeCustomer(user, paymentMethod.Token);
+                    customerCreateOptions.Metadata[BraintreeCustomerIdKey] = braintreeCustomerId;
+                    break;
+                }
+            default:
+                {
+                    _logger.LogError("Cannot create customer for user ({UserID}) using payment method type ({PaymentMethodType}) as it is not supported", user.Id, paymentMethod.Type.ToString());
+                    throw new BillingException();
+                }
+        }
+
+        try
+        {
+            return await stripeAdapter.CustomerCreateAsync(customerCreateOptions);
+        }
+        catch
+        {
+            await Revert();
+            throw;
+        }
+
+        async Task Revert()
+        {
+            // ReSharper disable once SwitchStatementMissingSomeEnumCasesNoDefault
+            switch (paymentMethod.Type)
+            {
+                case TokenizablePaymentMethodType.BankAccount:
+                    {
+                        await setupIntentCache.Remove(user.Id);
+                        break;
+                    }
+                case TokenizablePaymentMethodType.PayPal when !string.IsNullOrEmpty(braintreeCustomerId):
+                    {
+                        await braintreeGateway.Customer.DeleteAsync(braintreeCustomerId);
+                        break;
+                    }
+            }
+        }
+    }
+
+    private async Task<Customer> ReconcileBillingLocationAsync(
+        Customer customer,
+        BillingAddress billingAddress)
+    {
+        /*
+         * If the customer was previously set up with credit, which does not require a billing location,
+         * we need to update the customer on the fly before we start the subscription.
+         */
+        if (customer is { Address: { Country: not null and not "", PostalCode: not null and not "" } })
+        {
+            return customer;
+        }
+
+        var options = new CustomerUpdateOptions
+        {
+            Address = new AddressOptions
+            {
+                Line1 = billingAddress.Line1,
+                Line2 = billingAddress.Line2,
+                City = billingAddress.City,
+                PostalCode = billingAddress.PostalCode,
+                State = billingAddress.State,
+                Country = billingAddress.Country
+            },
+            Expand = _expand,
+            Tax = new CustomerTaxOptions
+            {
+                ValidateLocation = StripeConstants.ValidateTaxLocationTiming.Immediately
+            }
+        };
+        return await stripeAdapter.CustomerUpdateAsync(customer.Id, options);
+    }
+
+    private async Task<Subscription> CreateSubscriptionAsync(
+        Guid userId,
+        Customer customer,
+        int? storage)
+    {
+        var subscriptionItemOptionsList = new List<SubscriptionItemOptions>
+        {
+            new ()
+            {
+                Price = StripeConstants.Prices.PremiumAnnually,
+                Quantity = 1
+            }
+        };
+
+        if (storage is > 0)
+        {
+            subscriptionItemOptionsList.Add(new SubscriptionItemOptions
+            {
+                Price = StripeConstants.Prices.StoragePlanPersonal,
+                Quantity = storage
+            });
+        }
+
+        var usingPayPal = customer.Metadata?.ContainsKey(BraintreeCustomerIdKey) ?? false;
+
+        var subscriptionCreateOptions = new SubscriptionCreateOptions
+        {
+            AutomaticTax = new SubscriptionAutomaticTaxOptions
+            {
+                Enabled = true
+            },
+            CollectionMethod = StripeConstants.CollectionMethod.ChargeAutomatically,
+            Customer = customer.Id,
+            Items = subscriptionItemOptionsList,
+            Metadata = new Dictionary<string, string>
+            {
+                [StripeConstants.MetadataKeys.UserId] = userId.ToString()
+            },
+            PaymentBehavior = usingPayPal
+                ? StripeConstants.PaymentBehavior.DefaultIncomplete
+                : null,
+            OffSession = true
+        };
+
+        var subscription = await stripeAdapter.SubscriptionCreateAsync(subscriptionCreateOptions);
+
+        if (usingPayPal)
+        {
+            await stripeAdapter.InvoiceUpdateAsync(subscription.LatestInvoiceId, new InvoiceUpdateOptions
+            {
+                AutoAdvance = false
+            });
+        }
+
+        return subscription;
+    }
+}

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

@@ -0,0 +1,67 @@
+using Bit.Core.Billing.Commands;
+using Bit.Core.Billing.Models.Business;
+using Bit.Core.Billing.Services;
+using Bit.Core.Entities;
+using Bit.Core.Platform.Push;
+using Bit.Core.Services;
+using Microsoft.Extensions.Logging;
+using OneOf.Types;
+
+namespace Bit.Core.Billing.Premium.Commands;
+
+/// <summary>
+/// Creates a premium subscription for a self-hosted user.
+/// Validates the license and applies premium benefits including storage limits based on the license terms.
+/// </summary>
+public interface ICreatePremiumSelfHostedSubscriptionCommand
+{
+    /// <summary>
+    /// Creates a premium self-hosted subscription for the specified user using the provided license.
+    /// </summary>
+    /// <param name="user">The user to create the premium subscription for. Must not already be a premium user.</param>
+    /// <param name="license">The user license containing the premium subscription details and verification data. Must be valid and usable by the specified user.</param>
+    /// <returns>A billing command result indicating success or failure with appropriate error details.</returns>
+    Task<BillingCommandResult<None>> Run(User user, UserLicense license);
+}
+
+public class CreatePremiumSelfHostedSubscriptionCommand(
+    ILicensingService licensingService,
+    IUserService userService,
+    IPushNotificationService pushNotificationService,
+    ILogger<CreatePremiumSelfHostedSubscriptionCommand> logger)
+    : BaseBillingCommand<CreatePremiumSelfHostedSubscriptionCommand>(logger), ICreatePremiumSelfHostedSubscriptionCommand
+{
+    public Task<BillingCommandResult<None>> Run(
+        User user,
+        UserLicense license) => HandleAsync<None>(async () =>
+    {
+        if (user.Premium)
+        {
+            return new BadRequest("Already a premium user.");
+        }
+
+        if (!licensingService.VerifyLicense(license))
+        {
+            return new BadRequest("Invalid license.");
+        }
+
+        var claimsPrincipal = licensingService.GetClaimsPrincipalFromLicense(license);
+        if (!license.CanUse(user, claimsPrincipal, out var exceptionMessage))
+        {
+            return new BadRequest(exceptionMessage);
+        }
+
+        await licensingService.WriteUserLicenseAsync(user, license);
+
+        user.Premium = true;
+        user.RevisionDate = DateTime.UtcNow;
+        user.MaxStorageGb = Core.Constants.SelfHostedMaxStorageGb;
+        user.LicenseKey = license.LicenseKey;
+        user.PremiumExpirationDate = license.Expires;
+
+        await userService.SaveUserAsync(user);
+        await pushNotificationService.PushSyncVaultAsync(user.Id);
+
+        return new None();
+    });
+}

src/Core/Billing/Services/ILicensingService.cs

@@ -26,4 +26,5 @@ public interface ILicensingService
         SubscriptionInfo subscriptionInfo);
 
     Task<string?> CreateUserTokenAsync(User user, SubscriptionInfo subscriptionInfo);
+    Task WriteUserLicenseAsync(User user, UserLicense license);
 }

src/Core/Billing/Services/Implementations/LicensingService.cs

@@ -389,4 +389,12 @@ public class LicensingService : ILicensingService
         var token = tokenHandler.CreateToken(tokenDescriptor);
         return tokenHandler.WriteToken(token);
     }
+
+    public async Task WriteUserLicenseAsync(User user, UserLicense license)
+    {
+        var dir = $"{_globalSettings.LicenseDirectory}/user";
+        Directory.CreateDirectory(dir);
+        await using var fs = File.OpenWrite(Path.Combine(dir, $"{user.Id}.json"));
+        await JsonSerializer.SerializeAsync(fs, license, JsonHelpers.Indented);
+    }
 }

src/Core/Billing/Services/Implementations/PremiumUserBillingService.cs

@@ -304,7 +304,7 @@ public class PremiumUserBillingService(
         {
             new ()
             {
-                Price = "premium-annually",
+                Price = StripeConstants.Prices.PremiumAnnually,
                 Quantity = 1
             }
         };

src/Core/Billing/Services/NoopImplementations/NoopLicensingService.cs

@@ -73,4 +73,9 @@ public class NoopLicensingService : ILicensingService
     {
         return Task.FromResult<string?>(null);
     }
+
+    public Task WriteUserLicenseAsync(User user, UserLicense license)
+    {
+        return Task.CompletedTask;
+    }
 }

src/Core/Constants.cs

@@ -10,6 +10,11 @@ public static class Constants
     public const int BypassFiltersEventId = 12482444;
     public const int FailedSecretVerificationDelay = 2000;
 
+    /// <summary>
+    /// Self-hosted max storage limit in GB (10 TB).
+    /// </summary>
+    public const short SelfHostedMaxStorageGb = 10240;
+
     // File size limits - give 1 MB extra for cushion.
     // Note: if request size limits are changed, 'client_max_body_size'
     // in nginx/proxy.conf may also need to be updated accordingly.
@@ -166,6 +171,7 @@ public static class FeatureFlagKeys
     public const string PM21881_ManagePaymentDetailsOutsideCheckout = "pm-21881-manage-payment-details-outside-checkout";
     public const string PM21821_ProviderPortalTakeover = "pm-21821-provider-portal-takeover";
     public const string PM22415_TaxIDWarnings = "pm-22415-tax-id-warnings";
+    public const string PM23385_UseNewPremiumFlow = "pm-23385-use-new-premium-flow";
 
     /* Key Management Team */
     public const string ReturnErrorOnExistingKeypair = "return-error-on-existing-keypair";

src/Core/Services/Implementations/StripePaymentService.cs

@@ -906,7 +906,7 @@ public class StripePaymentService : IPaymentService
                     new()
                     {
                         Quantity = 1,
-                        Plan = "premium-annually"
+                        Plan = StripeConstants.Prices.PremiumAnnually
                     },
 
                     new()

src/Core/Services/Implementations/UserService.cs

@@ -44,8 +44,6 @@ namespace Bit.Core.Services;
 
 public class UserService : UserManager<User>, IUserService
 {
-    private const string PremiumPlanId = "premium-annually";
-
     private readonly IUserRepository _userRepository;
     private readonly IOrganizationUserRepository _organizationUserRepository;
     private readonly IOrganizationRepository _organizationRepository;
@@ -930,7 +928,7 @@ public class UserService : UserManager<User>, IUserService
 
         if (_globalSettings.SelfHosted)
         {
-            user.MaxStorageGb = 10240; // 10 TB
+            user.MaxStorageGb = Constants.SelfHostedMaxStorageGb;
             user.LicenseKey = license.LicenseKey;
             user.PremiumExpirationDate = license.Expires;
         }
@@ -989,7 +987,7 @@ public class UserService : UserManager<User>, IUserService
 
         user.Premium = license.Premium;
         user.RevisionDate = DateTime.UtcNow;
-        user.MaxStorageGb = _globalSettings.SelfHosted ? 10240 : license.MaxStorageGb; // 10 TB
+        user.MaxStorageGb = _globalSettings.SelfHosted ? Constants.SelfHostedMaxStorageGb : license.MaxStorageGb;
         user.LicenseKey = license.LicenseKey;
         user.PremiumExpirationDate = license.Expires;
         await SaveUserAsync(user);

src/Core/Tools/SendFeatures/Services/SendValidationService.cs

@@ -125,7 +125,7 @@ public class SendValidationService : ISendValidationService
             {
                 // Users that get access to file storage/premium from their organization get the default
                 // 1 GB max storage.
-                short limit = _globalSettings.SelfHosted ? (short)10240 : (short)1;
+                short limit = _globalSettings.SelfHosted ? Constants.SelfHostedMaxStorageGb : (short)1;
                 storageBytesRemaining = user.StorageBytesRemaining(limit);
             }
         }

src/Core/Vault/Services/Implementations/CipherService.cs

@@ -933,7 +933,7 @@ public class CipherService : ICipherService
                 // Users that get access to file storage/premium from their organization get the default
                 // 1 GB max storage.
                 storageBytesRemaining = user.StorageBytesRemaining(
-                    _globalSettings.SelfHosted ? (short)10240 : (short)1);
+                    _globalSettings.SelfHosted ? Constants.SelfHostedMaxStorageGb : (short)1);
             }
         }
         else if (cipher.OrganizationId.HasValue)