From 443a147433d30c45ddff1f8b80e626e890253197 Mon Sep 17 00:00:00 2001
From: Alex Morask <144709477+amorask-bitwarden@users.noreply.github.com>
Date: Thu, 30 Jan 2025 11:55:05 -0500
Subject: [PATCH] Replace StripePaymentService with PremiumUserBillingService
 in ReplacePaymentMethodAsync call (#5350)

---
 .../Services/IPremiumUserBillingService.cs    |  8 ++++++-
 .../PremiumUserBillingService.cs              | 23 +++++++++++++++++++
 .../Services/Implementations/UserService.cs   | 11 +++++----
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/Core/Billing/Services/IPremiumUserBillingService.cs b/src/Core/Billing/Services/IPremiumUserBillingService.cs
index f74bf6c8da..2161b247b9 100644
--- a/src/Core/Billing/Services/IPremiumUserBillingService.cs
+++ b/src/Core/Billing/Services/IPremiumUserBillingService.cs
@@ -1,4 +1,5 @@
-using Bit.Core.Billing.Models.Sales;
+using Bit.Core.Billing.Models;
+using Bit.Core.Billing.Models.Sales;
 using Bit.Core.Entities;
 
 namespace Bit.Core.Billing.Services;
@@ -27,4 +28,9 @@ public interface IPremiumUserBillingService
     /// </code>
     /// </example>
     Task Finalize(PremiumUserSale sale);
+
+    Task UpdatePaymentMethod(
+        User user,
+        TokenizedPaymentSource tokenizedPaymentSource,
+        TaxInformation taxInformation);
 }
diff --git a/src/Core/Billing/Services/Implementations/PremiumUserBillingService.cs b/src/Core/Billing/Services/Implementations/PremiumUserBillingService.cs
index 0672a8d5e7..ed841c9576 100644
--- a/src/Core/Billing/Services/Implementations/PremiumUserBillingService.cs
+++ b/src/Core/Billing/Services/Implementations/PremiumUserBillingService.cs
@@ -1,5 +1,6 @@
 using Bit.Core.Billing.Caches;
 using Bit.Core.Billing.Constants;
+using Bit.Core.Billing.Models;
 using Bit.Core.Billing.Models.Sales;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
@@ -58,6 +59,28 @@ public class PremiumUserBillingService(
         await userRepository.ReplaceAsync(user);
     }
 
+    public async Task UpdatePaymentMethod(
+        User user,
+        TokenizedPaymentSource tokenizedPaymentSource,
+        TaxInformation taxInformation)
+    {
+        if (string.IsNullOrEmpty(user.GatewayCustomerId))
+        {
+            var customer = await CreateCustomerAsync(user,
+                new CustomerSetup { TokenizedPaymentSource = tokenizedPaymentSource, TaxInformation = taxInformation });
+
+            user.Gateway = GatewayType.Stripe;
+            user.GatewayCustomerId = customer.Id;
+
+            await userRepository.ReplaceAsync(user);
+        }
+        else
+        {
+            await subscriberService.UpdatePaymentSource(user, tokenizedPaymentSource);
+            await subscriberService.UpdateTaxInformation(user, taxInformation);
+        }
+    }
+
     private async Task<Customer> CreateCustomerAsync(
         User user,
         CustomerSetup customerSetup)
diff --git a/src/Core/Services/Implementations/UserService.cs b/src/Core/Services/Implementations/UserService.cs
index 157bfd3a6e..11d4042def 100644
--- a/src/Core/Services/Implementations/UserService.cs
+++ b/src/Core/Services/Implementations/UserService.cs
@@ -9,6 +9,7 @@ using Bit.Core.AdminConsole.Services;
 using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Models;
 using Bit.Core.Auth.Models.Business.Tokenables;
+using Bit.Core.Billing.Models;
 using Bit.Core.Billing.Models.Sales;
 using Bit.Core.Billing.Services;
 using Bit.Core.Context;
@@ -1044,11 +1045,11 @@ public class UserService : UserManager<User>, IUserService, IDisposable
             throw new BadRequestException("Invalid token.");
         }
 
-        var updated = await _paymentService.UpdatePaymentMethodAsync(user, paymentMethodType, paymentToken, taxInfo: taxInfo);
-        if (updated)
-        {
-            await SaveUserAsync(user);
-        }
+        var tokenizedPaymentSource = new TokenizedPaymentSource(paymentMethodType, paymentToken);
+        var taxInformation = TaxInformation.From(taxInfo);
+
+        await _premiumUserBillingService.UpdatePaymentMethod(user, tokenizedPaymentSource, taxInformation);
+        await SaveUserAsync(user);
     }
 
     public async Task CancelPremiumAsync(User user, bool? endOfPeriod = null)