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

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

@@ -20,6 +20,10 @@ public class OrganizationUserUserDetails : IExternal, ITwoFactorProvidersUser, I
     public string Email { get; set; }
     public string AvatarColor { get; set; }
     public string TwoFactorProviders { get; set; }
+    /// <summary>
+    /// User's personal premium subscription status. Does not reflect organization premium access.
+    /// 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/PremiumAccess/IPremiumAccessQuery.cs

@@ -0,0 +1,50 @@
+using Bit.Core.Entities;
+
+namespace Bit.Core.Auth.UserFeatures.PremiumAccess;
+
+/// <summary>
+/// Query for checking premium access status for users.
+/// This is the centralized location for determining if a user can access premium features
+/// (either through personal subscription or organization membership).
+/// 
+/// <para>
+/// <strong>Note:</strong> This is different from checking User.Premium, which only indicates
+/// personal subscription status. Use these methods to check actual premium feature access.
+/// </para>
+/// </summary>
+public interface IPremiumAccessQuery
+{
+    /// <summary>
+    /// Checks if a user has access to premium features (personal subscription or organization).
+    /// This is the definitive way to check premium access for a single user.
+    /// </summary>
+    /// <param name="user">The user to check for premium access</param>
+    /// <returns>True if user can access premium features; false otherwise</returns>
+    Task<bool> CanAccessPremiumAsync(User user);
+
+    /// <summary>
+    /// Checks if a user has access to premium features (personal subscription or organization).
+    /// Use this overload when you already know the personal premium status and only need to check organization premium.
+    /// </summary>
+    /// <param name="userId">The user ID to check for premium access</param>
+    /// <param name="hasPersonalPremium">Whether the user has a personal premium subscription</param>
+    /// <returns>True if user can access premium features; false otherwise</returns>
+    Task<bool> CanAccessPremiumAsync(Guid userId, bool hasPersonalPremium);
+
+    /// <summary>
+    /// Checks if a user has access to premium features through organization membership only.
+    /// This is useful for determining the source of premium access (personal vs organization).
+    /// </summary>
+    /// <param name="userId">The user ID to check for organization premium access</param>
+    /// <returns>True if user has premium access through any organization; false otherwise</returns>
+    Task<bool> HasPremiumFromOrganizationAsync(Guid userId);
+
+    /// <summary>
+    /// Checks if multiple users have access to premium features (optimized bulk operation).
+    /// Uses cached organization abilities and minimizes database queries.
+    /// </summary>
+    /// <param name="users">The users to check for premium access</param>
+    /// <returns>Dictionary mapping user IDs to their premium access status (personal or through organization)</returns>
+    Task<Dictionary<Guid, bool>> CanAccessPremiumBulkAsync(IEnumerable<User> users);
+}
+

src/Core/Auth/UserFeatures/PremiumAccess/PremiumAccessQuery.cs

@@ -0,0 +1,97 @@
+using Bit.Core.Entities;
+using Bit.Core.Repositories;
+using Bit.Core.Services;
+
+namespace Bit.Core.Auth.UserFeatures.PremiumAccess;
+
+/// <summary>
+/// Query for checking premium access status for users using cached organization abilities.
+/// </summary>
+public class PremiumAccessQuery : IPremiumAccessQuery
+{
+    private readonly IOrganizationUserRepository _organizationUserRepository;
+    private readonly IApplicationCacheService _applicationCacheService;
+
+    public PremiumAccessQuery(
+        IOrganizationUserRepository organizationUserRepository,
+        IApplicationCacheService applicationCacheService)
+    {
+        _organizationUserRepository = organizationUserRepository;
+        _applicationCacheService = applicationCacheService;
+    }
+
+    public async Task<bool> CanAccessPremiumAsync(User user)
+    {
+        return await CanAccessPremiumAsync(user.Id, user.Premium);
+    }
+
+    public async Task<bool> CanAccessPremiumAsync(Guid userId, bool hasPersonalPremium)
+    {
+        if (hasPersonalPremium)
+        {
+            return true;
+        }
+
+        return await HasPremiumFromOrganizationAsync(userId);
+    }
+
+    public async Task<bool> HasPremiumFromOrganizationAsync(Guid userId)
+    {
+        // Note: GetManyByUserAsync only returns Accepted and Confirmed status org users
+        var orgUsers = await _organizationUserRepository.GetManyByUserAsync(userId);
+        if (!orgUsers.Any())
+        {
+            return false;
+        }
+
+        var orgAbilities = await _applicationCacheService.GetOrganizationAbilitiesAsync();
+        return orgUsers.Any(ou =>
+            orgAbilities.TryGetValue(ou.OrganizationId, out var orgAbility) &&
+            orgAbility.UsersGetPremium &&
+            orgAbility.Enabled);
+    }
+
+    public async Task<Dictionary<Guid, bool>> CanAccessPremiumBulkAsync(IEnumerable<User> users)
+    {
+        var result = new Dictionary<Guid, bool>();
+        var usersList = users.ToList();
+
+        if (!usersList.Any())
+        {
+            return result;
+        }
+
+        var userIds = usersList.Select(u => u.Id).ToList();
+
+        // Get all org memberships for these users in one query
+        // Note: GetManyByManyUsersAsync only returns Accepted and Confirmed status org users
+        var allOrgUsers = await _organizationUserRepository.GetManyByManyUsersAsync(userIds);
+        var orgUsersGrouped = allOrgUsers
+            .Where(ou => ou.UserId.HasValue)
+            .GroupBy(ou => ou.UserId!.Value)
+            .ToDictionary(g => g.Key, g => g.ToList());
+
+        var orgAbilities = await _applicationCacheService.GetOrganizationAbilitiesAsync();
+
+        foreach (var user in usersList)
+        {
+            var hasPersonalPremium = user.Premium;
+            if (hasPersonalPremium)
+            {
+                result[user.Id] = true;
+                continue;
+            }
+
+            var hasPremiumFromOrg = orgUsersGrouped.TryGetValue(user.Id, out var userOrgs) &&
+                userOrgs.Any(ou =>
+                    orgAbilities.TryGetValue(ou.OrganizationId, out var orgAbility) &&
+                    orgAbility.UsersGetPremium &&
+                    orgAbility.Enabled);
+
+            result[user.Id] = hasPremiumFromOrg;
+        }
+
+        return result;
+    }
+}
+

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

@@ -3,14 +3,30 @@
 
 using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Models;
+using Bit.Core.Auth.UserFeatures.PremiumAccess;
 using Bit.Core.Auth.UserFeatures.TwoFactorAuth.Interfaces;
+using Bit.Core.Entities;
+using Bit.Core.Models.Data.Organizations.OrganizationUsers;
 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 IPremiumAccessQuery _premiumAccessQuery;
+    private readonly IFeatureService _featureService;
+
+    public TwoFactorIsEnabledQuery(
+        IUserRepository userRepository,
+        IPremiumAccessQuery premiumAccessQuery,
+        IFeatureService featureService)
+    {
+        _userRepository = userRepository;
+        _premiumAccessQuery = premiumAccessQuery;
+        _featureService = featureService;
+    }
 
     public async Task<IEnumerable<(Guid userId, bool twoFactorIsEnabled)>> TwoFactorIsEnabledAsync(IEnumerable<Guid> userIds)
     {
@@ -20,15 +36,34 @@ public class TwoFactorIsEnabledQuery(IUserRepository userRepository) : ITwoFacto
             return result;
         }
 
-        var userDetails = await _userRepository.GetManyWithCalculatedPremiumAsync([.. userIds]);
-        foreach (var userDetail in userDetails)
+        if (_featureService.IsEnabled(FeatureFlagKeys.PremiumAccessCacheCheck))
         {
-            result.Add(
-                (userDetail.Id,
-                 await TwoFactorEnabledAsync(userDetail.GetTwoFactorProviders(),
-                                () => Task.FromResult(userDetail.HasPremiumAccess))
-                )
-            );
+            var users = await _userRepository.GetManyAsync([.. userIds]);
+            var premiumStatus = await _premiumAccessQuery.CanAccessPremiumBulkAsync(users);
+
+            foreach (var user in users)
+            {
+                result.Add(
+                    (user.Id,
+                     await TwoFactorEnabledAsync(
+                         user.GetTwoFactorProviders(),
+                         () => Task.FromResult(premiumStatus.GetValueOrDefault(user.Id, false))
+                     ))
+                );
+            }
+        }
+        else
+        {
+            var userDetails = await _userRepository.GetManyWithCalculatedPremiumAsync([.. userIds]);
+            foreach (var userDetail in userDetails)
+            {
+                result.Add(
+                    (userDetail.Id,
+                     await TwoFactorEnabledAsync(userDetail.GetTwoFactorProviders(),
+                                    () => Task.FromResult(userDetail.HasPremiumAccess))
+                    )
+                );
+            }
         }
 
         return result;
@@ -71,13 +106,43 @@ public async Task<bool> TwoFactorIsEnabledAsync(ITwoFactorProvidersUser user)
             return false;
         }
 
-        return await TwoFactorEnabledAsync(
-                        user.GetTwoFactorProviders(),
-                        async () =>
-                        {
-                            var calcUser = await _userRepository.GetCalculatedPremiumAsync(userId.Value);
-                            return calcUser?.HasPremiumAccess ?? false;
-                        });
+        if (_featureService.IsEnabled(FeatureFlagKeys.PremiumAccessCacheCheck))
+        {
+            // Try to get premium status without fetching User entity if possible
+            bool hasPersonalPremium;
+            if (user is User userEntity)
+            {
+                hasPersonalPremium = userEntity.Premium;
+            }
+            else if (user is OrganizationUserUserDetails orgUserDetails)
+            {
+                hasPersonalPremium = orgUserDetails.Premium.GetValueOrDefault(false);
+            }
+            else
+            {
+                // Fallback: fetch the User entity
+                var fetchedUser = await _userRepository.GetByIdAsync(userId.Value);
+                if (fetchedUser == null)
+                {
+                    return false;
+                }
+                hasPersonalPremium = fetchedUser.Premium;
+            }
+
+            return await TwoFactorEnabledAsync(
+                user.GetTwoFactorProviders(),
+                async () => await _premiumAccessQuery.CanAccessPremiumAsync(userId.Value, hasPersonalPremium));
+        }
+        else
+        {
+            return await TwoFactorEnabledAsync(
+                user.GetTwoFactorProviders(),
+                async () =>
+                {
+                    var calcUser = await _userRepository.GetCalculatedPremiumAsync(userId.Value);
+                    return calcUser?.HasPremiumAccess ?? false;
+                });
+        }
     }
 
     /// <summary>

src/Core/Auth/UserFeatures/UserServiceCollectionExtensions.cs

@@ -1,5 +1,6 @@
 using Bit.Core.Auth.Sso;
 using Bit.Core.Auth.UserFeatures.DeviceTrust;
+using Bit.Core.Auth.UserFeatures.PremiumAccess;
 using Bit.Core.Auth.UserFeatures.Registration;
 using Bit.Core.Auth.UserFeatures.Registration.Implementations;
 using Bit.Core.Auth.UserFeatures.TdeOffboardingPassword.Interfaces;
@@ -27,6 +28,7 @@ public static void AddUserServices(this IServiceCollection services, IGlobalSett
         services.AddUserRegistrationCommands();
         services.AddWebAuthnLoginCommands();
         services.AddTdeOffboardingPasswordCommands();
+        services.AddPremiumAccessQueries();
         services.AddTwoFactorQueries();
         services.AddSsoQueries();
     }
@@ -65,6 +67,11 @@ private static void AddWebAuthnLoginCommands(this IServiceCollection services)
         services.AddScoped<IAssertWebAuthnLoginCredentialCommand, AssertWebAuthnLoginCredentialCommand>();
     }
 
+    private static void AddPremiumAccessQueries(this IServiceCollection services)
+    {
+        services.AddScoped<IPremiumAccessQuery, PremiumAccessQuery>();
+    }
+
     private static void AddTwoFactorQueries(this IServiceCollection services)
     {
         services.AddScoped<ITwoFactorIsEnabledQuery, TwoFactorIsEnabledQuery>();

src/Core/Constants.cs

@@ -144,6 +144,7 @@ public static class FeatureFlagKeys
     public const string BlockClaimedDomainAccountCreation = "pm-28297-block-uninvited-claimed-domain-registration";
     public const string PolicyValidatorsRefactor = "pm-26423-refactor-policy-side-effects";
     public const string IncreaseBulkReinviteLimitForCloud = "pm-28251-increase-bulk-reinvite-limit-for-cloud";
+    public const string PremiumAccessCacheCheck = "pm-21411-premium-access-cache-check";
 
     /* Architecture */
     public const string DesktopMigrationMilestone1 = "desktop-ui-migration-milestone-1";

src/Core/Entities/User.cs

@@ -69,6 +69,10 @@ 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.
+    /// </summary>
     public bool Premium { get; set; }
     public DateTime? PremiumExpirationDate { get; set; }
     public DateTime? RenewalReminderDate { get; set; }

src/Core/Repositories/IUserRepository.cs

@@ -23,6 +23,7 @@ public interface IUserRepository : IRepository<User, Guid>
     /// Retrieves the data for the requested user IDs and includes an additional property indicating
     /// whether the user has premium access directly or through an organization.
     /// </summary>
+    [Obsolete("Use IUserService.CanAccessPremiumBulk instead. This method is only used when feature flag 'PremiumAccessCacheCheck' is disabled.")]
     Task<IEnumerable<UserWithCalculatedPremium>> GetManyWithCalculatedPremiumAsync(IEnumerable<Guid> ids);
     /// <summary>
     /// Retrieves the data for the requested user ID and includes additional property indicating
@@ -33,6 +34,7 @@ public interface IUserRepository : IRepository<User, Guid>
     /// </summary>
     /// <param name="userId">The user ID to retrieve data for.</param>
     /// <returns>User data with calculated premium access; null if nothing is found</returns>
+    [Obsolete("Use IUserService.CanAccessPremium instead. This method is only used when feature flag 'PremiumAccessCacheCheck' is disabled.")]
     Task<UserWithCalculatedPremium?> GetCalculatedPremiumAsync(Guid userId);
     /// <summary>
     /// Sets a new user key and updates all encrypted data.

src/Core/Services/IUserService.cs

@@ -64,6 +64,7 @@ Task<UserLicense> GenerateLicenseAsync(User user, SubscriptionInfo subscriptionI
     /// </summary>
     /// <param name="user">user being acted on</param>
     /// <returns>true if they can access premium; false otherwise.</returns>
+    [Obsolete("Use IPremiumAccessQuery.CanAccessPremiumAsync instead. This method is only used when feature flag 'PremiumAccessCacheCheck' is disabled.")]
     Task<bool> CanAccessPremium(User user);
 
     /// <summary>
@@ -74,6 +75,7 @@ Task<UserLicense> GenerateLicenseAsync(User user, SubscriptionInfo subscriptionI
     /// </summary>
     /// <param name="user"></param>
     /// <returns></returns>
+    [Obsolete("Use IPremiumAccessQuery.HasPremiumFromOrganizationAsync instead. This method is only used when feature flag 'PremiumAccessCacheCheck' is disabled.")]
     Task<bool> HasPremiumFromOrganization(User user);
     Task<string> GenerateSignInTokenAsync(User user, string purpose);