[PM-21411] Refactor interface for determining premium status and features

src/Core/AdminConsole/Models/Data/Organizations/OrganizationUsers/OrganizationUserUserDetails.cs

@@ -20,6 +20,12 @@ public class OrganizationUserUserDetails : IExternal, ITwoFactorProvidersUser, I
     public string Email { get; set; }
     public string AvatarColor { get; set; }
     public string TwoFactorProviders { get; set; }
+    /// <summary>
+    /// Indicates whether the user has a personal premium subscription.
+    /// Does not include premium access from organizations -
+    /// do not use this to check whether the user can access premium features.
+    /// Null when the organization user is in Invited status (UserId is null).
+    /// </summary>
     public bool? Premium { get; set; }
     public OrganizationUserStatusType Status { get; set; }
     public OrganizationUserType Type { get; set; }

src/Core/Auth/UserFeatures/TwoFactorAuth/TwoFactorIsEnabledQuery.cs

@@ -4,16 +4,37 @@
 using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Models;
 using Bit.Core.Auth.UserFeatures.TwoFactorAuth.Interfaces;
+using Bit.Core.Billing.Premium.Queries;
+using Bit.Core.Entities;
+using Bit.Core.Exceptions;
 using Bit.Core.Repositories;
+using Bit.Core.Services;
 
 namespace Bit.Core.Auth.UserFeatures.TwoFactorAuth;
 
-public class TwoFactorIsEnabledQuery(IUserRepository userRepository) : ITwoFactorIsEnabledQuery
+public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
 {
-    private readonly IUserRepository _userRepository = userRepository;
+    private readonly IUserRepository _userRepository;
+    private readonly IHasPremiumAccessQuery _hasPremiumAccessQuery;
+    private readonly IFeatureService _featureService;
+
+    public TwoFactorIsEnabledQuery(
+        IUserRepository userRepository,
+        IHasPremiumAccessQuery hasPremiumAccessQuery,
+        IFeatureService featureService)
+    {
+        _userRepository = userRepository;
+        _hasPremiumAccessQuery = hasPremiumAccessQuery;
+        _featureService = featureService;
+    }
 
     public async Task<IEnumerable<(Guid userId, bool twoFactorIsEnabled)>> TwoFactorIsEnabledAsync(IEnumerable<Guid> userIds)
     {
+        if (_featureService.IsEnabled(FeatureFlagKeys.PremiumAccessQuery))
+        {
+            return await TwoFactorIsEnabledVNextAsync(userIds);
+        }
+
         var result = new List<(Guid userId, bool hasTwoFactor)>();
         if (userIds == null || !userIds.Any())
         {
@@ -36,6 +57,11 @@ await TwoFactorEnabledAsync(userDetail.GetTwoFactorProviders(),
 
     public async Task<IEnumerable<(T user, bool twoFactorIsEnabled)>> TwoFactorIsEnabledAsync<T>(IEnumerable<T> users) where T : ITwoFactorProvidersUser
     {
+        if (_featureService.IsEnabled(FeatureFlagKeys.PremiumAccessQuery))
+        {
+            return await TwoFactorIsEnabledVNextAsync(users);
+        }
+
         var userIds = users
             .Select(u => u.GetUserId())
             .Where(u => u.HasValue)
@@ -71,13 +97,134 @@ public async Task<bool> TwoFactorIsEnabledAsync(ITwoFactorProvidersUser user)
             return false;
         }
 
+        if (_featureService.IsEnabled(FeatureFlagKeys.PremiumAccessQuery))
+        {
+            var userEntity = user as User ?? await _userRepository.GetByIdAsync(userId.Value);
+            if (userEntity == null)
+            {
+                throw new NotFoundException();
+            }
+
+            return await TwoFactorIsEnabledVNextAsync(userEntity);
+        }
+
         return await TwoFactorEnabledAsync(
-                        user.GetTwoFactorProviders(),
-                        async () =>
-                        {
-                            var calcUser = await _userRepository.GetCalculatedPremiumAsync(userId.Value);
-                            return calcUser?.HasPremiumAccess ?? false;
-                        });
+            user.GetTwoFactorProviders(),
+            async () =>
+            {
+                var calcUser = await _userRepository.GetCalculatedPremiumAsync(userId.Value);
+                return calcUser?.HasPremiumAccess ?? false;
+            });
+    }
+
+    private async Task<IEnumerable<(Guid userId, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync(IEnumerable<Guid> userIds)
+    {
+        var result = new List<(Guid userId, bool hasTwoFactor)>();
+        if (userIds == null || !userIds.Any())
+        {
+            return result;
+        }
+
+        var users = await _userRepository.GetManyAsync([.. userIds]);
+
+        // Get enabled providers for each user
+        var usersTwoFactorProvidersMap = users.ToDictionary(u => u.Id, GetEnabledTwoFactorProviders);
+
+        // Bulk fetch premium status only for users who need it (those with only premium providers)
+        var userIdsNeedingPremium = usersTwoFactorProvidersMap
+            .Where(kvp => kvp.Value.Any() && kvp.Value.All(TwoFactorProvider.RequiresPremium))
+            .Select(kvp => kvp.Key)
+            .ToList();
+
+        var premiumStatusMap = userIdsNeedingPremium.Count > 0
+            ? await _hasPremiumAccessQuery.HasPremiumAccessAsync(userIdsNeedingPremium)
+            : new Dictionary<Guid, bool>();
+
+        foreach (var user in users)
+        {
+            var userTwoFactorProviders = usersTwoFactorProvidersMap[user.Id];
+
+            if (!userTwoFactorProviders.Any())
+            {
+                result.Add((user.Id, false));
+                continue;
+            }
+
+            // User has providers. If they're in the premium check map, verify premium status
+            var twoFactorIsEnabled = !premiumStatusMap.TryGetValue(user.Id, out var hasPremium) || hasPremium;
+            result.Add((user.Id, twoFactorIsEnabled));
+        }
+
+        return result;
+    }
+
+    private async Task<IEnumerable<(T user, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync<T>(IEnumerable<T> users)
+        where T : ITwoFactorProvidersUser
+    {
+        var userIds = users
+            .Select(u => u.GetUserId())
+            .Where(u => u.HasValue)
+            .Select(u => u.Value)
+            .ToList();
+
+        var twoFactorResults = await TwoFactorIsEnabledVNextAsync(userIds);
+
+        var result = new List<(T user, bool twoFactorIsEnabled)>();
+
+        foreach (var user in users)
+        {
+            var userId = user.GetUserId();
+            if (userId.HasValue)
+            {
+                var hasTwoFactor = twoFactorResults.FirstOrDefault(res => res.userId == userId.Value).twoFactorIsEnabled;
+                result.Add((user, hasTwoFactor));
+            }
+            else
+            {
+                result.Add((user, false));
+            }
+        }
+
+        return result;
+    }
+
+    private async Task<bool> TwoFactorIsEnabledVNextAsync(User user)
+    {
+        var enabledProviders = GetEnabledTwoFactorProviders(user);
+
+        if (!enabledProviders.Any())
+        {
+            return false;
+        }
+
+        // If all providers require premium, check if user has premium access
+        if (enabledProviders.All(TwoFactorProvider.RequiresPremium))
+        {
+            return await _hasPremiumAccessQuery.HasPremiumAccessAsync(user.Id);
+        }
+
+        // User has at least one non-premium provider
+        return true;
+    }
+
+    /// <summary>
+    /// Gets all enabled two-factor provider types for a user.
+    /// </summary>
+    /// <param name="user">user with two factor providers</param>
+    /// <returns>list of enabled provider types</returns>
+    private static IList<TwoFactorProviderType> GetEnabledTwoFactorProviders(User user)
+    {
+        var providers = user.GetTwoFactorProviders();
+
+        if (providers == null || providers.Count == 0)
+        {
+            return Array.Empty<TwoFactorProviderType>();
+        }
+
+        // TODO: PM-21210: In practice we don't save disabled providers to the database, worth looking into.
+        return (from provider in providers
+                where provider.Value?.Enabled ?? false
+                select provider.Key).ToList();
     }
 
     /// <summary>

src/Core/Billing/Extensions/ServiceCollectionExtensions.cs

@@ -6,6 +6,7 @@
 using Bit.Core.Billing.Organizations.Services;
 using Bit.Core.Billing.Payment;
 using Bit.Core.Billing.Premium.Commands;
+using Bit.Core.Billing.Premium.Queries;
 using Bit.Core.Billing.Pricing;
 using Bit.Core.Billing.Services;
 using Bit.Core.Billing.Services.Implementations;
@@ -31,6 +32,7 @@ public static void AddBillingOperations(this IServiceCollection services)
         services.AddPaymentOperations();
         services.AddOrganizationLicenseCommandsQueries();
         services.AddPremiumCommands();
+        services.AddPremiumQueries();
         services.AddTransient<IGetOrganizationMetadataQuery, GetOrganizationMetadataQuery>();
         services.AddTransient<IGetOrganizationWarningsQuery, GetOrganizationWarningsQuery>();
         services.AddTransient<IRestartSubscriptionCommand, RestartSubscriptionCommand>();
@@ -50,4 +52,9 @@ private static void AddPremiumCommands(this IServiceCollection services)
         services.AddScoped<ICreatePremiumSelfHostedSubscriptionCommand, CreatePremiumSelfHostedSubscriptionCommand>();
         services.AddTransient<IPreviewPremiumTaxCommand, PreviewPremiumTaxCommand>();
     }
+
+    private static void AddPremiumQueries(this IServiceCollection services)
+    {
+        services.AddScoped<IHasPremiumAccessQuery, HasPremiumAccessQuery>();
+    }
 }

src/Core/Billing/Premium/Models/UserPremiumAccess.cs

@@ -0,0 +1,29 @@
+namespace Bit.Core.Billing.Premium.Models;
+
+/// <summary>
+/// Represents user premium access status from personal subscriptions and organization memberships.
+/// </summary>
+public class UserPremiumAccess
+{
+    /// <summary>
+    /// The unique identifier for the user.
+    /// </summary>
+    public Guid Id { get; set; }
+
+    /// <summary>
+    /// Indicates whether the user has a personal premium subscription.
+    /// This does NOT include premium access from organizations.
+    /// </summary>
+    public bool PersonalPremium { get; set; }
+
+    /// <summary>
+    /// Indicates whether the user has premium access through any organization membership.
+    /// This is true if the user is a member of at least one enabled organization that grants premium access to users.
+    /// </summary>
+    public bool OrganizationPremium { get; set; }
+
+    /// <summary>
+    /// Indicates whether the user has premium access from any source (personal subscription or organization).
+    /// </summary>
+    public bool HasPremiumAccess => PersonalPremium || OrganizationPremium;
+}

src/Core/Billing/Premium/Queries/HasPremiumAccessQuery.cs

@@ -0,0 +1,49 @@
+using Bit.Core.Exceptions;
+using Bit.Core.Repositories;
+
+namespace Bit.Core.Billing.Premium.Queries;
+
+public class HasPremiumAccessQuery : IHasPremiumAccessQuery
+{
+    private readonly IUserRepository _userRepository;
+
+    public HasPremiumAccessQuery(IUserRepository userRepository)
+    {
+        _userRepository = userRepository;
+    }
+
+    public async Task<bool> HasPremiumAccessAsync(Guid userId)
+    {
+        var user = await _userRepository.GetPremiumAccessAsync(userId);
+        if (user == null)
+        {
+            throw new NotFoundException();
+        }
+
+        return user.HasPremiumAccess;
+    }
+
+    public async Task<Dictionary<Guid, bool>> HasPremiumAccessAsync(IEnumerable<Guid> userIds)
+    {
+        var distinctUserIds = userIds.Distinct().ToList();
+        var usersWithPremium = await _userRepository.GetPremiumAccessByIdsAsync(distinctUserIds);
+
+        if (usersWithPremium.Count() != distinctUserIds.Count)
+        {
+            throw new NotFoundException();
+        }
+
+        return usersWithPremium.ToDictionary(u => u.Id, u => u.HasPremiumAccess);
+    }
+
+    public async Task<bool> HasPremiumFromOrganizationAsync(Guid userId)
+    {
+        var user = await _userRepository.GetPremiumAccessAsync(userId);
+        if (user == null)
+        {
+            throw new NotFoundException();
+        }
+
+        return user.OrganizationPremium;
+    }
+}

src/Core/Billing/Premium/Queries/IHasPremiumAccessQuery.cs

@@ -0,0 +1,30 @@
+namespace Bit.Core.Billing.Premium.Queries;
+
+/// <summary>
+/// Centralized query for checking if users have premium access through personal subscriptions or organizations.
+/// Note: Different from User.Premium which only checks personal subscriptions.
+/// </summary>
+public interface IHasPremiumAccessQuery
+{
+    /// <summary>
+    /// Checks if a user has premium access (personal or organization).
+    /// </summary>
+    /// <param name="userId">The user ID to check</param>
+    /// <returns>True if user can access premium features</returns>
+    Task<bool> HasPremiumAccessAsync(Guid userId);
+
+    /// <summary>
+    /// Checks premium access for multiple users.
+    /// </summary>
+    /// <param name="userIds">The user IDs to check</param>
+    /// <returns>Dictionary mapping user IDs to their premium access status</returns>
+    Task<Dictionary<Guid, bool>> HasPremiumAccessAsync(IEnumerable<Guid> userIds);
+
+    /// <summary>
+    /// Checks if a user belongs to any organization that grants premium (enabled org with UsersGetPremium).
+    /// Returns true regardless of personal subscription. Useful for UI decisions like showing subscription options.
+    /// </summary>
+    /// <param name="userId">The user ID to check</param>
+    /// <returns>True if user is in any organization that grants premium</returns>
+    Task<bool> HasPremiumFromOrganizationAsync(Guid userId);
+}

src/Core/Constants.cs

@@ -143,6 +143,7 @@ public static class FeatureFlagKeys
     public const string BlockClaimedDomainAccountCreation = "pm-28297-block-uninvited-claimed-domain-registration";
     public const string IncreaseBulkReinviteLimitForCloud = "pm-28251-increase-bulk-reinvite-limit-for-cloud";
     public const string BulkRevokeUsersV2 = "pm-28456-bulk-revoke-users-v2";
+    public const string PremiumAccessQuery = "pm-21411-premium-access-query";
 
     /* Architecture */
     public const string DesktopMigrationMilestone1 = "desktop-ui-migration-milestone-1";

src/Core/Entities/User.cs

@@ -69,6 +69,11 @@ public class User : ITableObject<Guid>, IStorableSubscriber, IRevisable, ITwoFac
     /// The security state is a signed object attesting to the version of the user's account.
     /// </summary>
     public string? SecurityState { get; set; }
+    /// <summary>
+    /// Indicates whether the user has a personal premium subscription.
+    /// Does not include premium access from organizations -
+    /// do not use this to check whether the user can access premium features.
+    /// </summary>
     public bool Premium { get; set; }
     public DateTime? PremiumExpirationDate { get; set; }
     public DateTime? RenewalReminderDate { get; set; }