Refresh Token using Polly with Named Client

Question

I have a policy that looks like this

var retryPolicy = Policy
    .Handle<HttpRequestException>()
    .OrResult<HttpResponseMessage>(resp => resp.StatusCode == HttpStatusCode.Unauthorized)
    .WaitAndRetryAsync(3, 
        retryAttempt => TimeSpan.FromSeconds(Math.Pow(2, retryAttempt)),
        onRetry: (resp, timeSpan, context) =>
        {
            // not sure what to put here
        });

Then I have a named client that looks like this

services.AddHttpClient("MyClient", client =>
    {
        client.BaseAddress = new Uri("http://some-url.com");
        client.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", authToken);
        client.Timeout = 30000;
    })
    .AddPolicyHandler(retryPolicy);

I need to refresh the bearer token on the http client in the event I receive a 401. So in a perfect world the following code would do exactly what I'm trying to accomplish

var retryPolicy = Policy
    .Handle<HttpRequestException>()
    .OrResult<HttpResponseMessage>(resp => resp.StatusCode == HttpStatusCode.Unauthorized)
    .WaitAndRetryAsync(3, 
        retryAttempt => TimeSpan.FromSeconds(Math.Pow(2, retryAttempt)),
        onRetry: (resp, timeSpan, context) =>
        {
            var newToken = GetNewToken();
            
            //httpClient doesn't exists here so I need to grab it some how
            httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", newToken);
        });

I have read the following articles:

Re-establishing authentication using Retry

Refresh Token using Polly with Typed Client

retry-to-refresh-authorization

and a couple others as well. However, they all seem use policy.ExecuteAsync() which I don't want to use because then I would have to change all the HttpClient calls throughout my solution. I'm trying to find a way to simply add this functionality to every request by only changing code in the StartUp.cs.

Answer 1

TL;DR: You need to define a communication protocol between a RetryPolicy, a DelegatingHandler and a TokenService.

---

In case of Typed Clients you can explicitly call the ExecuteAsync and use the Context to exchange data between the to-be-decorated method and the onRetry(Async) delegate.

This trick can't be used in a named client situation. What you need to do instead:

• Separate out the Token management into a dedicated service
• Use a DelegatingHandler to intercept the HttpClient's communication

This sequence diagram depicts the communication between the different components

Token Service

The DTO

public class Token
{
    public string Scheme { get; set; }
    public string AccessToken { get; set; }
}

The interface

public interface ITokenService
{
    Token GetToken();
    Task RefreshToken();
}

The dummy implementation

public class TokenService : ITokenService
{
    private DateTime lastRefreshed = DateTime.UtcNow;
    public Token GetToken()
        => new Token { Scheme = "Bearer", AccessToken = lastRefreshed.ToString("HH:mm:ss")}; 

    public Task RefreshToken()
    {
        lastRefreshed = DateTime.UtcNow;
        return Task.CompletedTask;
    }
}

The registration into the DI as Singleton

public void ConfigureServices(IServiceCollection services)
{
    services.AddSingleton<ITokenService, TokenService>();
    ...
}

Delegating Handler

The custom exception

public class OutdatedTokenException : Exception
{

}

The handler (interceptor)

public class TokenFreshnessHandler : DelegatingHandler
{
    private readonly ITokenService tokenService;
    public TokenFreshnessHandler(ITokenService service)
    {
        tokenService = service;
    }

    protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
    {
        var token = tokenService.GetToken();
        request.Headers.Authorization = new AuthenticationHeaderValue(token.Scheme, token.AccessToken);

        var response = await base.SendAsync(request, cancellationToken);
        if (response.StatusCode == HttpStatusCode.Unauthorized)
        {
            throw new OutdatedTokenException();
        }
        return response;
    }
}

• It retrieves the current token from the TokenService
• It sets the authorization header
• It executes the base method
• It checks the response's status
  • If 401 then it throws the custom exception
  • If other than 401 then it returns with the response

The registration into the DI as Transient

public void ConfigureServices(IServiceCollection services)
{
    services.AddSingleton<ITokenService, TokenService>();
    services.AddTransient<TokenFreshnessHandler>();
    ...
}

Retry Policy

The policy definition

public IAsyncPolicy<HttpResponseMessage> GetTokenRefresher(IServiceProvider provider)
{
    return Policy<HttpResponseMessage>
        .Handle<OutdatedTokenException>()
        .RetryAsync(async (_, __) => await provider.GetRequiredService<ITokenService>().RefreshToken());
}

• It receives an IServiceProvider to be able to access the TokenService
• It performs a single retry if an OutdatedTokenException was thrown
• Inside the onRetryAsync delegate it calls the TokenService's RefreshToken method

Putting all things together

public void ConfigureServices(IServiceCollection services)
{
    services.AddSingleton<ITokenService, TokenService>();
    services.AddTransient<TokenFreshnessHandler>();
    services.AddHttpClient("TestClient")
        .AddPolicyHandler((provider, _) => GetTokenRefresher(provider))
        .AddHttpMessageHandler<TokenFreshnessHandler>();
    ...
}

• Please bear in mind that the ordering of AddPolicyHandler and AddHttpMessageHandler matters
• If you would call the AddHttpMessageHandler first and then the AddPolicyHandler in that case your retry would not be triggered

Answer 2

This post contains an alternative version of my previously suggested solution.

I'm posting this as a separate answer (rather than editing the previous one) because both solutions are viable and the other post is already a lengthy one.

---

Why do we need an alternative version?

Because the TokenFreshnessHandler has too much responsibility whereas the Retry policy has too few.

If you look at the SendAsync method overridden implementation then you can see that it perform some operation on the request and on the response as well.

If we could make a separation where

• the handler deals only with the request
• and policy make its assessment on the response

then we would end up with a much cleaner solution (IMHO).

How can we achieve this separation?

If we could use the Polly's Context as an intermediate storage between the retry attempts then we were able to do this separation. Fortunately the Microsoft.Extensions.Http.Polly package defines two extension methods against the HttpRequestMessage:

• SetPolicyExecutionContext
• GetPolicyExecutionContext

These are under-documented features. On the docs.microsoft I could not even find the related pages. I have only found them under the dotnet-api-docs repo.

These can be useful if we know that the AddPolicyHandler attaches a new Context to the request only if it did not have one already. Unfortunately, this is yet again not documented, so it is an implementation detail which might change in the future. But currently we can rely on this.

How will this change the protocol?

As you see only difference here is the usage of the Context.

How should we change the handler?

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, tokenService.GetToken() } });
            request.SetPolicyExecutionContext(context);
        }

        var token = (Token)context[TokenKey];
        request.Headers.Authorization = new AuthenticationHeaderValue(token.Scheme, token.AccessToken);

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

• I've changed the name of the handler since its responsibilities have changed
• Now, the handler's implementation only cares about the request (and does not care about the response)
• As it was said previously: the PolicyHttpMessageHandler creates a new Context if there wasn't any
  • Because of this the GetPolicyExecutionContext does not return null (even for the very first attempt) rather than a Context with an empty context data collection (context.Count == 0)

How should we change the policy?

public IAsyncPolicy<HttpResponseMessage> GetTokenRefresher(IServiceProvider provider, HttpRequestMessage request)
{
    return Policy<HttpResponseMessage>
        .HandleResult(response => response.StatusCode == HttpStatusCode.Unauthorized)
        .RetryAsync(async (_, __) =>
        {
            await provider.GetRequiredService<ITokenService>().RefreshToken();
            request.SetPolicyExecutionContext(new Context());
        });
}

• Rather than triggering the policy for a custom exception, now it triggers in case of 401 response's status code
• The onRetryAsync has been modified in the way that it clears the attached context of the request

The registration code should be adjusted as well

services.AddHttpClient("TestClient")
    .AddPolicyHandler((sp, request) => GetTokenRefresher(sp, request))
    .AddHttpMessageHandler<TokenRetrievalHandler>()

• Now, we should pass to the GetTokenRefresher method not just the IServiceProvider but also HttpRequestMessage as well

Which solution should I use?

• This solution offers nicer separation but it relies on an implementation detail
• The other solution makes the handler smart whereas the policy dumb

Answer 3

As it was asked by Chris Harrington let me present here yet another variant: using typed HttpClient

In case of Typed Clients you can explicitly call the ExecuteAsync and use the Context to exchange data between the to-be-decorated method and the onRetry(Async) delegate.

The good news is that we don't need to use Context. We can have a solution which does not require us to populate, propagate and fetch context data.

As always let's start with a sequence diagram:

• As you can see we don't have a DelegatingHandler component here
  • Because the retry is defined and used inside the typed client

Token Service

This component is exactly the same as in the other two variants.
So, I will not copy here the dummy implementation.

Typed Client

public interface IClient
{
    Task<string> GetAsync(string url);
}

public class Client : IClient
{
    private readonly HttpClient _client;
    private readonly ITokenService _tokenService;
    public Client(HttpClient client, ITokenService tokenService)
      => (_client, _tokenService) = (client, tokenService);

    public async Task<string> GetAsync(string url)
    {
        var response = await GetRetryPolicy().ExecuteAsync(() =>
        {
            var token = _tokenService.GetToken();
            var request = new HttpRequestMessage() { RequestUri = new Uri(url) };
            request.Headers.Authorization = new AuthenticationHeaderValue(token.Scheme, token.AccessToken);

            return _client.SendAsync(request);
        });
            
        return await response.Content.ReadAsStringAsync();
    }

    private IAsyncPolicy<HttpResponseMessage> GetRetryPolicy()
    => Policy<HttpResponseMessage>
        .HandleResult(res => res.StatusCode == HttpStatusCode.Unauthorized)
        .RetryAsync((dr, _) => _tokenService.RefreshToken());
}

• We have defined a retry policy (GetRetryPolicy) which triggers for 401 and calls the RefreshToken before the next attempt
• The HttpClient call is wrapped into the previous policy
  • First we retrieve the current token from the TokenService then we set it to the appropriate request header
  • We issue the request against the downstream and we are returning with the result

Component registration

This part becomes super simple

builder.Services.AddSingleton<ITokenService, TokenService>();
builder.Services.AddHttpClient<IClient, Client>();

And that's it. IClient implementation is ready to be used :)