Refactor TwoFactorIsEnabledQuery to utilize IFeatureService for premium access checks.

This commit is contained in:
Rui Tome 2025-12-08 16:55:39 +00:00
parent 353a07a040
commit da9f0a1c20
No known key found for this signature in database
GPG Key ID: 526239D96A8EC066
3 changed files with 84 additions and 45 deletions

View File

@ -1,5 +1,4 @@
using Bit.Core.Auth.Models;
using Bit.Core.Entities;
namespace Bit.Core.Auth.UserFeatures.TwoFactorAuth.Interfaces;
@ -23,25 +22,4 @@ public interface ITwoFactorIsEnabledQuery
/// </summary>
/// <param name="user">The user to check.</param>
Task<bool> TwoFactorIsEnabledAsync(ITwoFactorProvidersUser user);
/// <summary>
/// Returns a list of user IDs and whether two factor is enabled for each user.
/// This version uses HasPremiumAccessQuery with cached organization abilities for better performance.
/// </summary>
/// <param name="userIds">The list of user IDs to check.</param>
Task<IEnumerable<(Guid userId, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync(IEnumerable<Guid> userIds);
/// <summary>
/// Returns a list of users and whether two factor is enabled for each user.
/// This version uses HasPremiumAccessQuery with cached organization abilities for better performance.
/// </summary>
/// <param name="users">The list of users to check.</param>
/// <typeparam name="T">The type of user in the list. Must implement <see cref="ITwoFactorProvidersUser"/>.</typeparam>
Task<IEnumerable<(T user, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync<T>(IEnumerable<T> users) where T : ITwoFactorProvidersUser;
/// <summary>
/// Returns whether two factor is enabled for the user. A user is able to have a TwoFactorProvider that is enabled but requires Premium.
/// If the user does not have premium then the TwoFactorProvider is considered _not_ enabled.
/// This version uses HasPremiumAccessQuery with cached organization abilities for better performance.
/// </summary>
/// <param name="user">The user to check.</param>
Task<bool> TwoFactorIsEnabledVNextAsync(User user);
}

View File

@ -7,6 +7,7 @@ using Bit.Core.Auth.UserFeatures.TwoFactorAuth.Interfaces;
using Bit.Core.Billing.Premium.Queries;
using Bit.Core.Entities;
using Bit.Core.Repositories;
using Bit.Core.Services;
namespace Bit.Core.Auth.UserFeatures.TwoFactorAuth;
@ -14,17 +15,25 @@ public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
{
private readonly IUserRepository _userRepository;
private readonly IHasPremiumAccessQuery _hasPremiumAccessQuery;
private readonly IFeatureService _featureService;
public TwoFactorIsEnabledQuery(
IUserRepository userRepository,
IHasPremiumAccessQuery hasPremiumAccessQuery)
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())
{
@ -47,6 +56,11 @@ public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
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)
@ -76,8 +90,25 @@ public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
public async Task<bool> TwoFactorIsEnabledAsync(ITwoFactorProvidersUser user)
{
var userId = user.GetUserId();
if (!userId.HasValue)
if (_featureService.IsEnabled(FeatureFlagKeys.PremiumAccessQuery))
{
var userId = user.GetUserId();
if (!userId.HasValue)
{
return false;
}
var userEntity = user as User ?? await _userRepository.GetByIdAsync(userId.Value);
if (userEntity == null)
{
return false;
}
return await TwoFactorIsEnabledVNextAsync(userEntity);
}
var id = user.GetUserId();
if (!id.HasValue)
{
return false;
}
@ -86,12 +117,12 @@ public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
user.GetTwoFactorProviders(),
async () =>
{
var calcUser = await _userRepository.GetCalculatedPremiumAsync(userId.Value);
var calcUser = await _userRepository.GetCalculatedPremiumAsync(id.Value);
return calcUser?.HasPremiumAccess ?? false;
});
}
public async Task<IEnumerable<(Guid userId, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync(IEnumerable<Guid> userIds)
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())
@ -126,7 +157,7 @@ public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
return result;
}
public async Task<IEnumerable<(T user, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync<T>(IEnumerable<T> users)
private async Task<IEnumerable<(T user, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync<T>(IEnumerable<T> users)
where T : ITwoFactorProvidersUser
{
var userIds = users
@ -156,7 +187,7 @@ public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
return result;
}
public async Task<bool> TwoFactorIsEnabledVNextAsync(User user)
private async Task<bool> TwoFactorIsEnabledVNextAsync(User user)
{
var enabledProviders = GetEnabledTwoFactorProviders(user);

View File

@ -6,6 +6,7 @@ using Bit.Core.Entities;
using Bit.Core.Models.Data;
using Bit.Core.Models.Data.Organizations.OrganizationUsers;
using Bit.Core.Repositories;
using Bit.Core.Services;
using Bit.Core.Utilities;
using Bit.Test.Common.AutoFixture;
using Bit.Test.Common.AutoFixture.Attributes;
@ -408,12 +409,17 @@ public class TwoFactorIsEnabledQueryTests
[Theory]
[BitAutoData((IEnumerable<Guid>)null)]
[BitAutoData([])]
public async Task TwoFactorIsEnabledVNextAsync_NullOrEmptyUserIds_ReturnsEmpty(
public async Task TwoFactorIsEnabledAsync_WhenPremiumAccessQueryEnabled_WithNoUserIds_ReturnsEmpty(
IEnumerable<Guid> userIds,
SutProvider<TwoFactorIsEnabledQuery> sutProvider)
{
// Arrange
sutProvider.GetDependency<IFeatureService>()
.IsEnabled(FeatureFlagKeys.PremiumAccessQuery)
.Returns(true);
// Act
var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(userIds);
var result = await sutProvider.Sut.TwoFactorIsEnabledAsync(userIds);
// Assert
Assert.Empty(result);
@ -422,7 +428,7 @@ public class TwoFactorIsEnabledQueryTests
[Theory]
[BitAutoData(TwoFactorProviderType.Duo)]
[BitAutoData(TwoFactorProviderType.YubiKey)]
public async Task TwoFactorIsEnabledVNextAsync_MixedScenarios_ReturnsCorrectResults(
public async Task TwoFactorIsEnabledAsync_WhenPremiumAccessQueryEnabled_WithMixedScenarios_ReturnsCorrectResults(
TwoFactorProviderType premiumProviderType,
SutProvider<TwoFactorIsEnabledQuery> sutProvider,
User user1,
@ -430,6 +436,10 @@ public class TwoFactorIsEnabledQueryTests
User user3)
{
// Arrange
sutProvider.GetDependency<IFeatureService>()
.IsEnabled(FeatureFlagKeys.PremiumAccessQuery)
.Returns(true);
var users = new List<User> { user1, user2, user3 };
var userIds = users.Select(u => u.Id).ToList();
@ -467,7 +477,7 @@ public class TwoFactorIsEnabledQueryTests
.Returns(premiumStatus);
// Act
var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(userIds);
var result = await sutProvider.Sut.TwoFactorIsEnabledAsync(userIds);
// Assert
Assert.Contains(result, res => res.userId == user1.Id && res.twoFactorIsEnabled == true); // Non-premium provider
@ -478,7 +488,7 @@ public class TwoFactorIsEnabledQueryTests
[Theory]
[BitAutoData(TwoFactorProviderType.Duo)]
[BitAutoData(TwoFactorProviderType.YubiKey)]
public async Task TwoFactorIsEnabledVNextAsync_OnlyChecksPremiumAccessForUsersWhoNeedIt(
public async Task TwoFactorIsEnabledAsync_WhenPremiumAccessQueryEnabled_OnlyChecksPremiumAccessForUsersWhoNeedIt(
TwoFactorProviderType premiumProviderType,
SutProvider<TwoFactorIsEnabledQuery> sutProvider,
User user1,
@ -486,6 +496,10 @@ public class TwoFactorIsEnabledQueryTests
User user3)
{
// Arrange
sutProvider.GetDependency<IFeatureService>()
.IsEnabled(FeatureFlagKeys.PremiumAccessQuery)
.Returns(true);
var users = new List<User> { user1, user2, user3 };
var userIds = users.Select(u => u.Id).ToList();
@ -521,7 +535,7 @@ public class TwoFactorIsEnabledQueryTests
.Returns(premiumStatus);
// Act
var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(userIds);
var result = await sutProvider.Sut.TwoFactorIsEnabledAsync(userIds);
// Assert - Verify optimization: premium checked ONLY for users 2 and 3 (not user 1)
await sutProvider.GetDependency<IHasPremiumAccessQuery>()
@ -532,18 +546,22 @@ public class TwoFactorIsEnabledQueryTests
[Theory]
[BitAutoData]
public async Task TwoFactorIsEnabledVNextAsync_Generic_WithNoUserIds_ReturnsAllTwoFactorDisabled(
public async Task TwoFactorIsEnabledAsync_WhenPremiumAccessQueryEnabled_WithNoUserIds_ReturnsAllTwoFactorDisabled(
SutProvider<TwoFactorIsEnabledQuery> sutProvider,
List<OrganizationUserUserDetails> users)
{
// Arrange
sutProvider.GetDependency<IFeatureService>()
.IsEnabled(FeatureFlagKeys.PremiumAccessQuery)
.Returns(true);
foreach (var user in users)
{
user.UserId = null;
}
// Act
var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(users);
var result = await sutProvider.Sut.TwoFactorIsEnabledAsync(users);
// Assert
foreach (var user in users)
@ -559,15 +577,19 @@ public class TwoFactorIsEnabledQueryTests
[Theory]
[BitAutoData(TwoFactorProviderType.Authenticator, true)] // Non-premium provider
[BitAutoData(TwoFactorProviderType.Duo, true)] // Premium provider with premium access
[BitAutoData(TwoFactorProviderType.Duo, true)] // Premium provider with premium access
[BitAutoData(TwoFactorProviderType.YubiKey, false)] // Premium provider without premium access
public async Task TwoFactorIsEnabledVNextAsync_SingleUser_VariousScenarios(
public async Task TwoFactorIsEnabledAsync_WhenPremiumAccessQueryEnabled_SingleUser_VariousScenarios(
TwoFactorProviderType providerType,
bool hasPremiumAccess,
SutProvider<TwoFactorIsEnabledQuery> sutProvider,
User user)
{
// Arrange
sutProvider.GetDependency<IFeatureService>()
.IsEnabled(FeatureFlagKeys.PremiumAccessQuery)
.Returns(true);
user.SetTwoFactorProviders(new Dictionary<TwoFactorProviderType, TwoFactorProvider>
{
{ providerType, new TwoFactorProvider { Enabled = true } }
@ -578,7 +600,7 @@ public class TwoFactorIsEnabledQueryTests
.Returns(hasPremiumAccess);
// Act
var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
var result = await sutProvider.Sut.TwoFactorIsEnabledAsync(user);
// Assert
var requiresPremium = TwoFactorProvider.RequiresPremium(providerType);
@ -588,18 +610,22 @@ public class TwoFactorIsEnabledQueryTests
[Theory]
[BitAutoData]
public async Task TwoFactorIsEnabledVNextAsync_SingleUser_NoEnabledProviders_ReturnsFalse(
public async Task TwoFactorIsEnabledAsync_WhenPremiumAccessQueryEnabled_WithNoEnabledProviders_ReturnsFalse(
SutProvider<TwoFactorIsEnabledQuery> sutProvider,
User user)
{
// Arrange - test both disabled and null providers
// Arrange
sutProvider.GetDependency<IFeatureService>()
.IsEnabled(FeatureFlagKeys.PremiumAccessQuery)
.Returns(true);
user.SetTwoFactorProviders(new Dictionary<TwoFactorProviderType, TwoFactorProvider>
{
{ TwoFactorProviderType.Email, new TwoFactorProvider { Enabled = false } }
});
// Act
var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
var result = await sutProvider.Sut.TwoFactorIsEnabledAsync(user);
// Assert
Assert.False(result);
@ -607,15 +633,19 @@ public class TwoFactorIsEnabledQueryTests
[Theory]
[BitAutoData]
public async Task TwoFactorIsEnabledVNextAsync_SingleUser_NullProviders_ReturnsFalse(
public async Task TwoFactorIsEnabledAsync_WhenPremiumAccessQueryEnabled_WithNullProviders_ReturnsFalse(
SutProvider<TwoFactorIsEnabledQuery> sutProvider,
User user)
{
// Arrange
sutProvider.GetDependency<IFeatureService>()
.IsEnabled(FeatureFlagKeys.PremiumAccessQuery)
.Returns(true);
user.TwoFactorProviders = null;
// Act
var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
var result = await sutProvider.Sut.TwoFactorIsEnabledAsync(user);
// Assert
Assert.False(result);