pm-24210-v3 (#6148)

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> 
-    /// <param name="user">The authenticated user.</param> 
+    /// </summary>
+    /// <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;
-                await _authRequestRepository.ReplaceAsync(authRequest);
+                // We save the validated auth request so that we can set it's authentication date
+                // 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;