Azure MSI with AdlsClient: Access token expired

Question

I am using Azure Managed Service Identity (MSI) to create a static (singleton) AdlsClient.

I, then, use the AdlsClient in a Functions app to write to a Data Lake store.

The app works fine for about a day but then it stops working and I see this error.

The access token in the 'Authorization' header is expired.”

Operation: CREATE failed with HttpStatus:Unauthorized Error

Apparently, the MSI token expires every day without warning.

Unfortunately, the MSI token provider doesn't return an expiry date along with the token so, I can't check to see if the token is still valid.

What is the right way to deal with this? Any help is appreciated.

Here's my code.

public static class AzureDataLakeUploaderClient
{
    private static Lazy<AdlsClient> lazyClient = new Lazy<AdlsClient>(InitializeADLSClientAsync);

    public static AdlsClient AzureDataLakeClient => lazyClient.Value;

    private static AdlsClient InitializeADLSClientAsync()
    {

        var azureServiceTokenProvider = new AzureServiceTokenProvider();
        string accessToken = azureServiceTokenProvider.GetAccessTokenAsync("https://datalake.azure.net/").Result;
        var client = AdlsClient.CreateClient(GetAzureDataLakeConnectionString(), "Bearer " + accessToken);
        return client;
    }
}

Thanks!

Answer 1

The access token that GetAccessTokenAsync returns is guaranteed to not expire within the next 5 minutes. By default, Azure AD access tokens expire in an hour [1].

So, if you use the same token (with default expiration time) for more than an hour, you will get an "expired token" error message. Please initialize the AdlsClient with a token fetched from GetAccessTokenAsync every time you need to use the AdlsClient. GetAccessTokenAsync caches the access token in memory, and will automatically get a new token if it is within 5 minutes of expiry.

A lazy object always returns the same object that it was initialized with [2]. So, the AdlsClient continues to use old token.

References

[1] https://learn.microsoft.com/en-us/azure/active-directory/active-directory-configurable-token-lifetimes#token-types

[2] https://learn.microsoft.com/en-us/dotnet/framework/performance/lazy-initialization#basic-lazy-initialization

Answer 2

A recent update appeared in the link below to automatically refresh tokens for Storage Accounts: https://learn.microsoft.com/en-us/azure/storage/common/storage-auth-aad-msi

I've modified the code above and tested it with Azure Data Lake Store Gen1 successfully to auto refresh MSI tokens.

To implement the code for ADLS Gen1 I needed two libraries:

<PackageReference Include="Microsoft.Azure.Services.AppAuthentication" Version="1.2.0-preview3" />
<PackageReference Include="Microsoft.Azure.Storage.Common" Version="10.0.3" />

I then used this code to create an AdlsClient instance with a constantly refreshed token:

var miAuthentication = new AzureManagedIdentityAuthentication("https://datalake.azure.net/");
var tokenCredential = miAuthentication.GetAccessToken();
ServiceClientCredentials serviceClientCredential = new TokenCredentials(tokenCredential.Token);
var dataLakeClient = AdlsClient.CreateClient(clientAccountPath, serviceClientCredential);

Below is the class that I modified from the article to generically refresh tokens. This can now be used for auto refreshing MSI tokens for both ADLS Gen1("https://datalake.azure.net/") and Storage Accounts("https://storage.azure.com/") by providing the relevant resource address when instantiating AzureManagedIdentityAuthentication. Make sure to use the code in the link to create the StorageCredentials object for storage accounts.

using System;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Azure.Services.AppAuthentication;
using Microsoft.Azure.Storage.Auth;

namespace SharedCode.Authentication
{
    /// <summary>
    /// Class AzureManagedIdentityAuthentication.
    /// </summary>
    public class AzureManagedIdentityAuthentication
    {
        private string _resource = null;
        /// <summary>
        /// Initializes a new instance of the <see cref="AzureManagedIdentityAuthentication"/> class.
        /// </summary>
        /// <param name="resource">The resource.</param>
        public AzureManagedIdentityAuthentication(string resource)
        {
            _resource = resource;
        }
        /// <summary>
        /// Gets the access token.
        /// </summary>
        /// <returns>TokenCredential.</returns>
        public TokenCredential GetAccessToken()
        {
            // Get the initial access token and the interval at which to refresh it.
            AzureServiceTokenProvider azureServiceTokenProvider = new AzureServiceTokenProvider();
            var tokenAndFrequency = TokenRenewerAsync(azureServiceTokenProvider, CancellationToken.None).GetAwaiter().GetResult();

            // Create credentials using the initial token, and connect the callback function 
            // to renew the token just before it expires
            TokenCredential tokenCredential = new TokenCredential(tokenAndFrequency.Token,
                                                                    TokenRenewerAsync,
                                                                    azureServiceTokenProvider,
                                                                    tokenAndFrequency.Frequency.Value);
            return tokenCredential;
        }
        /// <summary>
        /// Renew the token
        /// </summary>
        /// <param name="state">The state.</param>
        /// <param name="cancellationToken">The cancellation token.</param>
        /// <returns>System.Threading.Tasks.Task&lt;Microsoft.Azure.Storage.Auth.NewTokenAndFrequency&gt;.</returns>
        private async Task<NewTokenAndFrequency> TokenRenewerAsync(Object state, CancellationToken cancellationToken)
        {
            // Use the same token provider to request a new token.
            var authResult = await ((AzureServiceTokenProvider)state).GetAuthenticationResultAsync(_resource);

            // Renew the token 5 minutes before it expires.
            var next = (authResult.ExpiresOn - DateTimeOffset.UtcNow) - TimeSpan.FromMinutes(5);
            if (next.Ticks < 0)
            {
                next = default(TimeSpan);
            }

            // Return the new token and the next refresh time.
            return new NewTokenAndFrequency(authResult.AccessToken, next);
        }
    }
}

Answer 3

If anyone else is hit with this issue, I was able to get this to work the following way.

We know from Varun's answer that "GetAccessTokenAsync caches the access token in memory, and will automatically get a new token if it is within 5 minutes of expiry"

So, we could just check whether the current access token isn't the same as the old one. This would only be true if we are within 5 minutes of the token expiry in which case we'd create a new static client. In all other cases, we'd just return the existing client.

Something like this...

    private static AzureServiceTokenProvider azureServiceTokenProvider = new AzureServiceTokenProvider();

    private static string accessToken = GetAccessToken();

    private static AdlsClient azureDataLakeClient = null;

    public static AdlsClient GetAzureDataLakeClient()
    {
        var newAccessToken = GetAccessToken();
        if (azureDataLakeClient == null || accessToken != newAccessToken)
        {
            // Create new AdlsClient with the new token
            CreateDataLakeClient(newAccessToken);
        }

        return azureDataLakeClient;
    }

    private static string GetAccessToken()
    {
        return azureServiceTokenProvider.GetAccessTokenAsync("https://datalake.azure.net/").Result;
    }

Answer 4

Prerequisites

We need to know the following information to come up with an efficient solution:

1. Your assembly in an Azure Function app is loaded upon Function startup. However, for each invocation, the same loaded assembly is used for invoking your function app's method. This implies that any singletons will be persisted across invocations of your Azure Function.
2. AzureServiceTokenProvider caches your token between calls to GetAccessTokenAsyncfor each resource.
3. AdlsClient saves the token in a thread-safe manner and only uses it when you ask it do something. Furthermore, it provides a way to update the token in a thread-safe manner.

Solution

    using System;
    using System.Collections.Concurrent;
    using System.Threading;
    using System.Threading.Tasks;

    using Microsoft.Azure.DataLake.Store;
    using Microsoft.Azure.Services.AppAuthentication;

    public class AdlsClientFactory
    {
        private readonly ConcurrentDictionary<string, Lazy<AdlsClient>> adlsClientDictionary;

        public AdlsClientFactory()
        {
            this.adlsClientDictionary = new ConcurrentDictionary<string, Lazy<AdlsClient>>();
        }

        public async Task<IDataStoreClient> CreateAsync(string fqdn)
        {
            Lazy<AdlsClient> lazyClient = this.adlsClientDictionary.GetOrAdd(fqdn, CreateLazyAdlsClient);
            AdlsClient adlsClient = lazyClient.Value;

            // Get new token if old token expired otherwise use same token
            var azureServiceTokenProvider = new AzureServiceTokenProvider();
            string freshSerializedToken = await azureServiceTokenProvider.GetAccessTokenAsync("https://datalake.azure.net/");

            // "Bearer" + accessToken is done by the <see cref="AdlsClient.SetToken" /> command.
            adlsClient.SetToken(freshSerializedToken);

            return new AdlDataStoreClient(adlsClient);
        }

        private Lazy<AdlsClient> CreateLazyAdlsClient(string fqdn)
        {
            // TODO: This is just a sample. Figure out how to remove thread blocking while using lazy if that's important to you.
            var azureServiceTokenProvider = new AzureServiceTokenProvider();
            string freshSerializedToken = azureServiceTokenProvider.GetAccessTokenAsync("https://datalake.azure.net/").Result;
            return new Lazy<AdlsClient>(() => AdlsClient.CreateClient(fqdn, "Bearer " + freshSerializedToken), LazyThreadSafetyMode.ExecutionAndPublication);
        }
    }