From 03b0dc04c40be09f17f0d3809defa944464312b6 Mon Sep 17 00:00:00 2001
From: Rui Tome <rtome@bitwarden.com>
Date: Thu, 22 May 2025 16:36:08 +0100
Subject: [PATCH] Enhance RegisterUserCommand to integrate organization policy
 requirements for two-factor authentication. Added support for feature flags
 and refactored email two-factor provider setup into a separate method.

---
 .../Implementations/RegisterUserCommand.cs    | 53 ++++++++++++++-----
 1 file changed, 40 insertions(+), 13 deletions(-)

diff --git a/src/Core/Auth/UserFeatures/Registration/Implementations/RegisterUserCommand.cs b/src/Core/Auth/UserFeatures/Registration/Implementations/RegisterUserCommand.cs
index e721649dc9..00d48c0241 100644
--- a/src/Core/Auth/UserFeatures/Registration/Implementations/RegisterUserCommand.cs
+++ b/src/Core/Auth/UserFeatures/Registration/Implementations/RegisterUserCommand.cs
@@ -1,4 +1,6 @@
 using Bit.Core.AdminConsole.Enums;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies.PolicyRequirements;
 using Bit.Core.AdminConsole.Repositories;
 using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Models;
@@ -42,6 +44,9 @@ public class RegisterUserCommand : IRegisterUserCommand
 
     private readonly IDataProtectorTokenFactory<EmergencyAccessInviteTokenable> _emergencyAccessInviteTokenDataFactory;
 
+    private readonly IOrganizationPolicyRequirementQuery _organizationPolicyRequirementQuery;
+    private readonly IFeatureService _featureService;
+
     private readonly string _disabledUserRegistrationExceptionMsg = "Open registration has been disabled by the system administrator.";
 
     public RegisterUserCommand(
@@ -56,7 +61,9 @@ public class RegisterUserCommand : IRegisterUserCommand
         IUserService userService,
         IMailService mailService,
         IValidateRedemptionTokenCommand validateRedemptionTokenCommand,
-        IDataProtectorTokenFactory<EmergencyAccessInviteTokenable> emergencyAccessInviteTokenDataFactory
+        IDataProtectorTokenFactory<EmergencyAccessInviteTokenable> emergencyAccessInviteTokenDataFactory,
+        IOrganizationPolicyRequirementQuery organizationPolicyRequirementQuery,
+        IFeatureService featureService
         )
     {
         _globalSettings = globalSettings;
@@ -76,6 +83,9 @@ public class RegisterUserCommand : IRegisterUserCommand
         _validateRedemptionTokenCommand = validateRedemptionTokenCommand;
         _emergencyAccessInviteTokenDataFactory = emergencyAccessInviteTokenDataFactory;
 
+        _organizationPolicyRequirementQuery = organizationPolicyRequirementQuery;
+        _featureService = featureService;
+
         _providerServiceDataProtector = dataProtectionProvider.CreateProtector("ProviderServiceDataProtector");
     }
 
@@ -214,24 +224,41 @@ public class RegisterUserCommand : IRegisterUserCommand
         var orgUser = await _organizationUserRepository.GetByIdAsync(orgUserId.Value);
         if (orgUser != null)
         {
-            var twoFactorPolicy = await _policyRepository.GetByOrganizationIdTypeAsync(orgUser.OrganizationId,
-                PolicyType.TwoFactorAuthentication);
-            if (twoFactorPolicy != null && twoFactorPolicy.Enabled)
+            if (_featureService.IsEnabled(FeatureFlagKeys.PolicyRequirements))
             {
-                user.SetTwoFactorProviders(new Dictionary<TwoFactorProviderType, TwoFactorProvider>
-                {
+                var twoFactorPolicyRequirement = await _organizationPolicyRequirementQuery
+                    .GetAsync<OrganizationTwoFactorPolicyRequirement>(orgUser.OrganizationId);
 
-                    [TwoFactorProviderType.Email] = new TwoFactorProvider
-                    {
-                        MetaData = new Dictionary<string, object> { ["Email"] = user.Email.ToLowerInvariant() },
-                        Enabled = true
-                    }
-                });
-                _userService.SetTwoFactorProvider(user, TwoFactorProviderType.Email);
+                if (twoFactorPolicyRequirement.IsRequired)
+                {
+                    EnableEmailTwoFactorForUser(user);
+                }
+            }
+            else
+            {
+                var twoFactorPolicy = await _policyRepository.GetByOrganizationIdTypeAsync(orgUser.OrganizationId,
+                    PolicyType.TwoFactorAuthentication);
+                if (twoFactorPolicy != null && twoFactorPolicy.Enabled)
+                {
+                    EnableEmailTwoFactorForUser(user);
+                }
             }
         }
     }
 
+    private void EnableEmailTwoFactorForUser(User user)
+    {
+        user.SetTwoFactorProviders(new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            [TwoFactorProviderType.Email] = new TwoFactorProvider
+            {
+                MetaData = new Dictionary<string, object> { ["Email"] = user.Email.ToLowerInvariant() },
+                Enabled = true
+            }
+        });
+        _userService.SetTwoFactorProvider(user, TwoFactorProviderType.Email);
+    }
+
 
     private async Task SendAppropriateWelcomeEmailAsync(User user, string initiationPath)
     {