Re-Authorize Efficiently Using Polly And .NET HttpClientFactory

In today’s world, we are using a lot of APIs to build our applications. To make your .NET based applications more ressilient and fault tolerant, go to solution is to use Polly. Probably, you are already using it. So,I’m not going to explain how to use it in general. You can find some excellent documentation here.

However, In this post I’m going to discuss about one of the typical scenario where we need to refresh the authentication token when using Polly. Let’s get started.

UPDATE: If you are using Polly v8 and .NET 8, you can find the updated post here

The Problem

To handle such situation what we do is retry on Unauthorized failure and call a method to reauthorize. Something like below from the Polly documentation -

FromDocumentation.cs

HttpClient httpClient = ...

var authEnsuringPolicy = Policy<HttpResponseMessage>  
  .HandleResult(r => r.StatusCode == StatusCode.Unauthorized)
  .RetryAsync(1, onRetryAsync: async (e, i) => await AuthHelper.RefreshAuthorization(httpClient));

var httpResponseMessage =  
    await authEnsuringPolicy.ExecuteAsync(() => httpClient.GetAsync(uri));

The above pattern works by sharing the httpClient variable across closures, but a significant drawback is that you have to declare and use the policy in the same scope. This makes impossible to use dependency injection approach where you define policies centrally on startup, then provide them by DI to the point of use. Using DI with Polly in this way is a powerful pattern for separation of concerns, and allows easy stubbing out of Polly in unit testing.

From Polly v5.1.0, with Context available as a state variable to every delegate, the policy declaration can be rewritten:

FromDocumentation.cs

var authEnsuringPolicy = Policy<HttpResponseMessage>  
  .HandleResult(r => r.StatusCode == StatusCode.Unauthorized)
  .RetryAsync(1, onRetryAsync: async (ex, i, context) => await AuthHelper.RefreshAuthorization(context["httpClient"]));

And the usage (elsewhere in the codebase):

FromDocumentation.cs

var httpResponseMessage =  
    await authEnsuringPolicy.ExecuteAsync(context => context["httpClient"].GetAsync(uri), 
    contextData: new Dictionary<string, object> {{"httpClient", httpClient}}
    );

Passing context as state-data parameters in the different parts of policy execution allows policy declaration and usage now to be separate.

The above example and explanation is taken from the Polly documentation. You can find the same here.

IMO, the above approach is not very clean as we are passing httpClient as a state variable. Instead if we can pass the Token as a state variable, it will be more cleaner. Let’s try to implement the same.

The Solution

The solution could be to make use of the Context , DelegationHandler and HttpClientFactory. Define a communication between them as below -

In our scenario, we are using Auth0 as Token Provider and our Downstream system (Todo API)is protected by it.

First, we will focus on Auth0 Service. This service is responsible for getting the actual Token from Auth0 and refresh it when it is expired. We will start by defining Token DTO as below -

ITokenService.cs

public record Token
{
    public static Token Empty => new();

    [JsonPropertyName("token_type")]
    public string Scheme { get; set; } = default!;

    [JsonPropertyName("access_token")]
    public string AccessToken { get; set; } = default!;

    [JsonPropertyName("expires_in")]
    public double ExpiresIn { get; set; } = default!;
}

We will also define a Auth0Service which will be responsible for the actual hard work -

IAuth0Service.cs

public interface IAuth0Service
{
    Task<Token> GetTokenAsync();
}
public class Auth0Service : IAuth0Service
{
    private readonly HttpClient _httpClient;
    private readonly Auth0Option _settings;
    public Auth0Service(HttpClient client, Auth0Option settings)
    {
        _httpClient = client;
        _settings = settings;
    }
    public async Task<Token> GetTokenAsync()
    {
        var content = new FormUrlEncodedContent(new Dictionary<string, string>
        {
            ["client_id"] = _settings.ClientId,
            ["client_secret"] = _settings.ClientSecret,
            ["audience"] = _settings.Audience,
            ["grant_type"] = _settings.GrantType,
            ["scope"] = _settings.Scope
        });

            var result= await  _httpClient.PostAsync(_settings.TokenUrl, content);
            if (!result.IsSuccessStatusCode)
            {
               return Token.Empty;
            }
            var token = await result.Content.ReadFromJsonAsync<Token>();

            return token!;
    }
}

It will work as Typed Client for Auth0.

Next, we will focus on Token Service. This service is responsible for getting the cached Token from memory and refresh it from Auth0 Service when it is expired. Below is the implementation -

ITokenService.cs

public interface ITokenService
{
    ValueTask<Token> GetToken();
    Task<Token> RefreshToken();
}
public class TokenService : ITokenService
{
    private const string CacheKey = nameof(TokenService);

    private readonly IAuth0Service _auth0Service;
    private readonly IMemoryCache _memoryCache;
    public TokenService(IAuth0Service auth0Service, IMemoryCache memoryCache)
    {
        _auth0Service = auth0Service;
        _memoryCache = memoryCache;
    }
    public async ValueTask<Token> GetToken()
    {
        if (!_memoryCache.TryGetValue(CacheKey, out Token? cacheValue))
        {
            cacheValue = await RefreshToken();
        }
        return cacheValue!;
    }

    public async Task<Token> RefreshToken()
    {
        var token = await _auth0Service.GetTokenAsync();
        if (token != Token.Empty)
        {
            var expires_in = token.ExpiresIn>0?token.ExpiresIn-10:token.ExpiresIn;
            _memoryCache.Set(CacheKey, token, new MemoryCacheEntryOptions()
                                                    .SetSlidingExpiration(TimeSpan.FromMinutes(5))
                                                    .SetAbsoluteExpiration(TimeSpan.FromSeconds(expires_in)));
        }

        return token;
    }
}

Here, couple of things to note -

1. We are using ValueTask for GetToken. This is because we know that the Hot path will be to get the Token from the memory cache and very few scenraio(e.g very first time when Downstream is getting hit) it will actualy call the Auth0 Service. So, we are using ValueTask to avoid the overhead of allocating a new Task in the heap.

2. We are using SetAbsoluteExpiration and SetSlidingExpiration both to make sure we are not overwhelming by storing the unused Token in the memory cache. Also to avoid edge case when the Token is expired but not yet refreshed we are making sure it removed from the cache before it’s actual expiry time.

Time to rewrite our Policy for getting the refreshed Token -

ITokenService.cs

public static IAsyncPolicy<HttpResponseMessage> GetTokenRefresher(IServiceProvider provider, HttpRequestMessage request)
    {
        var delay = Backoff.ConstantBackoff(TimeSpan.FromMilliseconds(100), retryCount: 3);

        return Policy<HttpResponseMessage>
                .HandleResult(msg => msg.StatusCode == System.Net.HttpStatusCode.Unauthorized)
                .WaitAndRetryAsync(delay, async (_, _, _, _) =>
            {
                await provider.GetRequiredService<ITokenService>().RefreshToken();
                request.SetPolicyExecutionContext(new Context());
            });
    }

    public static IAsyncPolicy<HttpResponseMessage> GetConstantBackofffPolicy()
    {
        var delay = Backoff.ConstantBackoff(TimeSpan.FromMilliseconds(100), retryCount: 3);

        return HttpPolicyExtensions
            .HandleTransientHttpError()
            .OrResult(msg => msg.StatusCode == System.Net.HttpStatusCode.NotFound)
            .WaitAndRetryAsync(delay);
    }

Very simple and clean. We are making sure, after refreshing the token, the context is clear so that in DelegationHandler we can get the updated token. Below is the implementation of custom handler -

ITokenService.cs

public class TokenRetrievalHandler : DelegatingHandler
{
    private readonly ITokenService tokenService;
    private const string TokenRetrieval = nameof(TokenRetrieval);
    private const string TokenKey = nameof(TokenKey);
    public TokenRetrievalHandler(ITokenService service)
    {
        tokenService = service;
    }

    protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
    {
        var context = request.GetPolicyExecutionContext();
        if (context.Count == 0)
        {
            context = new Context(TokenRetrieval, new Dictionary<string, object> { { TokenKey, await tokenService.GetToken() } });
            request.SetPolicyExecutionContext(context);
        }

        var token = (Token)context[TokenKey];

        if(token!= Token.Empty)
            request.Headers.Authorization = new AuthenticationHeaderValue(token.Scheme, token.AccessToken);

        return await base.SendAsync(request, cancellationToken);
    }
}

We are intercepting the actual call and making sure,the updated token is added in header always.

We have all the pieces in place. Now, we just need to orchestrate the flow. We will also make use of PolicyWrap to combine the same type of policies to make our application fault tolerant.

We will create ServiceCollection extension to keep our StartUp.cs clear and concise. First we will create the extension for Auth0 Service as below

ServiceCollectionExtensions.cs

public static IServiceCollection AddAuth0Service(this IServiceCollection services, HostBuilderContext context)
    {
        var auth0Option = context.Configuration.GetSection(Auth0Option.ConfigKey).Get<Auth0Option>()!;

        services.AddSingleton(auth0Option);

        services
                .AddHttpClient<IAuth0Service, Auth0Service>(client =>
                {
                    client.BaseAddress = new Uri(auth0Option.BaseAddress);
                    client.DefaultRequestHeaders.TryAddWithoutValidation("Content-Type", "application/json");
                })
                .AddPolicyHandler(PollyRetryPolicies.GetConstantBackofffPolicy());
        return services;
    }

Nothing fancy here, just making sure in case of any transient error, we are retrying 3 times with a constant backoff of 100ms.

Now, the intersting Typed Client for DownStream System which is in our case TodoService. Below is the extension method for the same -

ServiceCollectionExtensions.cs

public static IServiceCollection AddTodoService(this IServiceCollection services, HostBuilderContext context)
    {
        services.AddSingleton<ITokenService, TokenService>();
        services.AddTransient<TokenRetrievalHandler>();

        var todoServiceOption = context.Configuration.GetSection(TodoServiceOption.ConfigKey).Get<TodoServiceOption>()!;
        services
                .AddHttpClient<ITodoService, TodoService>(client => client.BaseAddress = new Uri(todoServiceOption.BaseUrl))
                .AddPolicyHandler((provider,request)=> Policy.WrapAsync(PollyRetryPolicies.GetConstantBackofffPolicy(), PollyRetryPolicies.GetTokenRefresher(provider, request)))
                .AddHttpMessageHandler<TokenRetrievalHandler>();
        return services;
    }

We have added the TokenRetrievalHandler to intercept any call from this Typed Client and attach the Token in the header. Also, we have added the PolicyWrap to make sure we are retrying 3 times with a constant backoff of 100ms in case of any transient error and also refresh the token in case of Unauthorized error.

And the last bit, The Console from where we are calling the API -

Program.cs

public static class Program
{
    public static async Task Main(string[] args)
    {
        var host = CreateHostBuilder(args).Build();

        var todoService = host.Services.GetRequiredService<ITodoService>();
        await GetAllTodo(todoService);

        await host.RunAsync();
    }

    static IHostBuilder CreateHostBuilder(string[] args) => Host
        .CreateDefaultBuilder(args)
         .ConfigureServices((context, services) => services.AddMemoryCache()
                                                           .AddAuth0Service(context)
                                                           .AddTodoService(context))
        .ConfigureAppConfiguration((_, config) => config.AddJsonFile("appsettings.json", optional: false, reloadOnChange: true));

    public static async Task GetAllTodo(ITodoService todoService)
    {
        var todos = await todoService.GetTodosAsync();
        if (todos is null)
        {
            Console.WriteLine("Todos is null");
            return;
        }
        Console.WriteLine($"Todos: {todos.Count()}");
    }
}

Here, we have configured the services for MemoryCache, Auth0Service and TodoService. Also, we have added the appsettings.json file to get the configuration for the same.

That’s it. We are done. Now, we can run the application.

Conclusion

IMO, the above approach is more cleaner and easy to understand. Also, it is more flexible and can be extended easily.Also, each service has it’s own responsibility and can be tested independently. I have used Auth0 as Token Provider but you can use any other provider as well with any other Downstream system.I have skipped the part for Auth0 configuration but let me know if you want me to explain it in seperate post. You can also use the same approach for Named Client as well. This post is inspired from below post

As always, You can find the source code here.