Add decorator or proxy to dependency injection

Published on Saturday, September 4, 2021, by Vincent Jacquet In Code Tags #depency-injection

By design, the dependency injection mechanism makes it difficult to know which type is instantiated when you get a service. Therefore it can become cumbersome to extend a service from a third party, not to mention difficult to maintain.

Basically, given a service interface, we would like to instantiate a decorator or a proxy on it, like this:

private class MyServiceDecorator : IAnyService
{
    private readonly IAnyService _underlying;

    public MyServiceDecorator(IAnyService underlying)
    {
        _underlying = underlying;
    }

    // ...
}

In order to do so, we have to remove the ServiceDescriptor of the existing service and use it to create another ServiceDescriptor with the same lifetime and a factory that instantiate our decorator. Finally, we have to keep in mind that the service container is responsible for disposing the the services it creates, which is handled by the Holder type.

public static IServiceCollection AddDecoratorOrProxy<TService, TImplementation>(this IServiceCollection services)
    where TService : class
    where TImplementation : class, TService
{
    var descriptor = services.Single(d => d.ServiceType == typeof(TService));
    services.Remove(descriptor);
    if (descriptor.ImplementationInstance != null)
    {
        Debug.Assert(descriptor.Lifetime == ServiceLifetime.Singleton);

        var instance = (TService)descriptor.ImplementationInstance;
        services.Add(new ServiceDescriptor(typeof(TService), sp => ActivatorUtilities.CreateInstance(sp, typeof(TImplementation), instance), ServiceLifetime.Singleton));
    }
    else if (descriptor.ImplementationFactory != null)
    {
        Func<IServiceProvider, object> factory = sp => descriptor.ImplementationFactory(sp);
        var disposerType = typeof(Holder<,>).MakeGenericType(descriptor.ServiceType, typeof(TImplementation));
        services.Add(new ServiceDescriptor(disposerType, sp => ActivatorUtilities.CreateInstance(sp, disposerType, factory(sp)), descriptor.Lifetime));
        services.Add(new ServiceDescriptor(typeof(TService), sp => ActivatorUtilities.CreateInstance<TImplementation>(sp, ((Holder)sp.GetRequiredService(disposerType)).Instance), descriptor.Lifetime));
    }
    else if (typeof(IDisposable).IsAssignableFrom(descriptor.ImplementationType))
    {
        Func<IServiceProvider, object> factory = sp => ActivatorUtilities.CreateInstance(sp, descriptor.ImplementationType);
        var disposerType = typeof(Holder<,>).MakeGenericType(descriptor.ServiceType, typeof(TImplementation));
        services.Add(new ServiceDescriptor(disposerType, sp => ActivatorUtilities.CreateInstance(sp, disposerType, factory(sp)), descriptor.Lifetime));
        services.Add(new ServiceDescriptor(typeof(TService), sp => ActivatorUtilities.CreateInstance<TImplementation>(sp, ((Holder)sp.GetRequiredService(disposerType)).Instance), descriptor.Lifetime));
    }
    else
    {
        Func<IServiceProvider, object> factory = sp => ActivatorUtilities.CreateInstance(sp, descriptor.ImplementationType);
        services.Add(new ServiceDescriptor(typeof(TService), sp => ActivatorUtilities.CreateInstance<TImplementation>(sp, factory(sp)), descriptor.Lifetime));
    }
    return services;
}

#region Holder

private class Holder : IDisposable
{
    public readonly object Instance;

    public Holder(object instance)
    {
        Instance = instance;
    }

    public void Dispose()
    {
        (Instance as IDisposable)?.Dispose();
    }
}

private class Holder<TService, TImplementation> : Holder
    where TService : class
    where TImplementation : class, TService
{
    public Holder(object instance)
        : base(instance)
    {
    }
}

#endregion

The implementation of IdentityServer4 used another implementation based on a Decorator type to decorate the original type and then inject the decorated version in the real decorator type.

internal class UserClaimsFactory<TUser> : IUserClaimsPrincipalFactory<TUser>
    where TUser : class
{
    private readonly Decorator<IUserClaimsPrincipalFactory<TUser>> _inner;
    private UserManager<TUser> _userManager;

    public UserClaimsFactory(Decorator<IUserClaimsPrincipalFactory<TUser>> inner, UserManager<TUser> userManager)
    {
        _inner = inner;
        _userManager = userManager;
    }

    // ...
}

The decorator is injected like this:

builder.Services.AddTransientDecorator<IUserClaimsPrincipalFactory<TUser>, UserClaimsFactory<TUser>>();

Compared to the implementation we just share, this version has two drawbacks:

The decorator is declared transient
The mechanism is not transparent as you have to use the Decorator type to pass in the underlying service.

The full code of AddTransientDecorator and Decorator can be found on github.

Leave a comment