Improve generic support for AddEntityFrameworkStores

aspnet · Nov 17, 2016 · 13ae7b2 · 13ae7b2
1 parent 3be87ce
commit 13ae7b2
Show file tree

Hide file tree

Showing 7 changed files with 121 additions and 15 deletions.
diff --git a/...osoft.AspNetCore.Identity.EntityFrameworkCore/IdentityEntityFrameworkBuilderExtensions.cs b/...osoft.AspNetCore.Identity.EntityFrameworkCore/IdentityEntityFrameworkBuilderExtensions.cs
@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Reflection;
 using Microsoft.AspNetCore.Identity;
 using Microsoft.AspNetCore.Identity.EntityFrameworkCore;
 using Microsoft.EntityFrameworkCore;
@@ -23,7 +24,8 @@ public static class IdentityEntityFrameworkBuilderExtensions
         public static IdentityBuilder AddEntityFrameworkStores<TContext>(this IdentityBuilder builder)
             where TContext : DbContext
         {
-            builder.Services.TryAdd(GetDefaultServices(builder.UserType, builder.RoleType, typeof(TContext)));
+            var keyType = InferKeyType(typeof(TContext));
+            AddStores(builder.Services, builder.UserType, builder.RoleType, typeof(TContext), keyType);
             return builder;
         }
 
@@ -38,26 +40,44 @@ public static IdentityBuilder AddEntityFrameworkStores<TContext, TKey>(this Iden
             where TContext : DbContext
             where TKey : IEquatable<TKey>
         {
-            builder.Services.TryAdd(GetDefaultServices(builder.UserType, builder.RoleType, typeof(TContext), typeof(TKey)));
+            AddStores(builder.Services, builder.UserType, builder.RoleType, typeof(TContext), typeof(TKey));
             return builder;
         }
 
-        private static IServiceCollection GetDefaultServices(Type userType, Type roleType, Type contextType, Type keyType = null)
+        private static void AddStores(IServiceCollection services, Type userType, Type roleType, Type contextType, Type keyType)
         {
-            Type userStoreType;
-            Type roleStoreType;
-            keyType = keyType ?? typeof(string);
-            userStoreType = typeof(UserStore<,,,>).MakeGenericType(userType, roleType, contextType, keyType);
-            roleStoreType = typeof(RoleStore<,,>).MakeGenericType(roleType, contextType, keyType);
-
-            var services = new ServiceCollection();
-            services.AddScoped(
+            var identityUserType = typeof(IdentityUser<>).MakeGenericType(keyType);
+            if (!identityUserType.GetTypeInfo().IsAssignableFrom(userType.GetTypeInfo()))
+            {
+                throw new InvalidOperationException(Resources.NotIdentityUser);
+            }
+            var identityRoleType = typeof(IdentityRole<>).MakeGenericType(keyType);
+            if (!identityRoleType.GetTypeInfo().IsAssignableFrom(roleType.GetTypeInfo()))
+            {
+                throw new InvalidOperationException(Resources.NotIdentityRole);
+            }
+            services.TryAddScoped(
                 typeof(IUserStore<>).MakeGenericType(userType),
-                userStoreType);
-            services.AddScoped(
+                typeof(UserStore<,,,>).MakeGenericType(userType, roleType, contextType, keyType));
+            services.TryAddScoped(
                 typeof(IRoleStore<>).MakeGenericType(roleType),
-                roleStoreType);
-            return services;
+                typeof(RoleStore<,,>).MakeGenericType(roleType, contextType, keyType));
+        }
+
+        private static Type InferKeyType(Type contextType)
+        {
+            var type = contextType.GetTypeInfo();
+            while (type.BaseType != null)
+            {
+                type = type.BaseType.GetTypeInfo();
+                var genericType = type.IsGenericType ? type.GetGenericTypeDefinition() : null;
+                if (genericType != null && genericType == typeof(IdentityDbContext<,,>))
+                {
+                    return type.GenericTypeArguments[2];
+                }
+            }
+            // Default is string
+            return typeof(string);
         }
     }
 }
diff --git a/src/Microsoft.AspNetCore.Identity.EntityFrameworkCore/Properties/Resources.Designer.cs b/src/Microsoft.AspNetCore.Identity.EntityFrameworkCore/Properties/Resources.Designer.cs
diff --git a/src/Microsoft.AspNetCore.Identity.EntityFrameworkCore/Resources.resx b/src/Microsoft.AspNetCore.Identity.EntityFrameworkCore/Resources.resx
@@ -117,6 +117,14 @@
   <resheader name="writer">
     <value>System.Resources.ResXResourceWriter, System.Windows.Forms, Version=4.0.0.0, Culture=neutral, PublicKeyToken=b77a5c561934e089</value>
   </resheader>
+  <data name="NotIdentityRole" xml:space="preserve">
+    <value>AddEntityFrameworkStores can only be called with a role that derives from IdentityRole&lt;TKey&gt;. If you are specifying more generic arguments, use AddRoleStore&lt;TStore&gt;() where TStore is your custom IRoleStore that uses your generics instead.</value>
+    <comment>error when the role does not derive from IdentityRole</comment>
+  </data>
+  <data name="NotIdentityUser" xml:space="preserve">
+    <value>AddEntityFrameworkStores can only be called with a user that derives from IdentityUser&lt;TKey&gt;. If you are specifying more generic arguments, use IdentityBuilder.AddUserStore&lt;TStore&gt;() where TStore is your custom IUserStore that uses your generics instead.</value>
+    <comment>error when the user does not derive from IdentityUser</comment>
+  </data>
   <data name="RoleNotFound" xml:space="preserve">
     <value>Role {0} does not exist.</value>
     <comment>error when a role does not exist</comment>

diff --git a/test/Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test/UserStoreGuidKeyTest.cs b/test/Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test/UserStoreGuidKeyTest.cs
@@ -51,5 +51,13 @@ protected override void AddRoleStore(IServiceCollection services, object context
         {
             services.AddSingleton<IRoleStore<GuidRole>>(new ApplicationRoleStore((TestDbContext)context));
         }
+
+        [Fact]
+        public void AddEntityFrameworkStoresCanInferKey()
+        {
+            var services = new ServiceCollection();
+            // This used to throw
+            var builder = services.AddIdentity<GuidUser, GuidRole>().AddEntityFrameworkStores<TestDbContext>();
+        }
     }
 }
diff --git a/test/Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test/UserStoreIntKeyTest.cs b/test/Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test/UserStoreIntKeyTest.cs
@@ -2,6 +2,8 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using Microsoft.Extensions.DependencyInjection;
+using Xunit;
 
 namespace Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test
 {
@@ -27,5 +29,13 @@ public UserStoreIntTest(ScratchDatabaseFixture fixture)
             : base(fixture)
         {
         }
+
+        [Fact]
+        public void AddEntityFrameworkStoresCanInferKey()
+        {
+            var services = new ServiceCollection();
+            // This used to throw
+            var builder = services.AddIdentity<IntUser, IntRole>().AddEntityFrameworkStores<TestDbContext>();
+        }
     }
 }
diff --git a/test/Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test/UserStoreStringKeyTest.cs b/test/Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test/UserStoreStringKeyTest.cs
@@ -2,6 +2,8 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using Microsoft.Extensions.DependencyInjection;
+using Xunit;
 
 namespace Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test
 {
@@ -27,7 +29,15 @@ public class UserStoreStringKeyTest : SqlStoreTestBase<StringUser, StringRole, s
     {
         public UserStoreStringKeyTest(ScratchDatabaseFixture fixture)
             : base(fixture)
+        { }
+
+        [Fact]
+        public void AddEntityFrameworkStoresCanInferKey()
         {
+            var services = new ServiceCollection();
+            // This used to throw
+            var builder = services.AddIdentity<StringUser, StringRole>().AddEntityFrameworkStores<TestDbContext>();
         }
+
     }
 }
diff --git a/test/Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test/UserStoreWithGenericsTest.cs b/test/Microsoft.AspNetCore.Identity.EntityFrameworkCore.Test/UserStoreWithGenericsTest.cs
@@ -83,6 +83,24 @@ protected override void SetUserPasswordHash(IdentityUserWithGenerics user, strin
 
         protected override Expression<Func<MyIdentityRole, bool>> RoleNameStartsWithPredicate(string roleName) => r => r.Name.StartsWith(roleName);
 
+        [Fact]
+        public void AddEntityFrameworkStoresWithInvalidUserThrows()
+        {
+            var services = new ServiceCollection();
+            var builder = services.AddIdentity<IdentityUserWithGenerics, IdentityRole>();
+            var e = Assert.Throws<InvalidOperationException>(() => builder.AddEntityFrameworkStores<ContextWithGenerics>());
+            Assert.Contains("AddEntityFrameworkStores", e.Message);
+        }
+
+        [Fact]
+        public void AddEntityFrameworkStoresWithInvalidRoleThrows()
+        {
+            var services = new ServiceCollection();
+            var builder = services.AddIdentity<IdentityUser, MyIdentityRole>();
+            var e = Assert.Throws<InvalidOperationException>(() => builder.AddEntityFrameworkStores<ContextWithGenerics>());
+            Assert.Contains("AddEntityFrameworkStores", e.Message);
+        }
+
         [Fact]
         public async Task CanAddRemoveUserClaimWithIssuer()
         {