pm-24210-v3 (#6148) + test merge conflict fix in base request validator tests

src/Identity/IdentityServer/CustomValidatorRequestContext.cs

@@ -1,6 +1,7 @@
 // FIXME: Update this file to be null safe and then delete the line below
 #nullable disable
 
+using Bit.Core.Auth.Entities;
 using Bit.Core.Entities;
 using Duende.IdentityServer.Validation;
 
@@ -41,4 +42,10 @@ public class CustomValidatorRequestContext
     /// This will be null if the authentication request is successful.
     /// </summary>
     public Dictionary<string, object> CustomResponse { get; set; }
+
+    /// <summary>
+    /// A validated auth request
+    /// <see cref="AuthRequest.IsValidForAuthentication"/>
+    /// </summary>
+    public AuthRequest ValidatedAuthRequest { get; set; }
 }

src/Identity/IdentityServer/RequestValidators/BaseRequestValidator.cs

@@ -35,6 +35,7 @@ public abstract class BaseRequestValidator<T> where T : class
     private readonly ILogger _logger;
     private readonly GlobalSettings _globalSettings;
     private readonly IUserRepository _userRepository;
+    private readonly IAuthRequestRepository _authRequestRepository;
 
     protected ICurrentContext CurrentContext { get; }
     protected IPolicyService PolicyService { get; }
@@ -59,7 +60,9 @@ public abstract class BaseRequestValidator<T> where T : class
         IFeatureService featureService,
         ISsoConfigRepository ssoConfigRepository,
         IUserDecryptionOptionsBuilder userDecryptionOptionsBuilder,
-        IPolicyRequirementQuery policyRequirementQuery)
+        IPolicyRequirementQuery policyRequirementQuery,
+        IAuthRequestRepository authRequestRepository
+        )
     {
         _userManager = userManager;
         _userService = userService;
@@ -76,6 +79,7 @@ public abstract class BaseRequestValidator<T> where T : class
         SsoConfigRepository = ssoConfigRepository;
         UserDecryptionOptionsBuilder = userDecryptionOptionsBuilder;
         PolicyRequirementQuery = policyRequirementQuery;
+        _authRequestRepository = authRequestRepository;
     }
 
     protected async Task ValidateAsync(T context, ValidatedTokenRequest request,
@@ -190,6 +194,14 @@ public abstract class BaseRequestValidator<T> where T : class
             return;
         }
 
+        // TODO: PM-24324 - This should be its own validator at some point.
+        // 6. Auth request handling
+        if (validatorContext.ValidatedAuthRequest != null)
+        {
+            validatorContext.ValidatedAuthRequest.AuthenticationDate = DateTime.UtcNow;
+            await _authRequestRepository.ReplaceAsync(validatorContext.ValidatedAuthRequest);
+        }
+
         await BuildSuccessResultAsync(user, context, validatorContext.Device, returnRememberMeToken);
     }
 
@@ -404,8 +416,8 @@ public abstract class BaseRequestValidator<T> where T : class
     /// <summary>
     /// Builds the custom response that will be sent to the client upon successful authentication, which
     /// includes the information needed for the client to initialize the user's account in state.
-    /// </summary> 
+    /// </summary>
-    /// <param name="user">The authenticated user.</param> 
+    /// <param name="user">The authenticated user.</param>
     /// <param name="context">The current request context.</param>
     /// <param name="device">The device used for authentication.</param>
     /// <param name="sendRememberToken">Whether to send a 2FA remember token.</param>

src/Identity/IdentityServer/RequestValidators/CustomTokenRequestValidator.cs

@@ -45,7 +45,8 @@ public class CustomTokenRequestValidator : BaseRequestValidator<CustomTokenReque
         ISsoConfigRepository ssoConfigRepository,
         IUserDecryptionOptionsBuilder userDecryptionOptionsBuilder,
         IUpdateInstallationCommand updateInstallationCommand,
-        IPolicyRequirementQuery policyRequirementQuery)
+        IPolicyRequirementQuery policyRequirementQuery,
+        IAuthRequestRepository authRequestRepository)
         : base(
             userManager,
             userService,
@@ -61,7 +62,8 @@ public class CustomTokenRequestValidator : BaseRequestValidator<CustomTokenReque
             featureService,
             ssoConfigRepository,
             userDecryptionOptionsBuilder,
-            policyRequirementQuery)
+            policyRequirementQuery,
+            authRequestRepository)
     {
         _userManager = userManager;
         _updateInstallationCommand = updateInstallationCommand;

src/Identity/IdentityServer/RequestValidators/ResourceOwnerPasswordValidator.cs

@@ -56,7 +56,8 @@ public class ResourceOwnerPasswordValidator : BaseRequestValidator<ResourceOwner
             featureService,
             ssoConfigRepository,
             userDecryptionOptionsBuilder,
-            policyRequirementQuery)
+            policyRequirementQuery,
+            authRequestRepository)
     {
         _userManager = userManager;
         _currentContext = currentContext;
@@ -108,8 +109,11 @@ public class ResourceOwnerPasswordValidator : BaseRequestValidator<ResourceOwner
             // Auth request is non-null so validate it
             if (authRequest.IsValidForAuthentication(validatorContext.User.Id, context.Password))
             {
-                authRequest.AuthenticationDate = DateTime.UtcNow;
+                // We save the validated auth request so that we can set it's authentication date
-                await _authRequestRepository.ReplaceAsync(authRequest);
+                // later on only upon successful authentication.
+                // For example, 2FA requires a resubmission so we can't mark the auth request
+                // as authenticated here.
+                validatorContext.ValidatedAuthRequest = authRequest;
                 return true;
             }
 

src/Identity/IdentityServer/RequestValidators/WebAuthnGrantValidator.cs

@@ -48,7 +48,8 @@ public class WebAuthnGrantValidator : BaseRequestValidator<ExtensionGrantValidat
         IFeatureService featureService,
         IUserDecryptionOptionsBuilder userDecryptionOptionsBuilder,
         IAssertWebAuthnLoginCredentialCommand assertWebAuthnLoginCredentialCommand,
-        IPolicyRequirementQuery policyRequirementQuery)
+        IPolicyRequirementQuery policyRequirementQuery,
+        IAuthRequestRepository authRequestRepository)
         : base(
             userManager,
             userService,
@@ -64,7 +65,8 @@ public class WebAuthnGrantValidator : BaseRequestValidator<ExtensionGrantValidat
             featureService,
             ssoConfigRepository,
             userDecryptionOptionsBuilder,
-            policyRequirementQuery)
+            policyRequirementQuery,
+            authRequestRepository)
     {
         _assertionOptionsDataProtector = assertionOptionsDataProtector;
         _assertWebAuthnLoginCredentialCommand = assertWebAuthnLoginCredentialCommand;

test/Identity.Test/IdentityServer/BaseRequestValidatorTests.cs

@@ -3,6 +3,8 @@ using Bit.Core.AdminConsole.Entities;
 using Bit.Core.AdminConsole.Enums;
 using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
 using Bit.Core.AdminConsole.Services;
+using Bit.Core.Auth.Entities;
+using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Repositories;
 using Bit.Core.Context;
 using Bit.Core.Entities;
@@ -42,6 +44,7 @@ public class BaseRequestValidatorTests
     private readonly ISsoConfigRepository _ssoConfigRepository;
     private readonly IUserDecryptionOptionsBuilder _userDecryptionOptionsBuilder;
     private readonly IPolicyRequirementQuery _policyRequirementQuery;
+    private readonly IAuthRequestRepository _authRequestRepository;
 
     private readonly BaseRequestValidatorTestWrapper _sut;
 
@@ -62,6 +65,7 @@ public class BaseRequestValidatorTests
         _ssoConfigRepository = Substitute.For<ISsoConfigRepository>();
         _userDecryptionOptionsBuilder = Substitute.For<IUserDecryptionOptionsBuilder>();
         _policyRequirementQuery = Substitute.For<IPolicyRequirementQuery>();
+        _authRequestRepository = Substitute.For<IAuthRequestRepository>();
 
         _sut = new BaseRequestValidatorTestWrapper(
             _userManager,
@@ -78,7 +82,8 @@ public class BaseRequestValidatorTests
             _featureService,
             _ssoConfigRepository,
             _userDecryptionOptionsBuilder,
-            _policyRequirementQuery);
+            _policyRequirementQuery,
+            _authRequestRepository);
     }
 
     /* Logic path
@@ -175,6 +180,99 @@ public class BaseRequestValidatorTests
         Assert.False(context.GrantResult.IsError);
     }
 
+    [Theory, BitAutoData]
+    public async Task ValidateAsync_ValidatedAuthRequest_ConsumedOnSuccess(
+        [AuthFixtures.ValidatedTokenRequest] ValidatedTokenRequest tokenRequest,
+        CustomValidatorRequestContext requestContext,
+        GrantValidationResult grantResult)
+    {
+        // Arrange
+        var context = CreateContext(tokenRequest, requestContext, grantResult);
+        // 1 -> to pass
+        _sut.isValid = true;
+
+        var authRequest = new AuthRequest
+        {
+            Type = AuthRequestType.AuthenticateAndUnlock,
+            RequestDeviceIdentifier = "",
+            RequestIpAddress = "1.1.1.1",
+            AccessCode = "password",
+            PublicKey = "test_public_key",
+            CreationDate = DateTime.UtcNow.AddMinutes(-5),
+            ResponseDate = DateTime.UtcNow.AddMinutes(-2),
+            Approved = true,
+            AuthenticationDate = null, // unused
+            UserId = requestContext.User.Id,
+        };
+        requestContext.ValidatedAuthRequest = authRequest;
+
+        // 2 -> will result to false with no extra configuration
+        // 3 -> set two factor to be false
+        _twoFactorAuthenticationValidator
+            .RequiresTwoFactorAsync(Arg.Any<User>(), tokenRequest)
+            .Returns(Task.FromResult(new Tuple<bool, Organization>(false, null)));
+
+        // 4 -> set up device validator to pass
+        _deviceValidator.ValidateRequestDeviceAsync(Arg.Any<ValidatedTokenRequest>(), Arg.Any<CustomValidatorRequestContext>())
+            .Returns(Task.FromResult(true));
+
+        // 5 -> not legacy user
+        _userService.IsLegacyUser(Arg.Any<string>())
+            .Returns(false);
+
+        // Act
+        await _sut.ValidateAsync(context);
+
+        // Assert
+        Assert.False(context.GrantResult.IsError);
+
+        // Check that the auth request was consumed
+        await _authRequestRepository.Received(1).ReplaceAsync(Arg.Is<AuthRequest>(ar =>
+            ar.AuthenticationDate.HasValue));
+    }
+
+    [Theory, BitAutoData]
+    public async Task ValidateAsync_ValidatedAuthRequest_NotConsumed_When2faRequired(
+        [AuthFixtures.ValidatedTokenRequest] ValidatedTokenRequest tokenRequest,
+        CustomValidatorRequestContext requestContext,
+        GrantValidationResult grantResult)
+    {
+        // Arrange
+        var context = CreateContext(tokenRequest, requestContext, grantResult);
+        // 1 -> to pass
+        _sut.isValid = true;
+
+        var authRequest = new AuthRequest
+        {
+            Type = AuthRequestType.AuthenticateAndUnlock,
+            RequestDeviceIdentifier = "",
+            RequestIpAddress = "1.1.1.1",
+            AccessCode = "password",
+            PublicKey = "test_public_key",
+            CreationDate = DateTime.UtcNow.AddMinutes(-5),
+            ResponseDate = DateTime.UtcNow.AddMinutes(-2),
+            Approved = true,
+            AuthenticationDate = null, // unused
+            UserId = requestContext.User.Id,
+        };
+        requestContext.ValidatedAuthRequest = authRequest;
+
+        // 2 -> will result to false with no extra configuration
+        // 3 -> set two factor to be required
+        _twoFactorAuthenticationValidator
+            .RequiresTwoFactorAsync(Arg.Any<User>(), tokenRequest)
+            .Returns(Task.FromResult(new Tuple<bool, Organization>(true, null)));
+
+        // Act
+        await _sut.ValidateAsync(context);
+
+        // Assert we errored for 2fa requirement
+        Assert.True(context.GrantResult.IsError);
+
+        // Assert that the auth request was NOT consumed
+        await _authRequestRepository.DidNotReceive().ReplaceAsync(Arg.Any<AuthRequest>());
+    }
+
     // Test grantTypes that require SSO when a user is in an organization that requires it
     [Theory]
     [BitAutoData("password")]

test/Identity.Test/Wrappers/BaseRequestValidatorTestWrapper.cs

@@ -62,7 +62,8 @@ IBaseRequestValidatorTestWrapper
         IFeatureService featureService,
         ISsoConfigRepository ssoConfigRepository,
         IUserDecryptionOptionsBuilder userDecryptionOptionsBuilder,
-        IPolicyRequirementQuery policyRequirementQuery) :
+        IPolicyRequirementQuery policyRequirementQuery,
+        IAuthRequestRepository authRequestRepository) :
          base(
             userManager,
             userService,
@@ -78,7 +79,8 @@ IBaseRequestValidatorTestWrapper
             featureService,
             ssoConfigRepository,
             userDecryptionOptionsBuilder,
-            policyRequirementQuery)
+            policyRequirementQuery,
+            authRequestRepository)
     {
     }