Fix FromKeyedServicesAttribute and FromServicesAttribute to support derived types across all generators (#63114)

* Initial plan

* Fix FromKeyedServicesAttribute and FromServicesAttribute to support derived types

- Replace OfType<FromKeyedServicesAttribute>() with IsAssignableFrom checks in BindingInfo.cs
- Replace OfType<FromKeyedServicesAttribute>() with IsAssignableFrom checks in RequestDelegateFactory.cs
- Update SignalR HubMethodDescriptor pattern matching to handle derived FromKeyedServicesAttribute types
- Add test for derived FromKeyedServicesAttribute detection

Co-authored-by: captainsafia <[email protected]>

* Update FromKeyedServicesAttribute support for derived types in generators and validation

- Add InheritsFrom extension method to check class inheritance in SymbolExtensions
- Add TryGetAttributeInheritingFrom methods to detect attributes that inherit from base types
- Update IsServiceParameter in ITypeSymbolExtensions to use inheritance checking
- Update EndpointParameter.cs to use TryGetAttributeInheritingFrom for FromKeyedServicesAttribute
- Add CustomFromKeyedServicesAttribute test type and SupportsDerivedFromKeyedServicesAttribute test
- All builds and tests pass successfully

Co-authored-by: captainsafia <[email protected]>

* Feedback

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: captainsafia <[email protected]>
Co-authored-by: Safia Abdalla <[email protected]>

Author: Copilot

src/Http/Http.Extensions/gen/Microsoft.AspNetCore.Http.RequestDelegateGenerator/StaticRouteHandlerModel/EndpointParameter.cs

@@ -147,13 +147,13 @@ private void ProcessEndpointParameterSource(Endpoint endpoint, ISymbol symbol, I
         else if (attributes.HasAttributeImplementingInterface(wellKnownTypes.Get(WellKnownType.Microsoft_AspNetCore_Http_Metadata_IFromServiceMetadata)))
         {
             Source = EndpointParameterSource.Service;
-            if (attributes.TryGetAttribute(wellKnownTypes.Get(WellKnownType.Microsoft_Extensions_DependencyInjection_FromKeyedServicesAttribute), out var keyedServicesAttribute))
+            if (attributes.TryGetAttributeInheritingFrom(wellKnownTypes.Get(WellKnownType.Microsoft_Extensions_DependencyInjection_FromKeyedServicesAttribute), out var keyedServicesAttribute))
             {
                 var location = endpoint.Operation.Syntax.GetLocation();
                 endpoint.Diagnostics.Add(Diagnostic.Create(DiagnosticDescriptors.KeyedAndNotKeyedServiceAttributesNotSupported, location));
             }
         }
-        else if (attributes.TryGetAttribute(wellKnownTypes.Get(WellKnownType.Microsoft_Extensions_DependencyInjection_FromKeyedServicesAttribute), out var keyedServicesAttribute))
+        else if (attributes.TryGetAttributeInheritingFrom(wellKnownTypes.Get(WellKnownType.Microsoft_Extensions_DependencyInjection_FromKeyedServicesAttribute), out var keyedServicesAttribute))
         {
             Source = EndpointParameterSource.KeyedService;
             var constructorArgument = keyedServicesAttribute.ConstructorArguments.FirstOrDefault();

src/Http/Http.Extensions/src/RequestDelegateFactory.cs

@@ -804,15 +804,15 @@ private static Expression CreateArgument(ParameterInfo parameter, RequestDelegat
         }
         else if (parameter.CustomAttributes.Any(a => typeof(IFromServiceMetadata).IsAssignableFrom(a.AttributeType)))
         {
-            if (parameterCustomAttributes.OfType<FromKeyedServicesAttribute>().FirstOrDefault() is not null)
+            if (parameterCustomAttributes.FirstOrDefault(a => typeof(FromKeyedServicesAttribute).IsAssignableFrom(a.GetType())) is not null)
             {
                 throw new NotSupportedException(
                     $"The {nameof(FromKeyedServicesAttribute)} is not supported on parameters that are also annotated with {nameof(IFromServiceMetadata)}.");
             }
             factoryContext.TrackedParameters.Add(parameter.Name, RequestDelegateFactoryConstants.ServiceAttribute);
             return BindParameterFromService(parameter, factoryContext);
         }
-        else if (parameterCustomAttributes.OfType<FromKeyedServicesAttribute>().FirstOrDefault() is { } keyedServicesAttribute)
+        else if (parameterCustomAttributes.FirstOrDefault(a => typeof(FromKeyedServicesAttribute).IsAssignableFrom(a.GetType())) is FromKeyedServicesAttribute keyedServicesAttribute)
         {
             if (factoryContext.ServiceProviderIsService is not IServiceProviderIsKeyedService)
             {

src/Http/Http.Extensions/test/RequestDelegateGenerator/RequestDelegateCreationTests.KeyServices.cs

@@ -252,6 +252,23 @@ public async Task RequestDelegateGeneratesCompilableCodeForKeyedServiceInNamespa
         await VerifyResponseBodyAsync(httpContext, "To be or not to be…");
     }
 
+    [Fact]
+    public async Task SupportsDerivedFromKeyedServicesAttribute()
+    {
+        var source = """
+app.MapGet("/", (HttpContext context, [CustomFromKeyedServices("customKey")] TestService arg) => context.Items["arg"] = arg);
+""";
+        var (_, compilation) = await RunGeneratorAsync(source);
+        var myOriginalService = new TestService();
+        var serviceProvider = CreateServiceProvider((serviceCollection) => serviceCollection.AddKeyedSingleton("customKey", myOriginalService));
+        var endpoint = GetEndpointFromCompilation(compilation, serviceProvider: serviceProvider);
+
+        var httpContext = CreateHttpContext(serviceProvider);
+        await endpoint.RequestDelegate(httpContext);
+
+        Assert.Same(myOriginalService, httpContext.Items["arg"]);
+    }
+
     private class MockServiceProvider : IServiceProvider, ISupportRequiredService
     {
         public object GetService(Type serviceType)

src/Http/Http.Extensions/test/RequestDelegateGenerator/SharedTypes.cs

@@ -110,6 +110,11 @@ public class CustomFromBodyAttribute : Attribute, IFromBodyMetadata
     public bool AllowEmpty { get; set; }
 }
 
+public class CustomFromKeyedServicesAttribute : FromKeyedServicesAttribute
+{
+    public CustomFromKeyedServicesAttribute(object key) : base(key) { }
+}
+
 public enum TodoStatus
 {
     Trap, // A trap for Enum.TryParse<T>!

src/Mvc/Mvc.Abstractions/src/ModelBinding/BindingInfo.cs

@@ -177,7 +177,7 @@ public Type? BinderType
         }
 
         // Keyed services
-        if (attributes.OfType<FromKeyedServicesAttribute>().FirstOrDefault() is { } fromKeyedServicesAttribute)
+        if (attributes.FirstOrDefault(a => typeof(FromKeyedServicesAttribute).IsAssignableFrom(a.GetType())) is FromKeyedServicesAttribute fromKeyedServicesAttribute)
         {
             if (bindingInfo.BindingSource != null)
             {

src/Mvc/Mvc.Abstractions/test/ModelBinding/BindingInfoTest.cs

@@ -326,4 +326,31 @@ public void GetBindingInfo_ThrowsWhenWithFromKeyedServicesAttributeAndIFromServi
         // Act and Assert
         Assert.Throws<NotSupportedException>(() => BindingInfo.GetBindingInfo(attributes, modelMetadata));
     }
+
+    [Fact]
+    public void GetBindingInfo_WithDerivedFromKeyedServicesAttribute()
+    {
+        // Arrange
+        var key = new object();
+        var attributes = new object[]
+        {
+                new CustomFromKeyedServicesAttribute(key),
+        };
+        var modelType = typeof(Guid);
+        var provider = new TestModelMetadataProvider();
+        var modelMetadata = provider.GetMetadataForType(modelType);
+
+        // Act
+        var bindingInfo = BindingInfo.GetBindingInfo(attributes, modelMetadata);
+
+        // Assert
+        Assert.NotNull(bindingInfo);
+        Assert.Same(BindingSource.Services, bindingInfo.BindingSource);
+        Assert.Same(key, bindingInfo.ServiceKey);
+    }
+
+    private class CustomFromKeyedServicesAttribute : FromKeyedServicesAttribute
+    {
+        public CustomFromKeyedServicesAttribute(object key) : base(key) { }
+    }
 }

src/Shared/RoslynUtils/SymbolExtensions.cs

@@ -141,6 +141,36 @@ public static bool TryGetAttributeImplementingInterface(this ImmutableArray<Attr
         return false;
     }
 
+    public static bool HasAttributeInheritingFrom(this ISymbol symbol, INamedTypeSymbol baseType)
+    {
+        return symbol.TryGetAttributeInheritingFrom(baseType, out var _);
+    }
+
+    public static bool TryGetAttributeInheritingFrom(this ISymbol symbol, INamedTypeSymbol baseType, [NotNullWhen(true)] out AttributeData? matchedAttribute)
+    {
+        return symbol.GetAttributes().TryGetAttributeInheritingFrom(baseType, out matchedAttribute);
+    }
+
+    public static bool HasAttributeInheritingFrom(this ImmutableArray<AttributeData> attributes, INamedTypeSymbol baseType)
+    {
+        return attributes.TryGetAttributeInheritingFrom(baseType, out var _);
+    }
+
+    public static bool TryGetAttributeInheritingFrom(this ImmutableArray<AttributeData> attributes, INamedTypeSymbol baseType, [NotNullWhen(true)] out AttributeData? matchedAttribute)
+    {
+        foreach (var attributeData in attributes)
+        {
+            if (attributeData.AttributeClass is not null && attributeData.AttributeClass.InheritsFrom(baseType))
+            {
+                matchedAttribute = attributeData;
+                return true;
+            }
+        }
+
+        matchedAttribute = null;
+        return false;
+    }
+
     public static bool Implements(this ITypeSymbol type, ITypeSymbol interfaceType)
     {
         foreach (var t in type.AllInterfaces)
@@ -153,6 +183,18 @@ public static bool Implements(this ITypeSymbol type, ITypeSymbol interfaceType)
         return false;
     }
 
+    public static bool InheritsFrom(this ITypeSymbol type, ITypeSymbol baseType)
+    {
+        foreach (var t in type.GetThisAndBaseTypes())
+        {
+            if (SymbolEqualityComparer.Default.Equals(t, baseType))
+            {
+                return true;
+            }
+        }
+        return false;
+    }
+
     public static bool IsType(this INamedTypeSymbol type, string typeName, SemanticModel semanticModel)
         => SymbolEqualityComparer.Default.Equals(type, semanticModel.Compilation.GetTypeByMetadataName(typeName));
 

src/Validation/gen/Extensions/ITypeSymbolExtensions.cs

@@ -137,7 +137,7 @@ internal static bool IsServiceParameter(this IParameterSymbol parameter, INamedT
         return parameter.GetAttributes().Any(attr =>
             attr.AttributeClass is not null &&
             (attr.AttributeClass.ImplementsInterface(fromServiceMetadataSymbol) ||
-             SymbolEqualityComparer.Default.Equals(attr.AttributeClass, fromKeyedServiceAttributeSymbol)));
+             attr.AttributeClass.InheritsFrom(fromKeyedServiceAttributeSymbol)));
     }
 
     /// <summary>