From dc271d807b18b7134ac6a029da9751ce9f363e82 Mon Sep 17 00:00:00 2001
From: maliming <malimings@gmail.com>
Date: Fri, 28 May 2021 14:30:21 +0800
Subject: [PATCH] Return 2FA info in AbpResourceOwnerPasswordValidator.

---
 .../AbpResourceOwnerPasswordValidator.cs      | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/modules/identityserver/src/Volo.Abp.IdentityServer.Domain/Volo/Abp/IdentityServer/AspNetIdentity/AbpResourceOwnerPasswordValidator.cs b/modules/identityserver/src/Volo.Abp.IdentityServer.Domain/Volo/Abp/IdentityServer/AspNetIdentity/AbpResourceOwnerPasswordValidator.cs
index e024976548..acbbaa1f9b 100644
--- a/modules/identityserver/src/Volo.Abp.IdentityServer.Domain/Volo/Abp/IdentityServer/AspNetIdentity/AbpResourceOwnerPasswordValidator.cs
+++ b/modules/identityserver/src/Volo.Abp.IdentityServer.Domain/Volo/Abp/IdentityServer/AspNetIdentity/AbpResourceOwnerPasswordValidator.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Resources;
 using System.Security.Claims;
 using System.Threading.Tasks;
 using IdentityModel;
@@ -130,6 +131,30 @@ namespace Volo.Abp.IdentityServer.AspNetIdentity
                 var result = await SignInManager.CheckPasswordSignInAsync(user, context.Password, true);
                 if (result.Succeeded)
                 {
+                    if (await IsTfaEnabled(user))
+                    {
+                        var twoFactorProvider = context.Request?.Raw?["TwoFactorProvider"];
+                        var twoFactorCode = context.Request?.Raw?["TwoFactorCode"];
+                        if (!twoFactorProvider.IsNullOrWhiteSpace() && !twoFactorCode.IsNullOrWhiteSpace())
+                        {
+                            if (await UserManager.VerifyTwoFactorTokenAsync(user, twoFactorProvider, twoFactorCode))
+                            {
+                                await SetSuccessResultAsync();
+                                return;
+                            }
+                        }
+
+                        Logger.LogInformation("Authentication failed for username: {username}, reason: RequiresTwoFactor", context.UserName);
+                        var twoFactorToken = await UserManager.GenerateUserTokenAsync(user, TokenOptions.DefaultProvider, nameof(SignInResult.RequiresTwoFactor));
+                        context.Result = new GrantValidationResult(TokenRequestErrors.InvalidGrant, nameof(SignInResult.RequiresTwoFactor),
+                            new Dictionary<string, object>()
+                            {
+                                {"userId", user.Id},
+                                {"twoFactorToken", twoFactorToken}
+                            });
+                        return;
+                    }
+
                     await SetSuccessResultAsync();
                     return;
                 }
@@ -196,6 +221,11 @@ namespace Volo.Abp.IdentityServer.AspNetIdentity
             context.UserName = userByEmail.UserName;
         }
 
+        protected virtual async Task<bool> IsTfaEnabled(IdentityUser user)
+            => UserManager.SupportsUserTwoFactor &&
+               await UserManager.GetTwoFactorEnabledAsync(user) &&
+               (await UserManager.GetValidTwoFactorProvidersAsync(user)).Count > 0;
+
         protected virtual Task AddCustomClaimsAsync(List<Claim> customClaims, IdentityUser user, ResourceOwnerPasswordValidationContext context)
         {
             if (user.TenantId.HasValue)