From d18ae5cad12700e3b18fab8feeb36eb486fc9db4 Mon Sep 17 00:00:00 2001
From: Hao Kung <haok@microsoft.com>
Date: Wed, 28 Sep 2016 13:33:17 -0700
Subject: [PATCH] Add ProviderInstance for use in 2FA TokenProvider

---
 .../TokenProviderDescriptor.cs                |  5 ++
 .../UserManager.cs                            |  5 +-
 .../UserManagerTest.cs                        | 46 +++++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/src/Microsoft.AspNetCore.Identity/TokenProviderDescriptor.cs b/src/Microsoft.AspNetCore.Identity/TokenProviderDescriptor.cs
index 993ecfbed..95c2920b3 100644
--- a/src/Microsoft.AspNetCore.Identity/TokenProviderDescriptor.cs
+++ b/src/Microsoft.AspNetCore.Identity/TokenProviderDescriptor.cs
@@ -23,5 +23,10 @@ public TokenProviderDescriptor(Type type)
         /// The type that will be used for this token provider.
         /// </summary>
         public Type ProviderType { get; }
+
+        /// <summary>
+        /// If specified, the instance to be used for the token provider.
+        /// </summary>
+        public object ProviderInstance { get; set; }
     }
 }
\ No newline at end of file
diff --git a/src/Microsoft.AspNetCore.Identity/UserManager.cs b/src/Microsoft.AspNetCore.Identity/UserManager.cs
index 30881435c..d3eff1b91 100644
--- a/src/Microsoft.AspNetCore.Identity/UserManager.cs
+++ b/src/Microsoft.AspNetCore.Identity/UserManager.cs
@@ -98,7 +98,10 @@ public UserManager(IUserStore<TUser> store,
                 _context = services.GetService<IHttpContextAccessor>()?.HttpContext;
                 foreach (var providerName in Options.Tokens.ProviderMap.Keys)
                 {
-                    var provider = services.GetRequiredService(Options.Tokens.ProviderMap[providerName].ProviderType) as IUserTwoFactorTokenProvider<TUser>;
+                    var description = Options.Tokens.ProviderMap[providerName];
+                    
+                    var provider = (description.ProviderInstance ?? services.GetRequiredService(description.ProviderType)) 
+                        as IUserTwoFactorTokenProvider<TUser>;
                     if (provider != null)
                     {
                         RegisterTokenProvider(providerName, provider);
diff --git a/test/Microsoft.AspNetCore.Identity.Test/UserManagerTest.cs b/test/Microsoft.AspNetCore.Identity.Test/UserManagerTest.cs
index 5be34eb51..664bc14f9 100644
--- a/test/Microsoft.AspNetCore.Identity.Test/UserManagerTest.cs
+++ b/test/Microsoft.AspNetCore.Identity.Test/UserManagerTest.cs
@@ -623,6 +623,52 @@ public async Task ClaimMethodsFailWhenStoreNotImplemented()
             await Assert.ThrowsAsync<NotSupportedException>(async () => await manager.GetClaimsAsync(null));
         }
 
+        private class ATokenProvider : IUserTwoFactorTokenProvider<TestUser>
+        {
+            public Task<bool> CanGenerateTwoFactorTokenAsync(UserManager<TestUser> manager, TestUser user)
+            {
+                throw new NotImplementedException();
+            }
+
+            public Task<string> GenerateAsync(string purpose, UserManager<TestUser> manager, TestUser user)
+            {
+                throw new NotImplementedException();
+            }
+
+            public Task<bool> ValidateAsync(string purpose, string token, UserManager<TestUser> manager, TestUser user)
+            {
+                throw new NotImplementedException();
+            }
+        }
+
+        [Fact]
+        public void UserManagerWillUseTokenProviderInstance()
+        {
+            var services = new ServiceCollection();
+            var provider = new ATokenProvider();
+            services.AddLogging()
+                .AddIdentity<TestUser, TestRole>(o => o.Tokens.ProviderMap.Add("A", new TokenProviderDescriptor(typeof(ATokenProvider))
+            {
+                ProviderInstance = provider
+            })).AddUserStore<NoopUserStore>();
+            var manager = services.BuildServiceProvider().GetService<UserManager<TestUser>>();
+            Assert.ThrowsAsync<NotImplementedException>(() => manager.GenerateUserTokenAsync(new TestUser(), "A", "purpose"));
+        }
+
+        [Fact]
+        public void UserManagerWillUseTokenProviderInstanceOverDefaults()
+        {
+            var services = new ServiceCollection();
+            var provider = new ATokenProvider();
+            services.AddLogging()
+                .AddIdentity<TestUser, TestRole>(o => o.Tokens.ProviderMap.Add(TokenOptions.DefaultProvider, new TokenProviderDescriptor(typeof(ATokenProvider))
+                {
+                    ProviderInstance = provider
+                })).AddUserStore<NoopUserStore>().AddDefaultTokenProviders();
+            var manager = services.BuildServiceProvider().GetService<UserManager<TestUser>>();
+            Assert.ThrowsAsync<NotImplementedException>(() => manager.GenerateUserTokenAsync(new TestUser(), TokenOptions.DefaultProvider, "purpose"));
+        }
+
         [Fact]
         public async Task TwoFactorStoreMethodsFailWhenStoreNotImplemented()
         {