OWIN OAuth JWT refresh token fails intermittently in load-balanced situation

Question

I'm currently working on an OWIN OAuth implementation which uses JWT and supports token refreshing. I'm having intermittent problems with the token refresh process. The process works reliably on my development environment, but when published onto our Azure Service Fabric test environment, which is setup in a 3-node load-balanced configuration, the refresh token request often fails (not always!), and I get the infamous "invalid_grant" error.

I've found that the refresh token works successfully when being handled by the same service fabric node that issued it originally. However, it always fails when handled by a different node.

My understanding is that by using JWT, having a micro-service infrastructure deliver a load-balanced authentication server get's around the "machine-key" related issues that arise from using the OOTB access token format provided by OWIN.

Failed refresh tokens are making their way into the IAuthenticationTokenProvider.ReceiveAsync method, but the OAuthAuthorizationServerProvider.GrantRefreshToken method is never being hit, suggesting something in the OWIN middle-ware is not happy with the refresh token. Can anyone offer any insight into what the cause may be?

Now for the code, there's quite a bit - apologies for all the reading!

The authentication server is a service fabric stateless service, here's the ConfigureApp method:

protected override void ConfigureApp(IAppBuilder appBuilder)
    {
        appBuilder.UseCors(CorsOptions.AllowAll);
        var oAuthAuthorizationServerOptions = InjectionContainer.GetInstance<OAuthAuthorizationServerOptions>();
        appBuilder.UseOAuthAuthorizationServer(oAuthAuthorizationServerOptions);
        appBuilder.UseJwtBearerAuthentication(InjectionContainer.GetInstance<JwtBearerAuthenticationOptions>());
        appBuilder.UseWebApi(GetHttpConfiguration(InjectionContainer));
    }

Here's the implementation of OAuthAuthorizationServerOptions:

public class AppOAuthOptions : OAuthAuthorizationServerOptions
{
    public AppOAuthOptions(IAppJwtConfiguration configuration,
        IAuthenticationTokenProvider authenticationTokenProvider,
        IOAuthAuthorizationServerProvider authAuthorizationServerProvider)
    {
        AllowInsecureHttp = true; 
        TokenEndpointPath = "/token";
        AccessTokenExpireTimeSpan = configuration.ExpirationMinutes;
        AccessTokenFormat = new AppJwtWriterFormat(this, configuration);
        Provider = authAuthorizationServerProvider;
        RefreshTokenProvider = authenticationTokenProvider;
    }
}

And here's the JwtBearerAuthenticationOptions implementation:

public class AppJwtOptions : JwtBearerAuthenticationOptions
{
    public AppJwtOptions(IAppJwtConfiguration config)
    {
        AuthenticationMode = AuthenticationMode.Active;
        AllowedAudiences = new[] {config.JwtAudience};
        IssuerSecurityTokenProviders = new[]
        {
            new SymmetricKeyIssuerSecurityTokenProvider(
                config.JwtIssuer,
                Convert.ToBase64String(Encoding.UTF8.GetBytes(config.JwtKey)))
        };
    }
}

public class InMemoryJwtConfiguration : IAppJwtConfiguration
{
    AppSettings _appSettings;

    public InMemoryJwtConfiguration(AppSettings appSettings)
    {
        _appSettings = appSettings;
    }

    public int ExpirationMinutes
    {
        get { return 15; }
        set { }
    }

    public string JwtAudience
    {
        get { return "CENSORED AUDIENCE"; }
        set { }
    }

    public string JwtIssuer
    {
        get { return "CENSORED ISSUER"; }
        set { }
    }

    public string JwtKey
    {
        get { return "CENSORED KEY :)"; }
        set { }
    }

    public int RefreshTokenExpirationMinutes
    {
        get { return 60; }
        set { }
    }

    public string TokenPath
    {
        get { return "/token"; }
        set { }
    }
}

And the ISecureData implementation:

public class AppJwtWriterFormat : ISecureDataFormat<AuthenticationTicket>
{
    public AppJwtWriterFormat(
        OAuthAuthorizationServerOptions options,
        IAppJwtConfiguration configuration)
    {
        _options = options;
        _configuration = configuration;
    }

    public string Protect(AuthenticationTicket data)
    {
        if (data == null)
            throw new ArgumentNullException(nameof(data));

        var now = DateTime.UtcNow;
        var expires = now.AddMinutes(_options.AccessTokenExpireTimeSpan.TotalMinutes);

        var symmetricKey = Encoding.UTF8.GetBytes(_configuration.JwtKey);

        var signingCredentials = new SigningCredentials(
            new InMemorySymmetricSecurityKey(symmetricKey),
            SignatureAlgorithm, DigestAlgorithm);

        var token = new JwtSecurityToken(
            _configuration.JwtIssuer,
            _configuration.JwtAudience,
            data.Identity.Claims,
            now,
            expires,
            signingCredentials);

        return new JwtSecurityTokenHandler().WriteToken(token);
    }

    public AuthenticationTicket Unprotect(string protectedText)
    {
        throw new NotImplementedException();
    }
}

This is the IAuthenticationTokenProvider implementation:

public class RefreshTokenProvider : IAuthenticationTokenProvider
{
    private readonly IAppJwtConfiguration _configuration;
    private readonly IContainer _container;

    public RefreshTokenProvider(IAppJwtConfiguration configuration, IContainer container)
    {
        _configuration = configuration;
        _container = container;
        _telemetry = telemetry;
    }

    public void Create(AuthenticationTokenCreateContext context)
    {
        CreateAsync(context).Wait();
    }

    public async Task CreateAsync(AuthenticationTokenCreateContext context)
    {
        try
        {
            var refreshTokenId = Guid.NewGuid().ToString("n");
            var now = DateTime.UtcNow;

            using (var container = _container.GetNestedContainer())
            {
                var hashLogic = container.GetInstance<IHashLogic>();
                var tokenStoreLogic = container.GetInstance<ITokenStoreLogic>();

                var userName = context.Ticket.Identity.FindFirst(ClaimTypes.UserData).Value;
                var userToken = new UserToken
                {
                    Email = userName,
                    RefreshTokenIdHash = hashLogic.HashInput(refreshTokenId),
                    Subject = context.Ticket.Identity.Name,
                    RefreshTokenExpiresUtc =
                        now.AddMinutes(Convert.ToDouble(_configuration.RefreshTokenExpirationMinutes)),
                    AccessTokenExpirationDateTime =
                        now.AddMinutes(Convert.ToDouble(_configuration.ExpirationMinutes))
                };

                context.Ticket.Properties.IssuedUtc = now;
                context.Ticket.Properties.ExpiresUtc = userToken.RefreshTokenExpiresUtc;
                context.Ticket.Properties.AllowRefresh = true;

                userToken.RefreshToken = context.SerializeTicket();

                await tokenStoreLogic.CreateUserTokenAsync(userToken);
                context.SetToken(refreshTokenId);
            }
        }
        catch (Exception ex)
        {
            // exception logging removed for brevity
            throw;
        }
    }

    public void Receive(AuthenticationTokenReceiveContext context)
    {
        ReceiveAsync(context).Wait();
    }

    public async Task ReceiveAsync(AuthenticationTokenReceiveContext context)
    {
        try
        {
            using (var container = _container.GetNestedContainer())
            {
                var hashLogic = container.GetInstance<IHashLogic>();
                var tokenStoreLogic = container.GetInstance<ITokenStoreLogic>();

                var hashedTokenId = hashLogic.HashInput(context.Token);
                var refreshToken = await tokenStoreLogic.FindRefreshTokenAsync(hashedTokenId);

                if (refreshToken == null)
                {
                    return;
                }

                context.DeserializeTicket(refreshToken.RefreshToken);
                await tokenStoreLogic.DeleteRefreshTokenAsync(hashedTokenId);
            }
        }
        catch (Exception ex)
        {
            // exception logging removed for brevity
            throw;
        }
    }
}

And finally, this is the OAuthAuthorizationServerProvider implementation:

public class AppOAuthProvider : OAuthAuthorizationServerProvider
{
    public override Task GrantRefreshToken(OAuthGrantRefreshTokenContext context)
    {
        if (context.ClientId != null)
        {
            context.Rejected();
            return Task.FromResult(0);
        }

        // Change authentication ticket for refresh token requests
        var newIdentity = new ClaimsIdentity(context.Ticket.Identity);
        newIdentity.AddClaim(new Claim("newClaim", "refreshToken"));

        var newTicket = new AuthenticationTicket(newIdentity, context.Ticket.Properties);
        context.Validated(newTicket);

        return Task.FromResult(0);
    }

    public override async Task GrantResourceOwnerCredentials(OAuthGrantResourceOwnerCredentialsContext context)
    {
        using (var container = _container.GetNestedContainer())
        {
            var requestedAuthenticationType = context.Request.Query["type"];
            var requiredAuthenticationType = (int)AuthenticationType.None;

            if (string.IsNullOrEmpty(requestedAuthenticationType) || !int.TryParse(requestedAuthenticationType, out requiredAuthenticationType))
            {
                context.SetError("Authentication Type Missing", "Type parameter is required to check which type of user you are trying to authenticate with.");
                context.Response.StatusCode = (int)HttpStatusCode.BadRequest;
                return;
            }

            var authenticationWorker = GetInstance<IAuthenticationWorker>(container);
            var result = await authenticationWorker.AuthenticateAsync(new AuthenticationRequestViewModel
            {
                UserName = context.UserName,
                Password = context.Password,
                IpAddress = context.Request.RemoteIpAddress ?? "",
                UserAgent = context.Request.Headers.ContainsKey("User-Agent") ? context.Request.Headers["User-Agent"] : ""
            });

            if (result.SignInStatus != SignInStatus.Success)
            {
                context.SetError(result.SignInStatus.ToString(), result.Message);
                context.Response.StatusCode = (int)HttpStatusCode.Unauthorized;
                return;
            }

            // After we have successfully logged in. Check the authentication type for the just authenticated user
            var userAuthenticationType = (int)result.AuthenticatedUserViewModel.Type;
            // Check if the auth types match
            if (userAuthenticationType != requiredAuthenticationType)
            {
                context.SetError("Invalid Account", "InvalidAccountForPortal");
                context.Response.StatusCode = (int)HttpStatusCode.Unauthorized;
                return;
            }

            var identity = SetClaimsIdentity(context, result.AuthenticatedUserViewModel);
            context.Validated(identity);
        }
    }

    public override async Task TokenEndpointResponse(OAuthTokenEndpointResponseContext context)
    {
        using (var container = GetNestedContainer())
        {
            var email = context.Identity.FindFirst(ClaimTypes.UserData).Value;

            var accessTokenHash = _hashLogic.HashInput(context.AccessToken);
            var tokenStoreLogic = GetInstance<ITokenStoreLogic>(container);

            await tokenStoreLogic.UpdateUserTokenAsync(email, accessTokenHash);

            var authLogic = GetInstance<IAuthenticationLogic>(container);
            var userDetail = await authLogic.GetDetailsAsync(email);

            context.AdditionalResponseParameters.Add("user_id", email);
            context.AdditionalResponseParameters.Add("user_name", userDetail.Name);
            context.AdditionalResponseParameters.Add("user_known_as", userDetail.KnownAs);
            context.AdditionalResponseParameters.Add("authentication_type", userDetail.Type);
        }

        await base.TokenEndpointResponse(context);
    }

    public override Task ValidateClientAuthentication(OAuthValidateClientAuthenticationContext context)
    {
        context.Validated();
        return Task.FromResult(0);
    }

    private ClaimsIdentity SetClaimsIdentity(OAuthGrantResourceOwnerCredentialsContext context, AuthenticatedUserViewModel user)
    {
        var identity = new ClaimsIdentity(
            new[]
            {
                new Claim(ClaimTypes.Name, context.UserName),
                new Claim(ClaimTypes.SerialNumber, user.SerialNumber),
                new Claim(ClaimTypes.UserData, user.Email.ToString(CultureInfo.InvariantCulture)),
                new Claim(ClaimTypeUrls.AdminScope, user.Scope.ToString()),
                new Claim(ClaimTypeUrls.DriverId, user.DriverId.ToString(CultureInfo.InvariantCulture)),
                new Claim(ClaimTypeUrls.AdministratorId, user.AdministratorId.ToString(CultureInfo.InvariantCulture))
            },
            _authenticationType
        );

        //add roles
        var roles = user.Roles;
        foreach (var role in roles)
            identity.AddClaim(new Claim(ClaimTypes.Role, role));

        return identity;
    }
}