[PM-29556] Fix: changing organization plan nulls out public and private keys

src/Api/AdminConsole/Controllers/ProviderClientsController.cs

@@ -57,8 +57,7 @@ public async Task<IResult> CreateAsync(
             Owner = user,
             BillingEmail = provider.BillingEmail,
             OwnerKey = requestBody.Key,
-            PublicKey = requestBody.KeyPair.PublicKey,
-            PrivateKey = requestBody.KeyPair.EncryptedPrivateKey,
+            Keys = requestBody.KeyPair.ToPublicKeyEncryptionKeyPairData(),
             CollectionName = requestBody.CollectionName,
             IsFromProvider = true
         };

src/Api/AdminConsole/Models/Request/Organizations/OrganizationCreateRequestModel.cs

@@ -113,11 +113,10 @@ public virtual OrganizationSignup ToOrganizationSignup(User user)
                 BillingAddressCountry = BillingAddressCountry,
             },
             InitiationPath = InitiationPath,
-            SkipTrial = SkipTrial
+            SkipTrial = SkipTrial,
+            Keys = Keys?.ToPublicKeyEncryptionKeyPairData()
         };
 
-        Keys?.ToOrganizationSignup(orgSignup);
-
         return orgSignup;
     }
 

src/Api/AdminConsole/Models/Request/Organizations/OrganizationKeysRequestModel.cs

@@ -2,8 +2,7 @@
 #nullable disable
 
 using System.ComponentModel.DataAnnotations;
-using Bit.Core.AdminConsole.Entities;
-using Bit.Core.Models.Business;
+using Bit.Core.KeyManagement.Models.Data;
 
 namespace Bit.Api.AdminConsole.Models.Request.Organizations;
 
@@ -14,48 +13,10 @@ public class OrganizationKeysRequestModel
     [Required]
     public string EncryptedPrivateKey { get; set; }
 
-    public OrganizationSignup ToOrganizationSignup(OrganizationSignup existingSignup)
+    public PublicKeyEncryptionKeyPairData ToPublicKeyEncryptionKeyPairData()
     {
-        if (string.IsNullOrWhiteSpace(existingSignup.PublicKey))
-        {
-            existingSignup.PublicKey = PublicKey;
-        }
-
-        if (string.IsNullOrWhiteSpace(existingSignup.PrivateKey))
-        {
-            existingSignup.PrivateKey = EncryptedPrivateKey;
-        }
-
-        return existingSignup;
-    }
-
-    public OrganizationUpgrade ToOrganizationUpgrade(OrganizationUpgrade existingUpgrade)
-    {
-        if (string.IsNullOrWhiteSpace(existingUpgrade.PublicKey))
-        {
-            existingUpgrade.PublicKey = PublicKey;
-        }
-
-        if (string.IsNullOrWhiteSpace(existingUpgrade.PrivateKey))
-        {
-            existingUpgrade.PrivateKey = EncryptedPrivateKey;
-        }
-
-        return existingUpgrade;
-    }
-
-    public Organization ToOrganization(Organization existingOrg)
-    {
-        if (string.IsNullOrWhiteSpace(existingOrg.PublicKey))
-        {
-            existingOrg.PublicKey = PublicKey;
-        }
-
-        if (string.IsNullOrWhiteSpace(existingOrg.PrivateKey))
-        {
-            existingOrg.PrivateKey = EncryptedPrivateKey;
-        }
-
-        return existingOrg;
+        return new PublicKeyEncryptionKeyPairData(
+            wrappedPrivateKey: EncryptedPrivateKey,
+            publicKey: PublicKey);
     }
 }

src/Api/AdminConsole/Models/Request/Organizations/OrganizationNoPaymentCreateRequest.cs

@@ -110,10 +110,9 @@ public virtual OrganizationSignup ToOrganizationSignup(User user)
                 BillingAddressCountry = BillingAddressCountry,
             },
             InitiationPath = InitiationPath,
+            Keys = Keys?.ToPublicKeyEncryptionKeyPairData()
         };
 
-        Keys?.ToOrganizationSignup(orgSignup);
-
         return orgSignup;
     }
 }

src/Api/AdminConsole/Models/Request/Organizations/OrganizationUpdateRequestModel.cs

@@ -22,7 +22,6 @@ public class OrganizationUpdateRequestModel
         OrganizationId = organizationId,
         Name = Name,
         BillingEmail = BillingEmail,
-        PublicKey = Keys?.PublicKey,
-        EncryptedPrivateKey = Keys?.EncryptedPrivateKey
+        Keys = Keys?.ToPublicKeyEncryptionKeyPairData()
     };
 }

src/Api/AdminConsole/Models/Request/Organizations/OrganizationUpgradeRequestModel.cs

@@ -43,11 +43,10 @@ public OrganizationUpgrade ToOrganizationUpgrade()
             {
                 BillingAddressCountry = BillingAddressCountry,
                 BillingAddressPostalCode = BillingAddressPostalCode
-            }
+            },
+            Keys = Keys?.ToPublicKeyEncryptionKeyPairData()
         };
 
-        Keys?.ToOrganizationUpgrade(orgUpgrade);
-
         return orgUpgrade;
     }
 }

src/Api/Billing/Models/Requests/KeyPairRequestBody.cs

@@ -2,6 +2,7 @@
 #nullable disable
 
 using System.ComponentModel.DataAnnotations;
+using Bit.Core.KeyManagement.Models.Data;
 
 namespace Bit.Api.Billing.Models.Requests;
 
@@ -12,4 +13,11 @@ public class KeyPairRequestBody
     public string PublicKey { get; set; }
     [Required(ErrorMessage = "'encryptedPrivateKey' must be provided")]
     public string EncryptedPrivateKey { get; set; }
+
+    public PublicKeyEncryptionKeyPairData ToPublicKeyEncryptionKeyPairData()
+    {
+        return new PublicKeyEncryptionKeyPairData(
+            wrappedPrivateKey: EncryptedPrivateKey,
+            publicKey: PublicKey);
+    }
 }

src/Core/AdminConsole/OrganizationFeatures/Organizations/CloudOrganizationSignUpCommand.cs

@@ -99,8 +99,8 @@ public async Task<SignUpOrganizationResponse> SignUpOrganizationAsync(Organizati
             ReferenceData = signup.Owner.ReferenceData,
             Enabled = true,
             LicenseKey = CoreHelpers.SecureRandomString(20),
-            PublicKey = signup.PublicKey,
-            PrivateKey = signup.PrivateKey,
+            PublicKey = signup.Keys?.PublicKey,
+            PrivateKey = signup.Keys?.WrappedPrivateKey,
             CreationDate = DateTime.UtcNow,
             RevisionDate = DateTime.UtcNow,
             Status = OrganizationStatusType.Created,

src/Core/AdminConsole/OrganizationFeatures/Organizations/OrganizationExtensions.cs

@@ -0,0 +1,28 @@
+using Bit.Core.AdminConsole.Entities;
+using Bit.Core.KeyManagement.Models.Data;
+
+namespace Bit.Core.AdminConsole.OrganizationFeatures.Organizations;
+
+public static class OrganizationExtensions
+{
+    /// <summary>
+    /// Updates the organization public and private keys if provided and not already set.
+    /// This is legacy code for old organizations that were not created with a public/private keypair.
+    /// It is a soft migration that will silently migrate organizations when they perform certain actions,
+    /// e.g. change their details or upgrade their plan.
+    /// </summary>
+    public static void BackfillPublicPrivateKeys(this Organization organization, PublicKeyEncryptionKeyPairData? keyPair)
+    {
+        // Only backfill if both new keys are provided and both old keys are missing.
+        if (string.IsNullOrWhiteSpace(keyPair?.PublicKey) ||
+            string.IsNullOrWhiteSpace(keyPair.WrappedPrivateKey) ||
+            !string.IsNullOrWhiteSpace(organization.PublicKey) ||
+            !string.IsNullOrWhiteSpace(organization.PrivateKey))
+        {
+            return;
+        }
+
+        organization.PublicKey = keyPair.PublicKey;
+        organization.PrivateKey = keyPair.WrappedPrivateKey;
+    }
+}

src/Core/AdminConsole/OrganizationFeatures/Organizations/ProviderClientOrganizationSignUpCommand.cs

@@ -93,8 +93,8 @@ public async Task<ProviderClientOrganizationSignUpResponse> SignUpClientOrganiza
             ReferenceData = signup.Owner.ReferenceData,
             Enabled = true,
             LicenseKey = CoreHelpers.SecureRandomString(20),
-            PublicKey = signup.PublicKey,
-            PrivateKey = signup.PrivateKey,
+            PublicKey = signup.Keys?.PublicKey,
+            PrivateKey = signup.Keys?.WrappedPrivateKey,
             CreationDate = DateTime.UtcNow,
             RevisionDate = DateTime.UtcNow,
             Status = OrganizationStatusType.Created,

src/Core/AdminConsole/OrganizationFeatures/Organizations/Update/OrganizationUpdateCommand.cs

@@ -39,8 +39,20 @@ private async Task<Organization> UpdateCloudAsync(Organization organization, Org
         var originalBillingEmail = organization.BillingEmail;
 
         // Apply updates to organization
-        organization.UpdateDetails(request);
-        organization.BackfillPublicPrivateKeys(request);
+        // These values may or may not be sent by the client depending on the operation being performed.
+        // Skip any values not provided.
+        if (request.Name is not null)
+        {
+            organization.Name = request.Name;
+        }
+
+        if (request.BillingEmail is not null)
+        {
+            organization.BillingEmail = request.BillingEmail.ToLowerInvariant().Trim();
+        }
+
+        organization.BackfillPublicPrivateKeys(request.Keys);
+
         await organizationService.ReplaceAndUpdateCacheAsync(organization, EventType.Organization_Updated);
 
         // Update billing information in Stripe if required
@@ -56,7 +68,7 @@ private async Task<Organization> UpdateCloudAsync(Organization organization, Org
     /// </summary>
     private async Task<Organization> UpdateSelfHostedAsync(Organization organization, OrganizationUpdateRequest request)
     {
-        organization.BackfillPublicPrivateKeys(request);
+        organization.BackfillPublicPrivateKeys(request.Keys);
         await organizationService.ReplaceAndUpdateCacheAsync(organization, EventType.Organization_Updated);
         return organization;
     }

src/Core/AdminConsole/OrganizationFeatures/Organizations/Update/OrganizationUpdateRequest.cs

@@ -1,4 +1,6 @@
-namespace Bit.Core.AdminConsole.OrganizationFeatures.Organizations.Update;
+using Bit.Core.KeyManagement.Models.Data;
+
+namespace Bit.Core.AdminConsole.OrganizationFeatures.Organizations.Update;
 
 /// <summary>
 /// Request model for updating the name, billing email, and/or public-private keys for an organization (legacy migration code).
@@ -22,12 +24,7 @@ public record OrganizationUpdateRequest
     public string? BillingEmail { get; init; }
 
     /// <summary>
-    /// The organization's public key to set (optional, only set if not already present on the organization).
-    /// </summary>
-    public string? PublicKey { get; init; }
-
-    /// <summary>
-    /// The organization's encrypted private key to set (optional, only set if not already present on the organization).
+    /// The organization's public/private key pair to set (optional, only set if not already present on the organization).
     /// </summary>
-    public string? EncryptedPrivateKey { get; init; }
+    public PublicKeyEncryptionKeyPairData? Keys { get; init; }
 }

src/Core/Models/Business/OrganizationUpgrade.cs

@@ -2,6 +2,7 @@
 #nullable disable
 
 using Bit.Core.Billing.Enums;
+using Bit.Core.KeyManagement.Models.Data;
 
 namespace Bit.Core.Models.Business;
 
@@ -13,8 +14,7 @@ public class OrganizationUpgrade
     public short AdditionalStorageGb { get; set; }
     public bool PremiumAccessAddon { get; set; }
     public TaxInfo TaxInfo { get; set; }
-    public string PublicKey { get; set; }
-    public string PrivateKey { get; set; }
+    public PublicKeyEncryptionKeyPairData Keys { get; set; }
     public int? AdditionalSmSeats { get; set; }
     public int? AdditionalServiceAccounts { get; set; }
     public bool UseSecretsManager { get; set; }

src/Core/OrganizationFeatures/OrganizationSubscriptions/UpgradeOrganizationPlanCommand.cs

@@ -4,6 +4,7 @@
 using Bit.Core.AdminConsole.Entities;
 using Bit.Core.AdminConsole.Enums;
 using Bit.Core.AdminConsole.Models.OrganizationConnectionConfigs;
+using Bit.Core.AdminConsole.OrganizationFeatures.Organizations;
 using Bit.Core.AdminConsole.Repositories;
 using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Repositories;
@@ -256,27 +257,20 @@ await _organizationUserRepository.GetManyByOrganizationAsync(organization.Id,
         organization.SelfHost = newPlan.HasSelfHost;
         organization.UsePolicies = newPlan.HasPolicies;
         organization.MaxStorageGb = (short)(newPlan.PasswordManager.BaseStorageGb + upgrade.AdditionalStorageGb);
-        organization.UseGroups = newPlan.HasGroups;
-        organization.UseDirectory = newPlan.HasDirectory;
-        organization.UseEvents = newPlan.HasEvents;
-        organization.UseTotp = newPlan.HasTotp;
-        organization.Use2fa = newPlan.Has2fa;
-        organization.UseApi = newPlan.HasApi;
         organization.UseSso = newPlan.HasSso;
         organization.UseOrganizationDomains = newPlan.HasOrganizationDomains;
         organization.UseKeyConnector = newPlan.HasKeyConnector ? organization.UseKeyConnector : false;
         organization.UseScim = newPlan.HasScim;
         organization.UseResetPassword = newPlan.HasResetPassword;
-        organization.SelfHost = newPlan.HasSelfHost;
         organization.UsersGetPremium = newPlan.UsersGetPremium || upgrade.PremiumAccessAddon;
         organization.UseCustomPermissions = newPlan.HasCustomPermissions;
         organization.Plan = newPlan.Name;
         organization.Enabled = success;
-        organization.PublicKey = upgrade.PublicKey;
-        organization.PrivateKey = upgrade.PrivateKey;
         organization.UsePasswordManager = true;
         organization.UseSecretsManager = upgrade.UseSecretsManager;
 
+        organization.BackfillPublicPrivateKeys(upgrade.Keys);
+
         if (upgrade.UseSecretsManager)
         {
             organization.SmSeats = newPlan.SecretsManager.BaseSeats + upgrade.AdditionalSmSeats.GetValueOrDefault();

test/Api.Test/AdminConsole/Controllers/ProviderClientsControllerTests.cs

@@ -66,8 +66,8 @@ public async Task CreateAsync_OK(
                     signup.Plan == requestBody.PlanType &&
                     signup.AdditionalSeats == requestBody.Seats &&
                     signup.OwnerKey == requestBody.Key &&
-                    signup.PublicKey == requestBody.KeyPair.PublicKey &&
-                    signup.PrivateKey == requestBody.KeyPair.EncryptedPrivateKey &&
+                    signup.Keys.PublicKey == requestBody.KeyPair.PublicKey &&
+                    signup.Keys.WrappedPrivateKey == requestBody.KeyPair.EncryptedPrivateKey &&
                     signup.CollectionName == requestBody.CollectionName),
                 requestBody.OwnerEmail,
                 user)