Lesson Learned #395:Token Management and Retry Strategy for Azure Active Directory Authentication

[Jose_Manuel_Jurado]

In this blog post, we'll explore how to optimize token retrieval from Azure Active Directory (AAD) by implementing token caching and a retry policy. By utilizing these techniques, developers can minimize unnecessary requests to the AAD directory and improve the performance of their applications. We'll walk through a PowerShell code example that demonstrates this approach.

$tenantId = "72f988bf-XXX"
$clientId = "4247a0c0-XXXX"
$clientSecret = "Agn8Q~XXXX"

$sqlServerUrl = "Servername.database.windows.net"
$database = "DBName"
$connectionString = "Server=tcp:$sqlServerUrl,1433;Initial Catalog=$database;Persist Security Info=False;MultipleActiveResultSets=False;Encrypt=True;TrustServerCertificate=False;"

function Get-AzureADToken {
    $URI = "https://login.microsoftonline.com/$tenantId/oauth2/token"
    $armresource = "https://database.windows.net/"
    $body = @{
        resource = $armresource
        grant_type = "client_credentials"
        client_id = $clientId
        client_secret = $clientSecret
    }
    $request = Invoke-RestMethod -Method POST -Uri $URI -Body $body -ContentType "application/x-www-form-urlencoded"
    return $request.access_token, $request.expires_on
}

$tokenCache = @{}  # Token cache dictionary

function Get-AccessToken {
    $token, $expiresOn = $tokenCache[$tenantId]

    if ($expiresOn -lt (Get-Date)) {
        # Token has expired, retrieve a new one
        $token, $expiresOn = Get-AzureADToken
        $expiresOnDateTime = [DateTimeOffset]::FromUnixTimeSeconds($expiresOn).UtcDateTime

        # Store the new token in the cache
        $tokenCache[$tenantId] = $token, $expiresOnDateTime
    }

    return $token
}

$connection = New-Object System.Data.SqlClient.SqlConnection($connectionString)

# Attempt to get the access token with retries
$maxRetries = 3
$retryDelay = 5000  # milliseconds

for ($retry = 1; $retry -le $maxRetries; $retry++) {
    try {
        $token = Get-AccessToken
        $connection.AccessToken = $token
        $connection.Open()
        break  # Connection successful, exit the retry loop
    }
    catch {
        if ($retry -eq $maxRetries) {
            throw $_  # Maximum retries reached, propagate the error
        }
        else {
            # Retry after a delay
            Start-Sleep -Milliseconds $retryDelay
        }
    }
}

$command = New-Object -TypeName System.Data.SqlClient.SqlCommand
$command.CommandText = "SELECT * FROM sys.databases"
$command.Connection = $connection
$command.ExecuteNonQuery() | Out-Null
$Results = $command.ExecuteReader()

while ($Results.Read()) {
    Write-Output $Results.GetValue(0) $Results.GetValue(1)
}

$Results.Close()
$connection.Close()

Explanation:

This code showcases an approach to obtain and cache Azure AD access tokens while implementing a retry policy for connection attempts. Let's break down the key components:

1. The Get-AzureADToken function retrieves an access token from AAD using client credentials (client ID and client secret).

2. The Get-AccessToken function checks if a valid access token is available in the token cache. If the token has expired, it retrieves a new one using Get-AzureADToken and stores it in the cache.

3. The $tokenCache dictionary serves as a simple in-memory cache for storing access tokens based on the $tenantId.

4. The connection logic attempts to acquire an access token using Get-AccessToken with a configurable retry mechanism. If the maximum number of retries is reached, an error is thrown.

Conclusion:

By implementing token caching and a retry policy in your Azure Active Directory authentication code, you can minimize unnecessary requests to the directory and improve performance.

This blog post provided a PowerShell example demonstrating these techniques, enabling developers to optimize their AAD token retrieval and enhance the reliability of their applications.

Remember to handle and store credentials securely and consider implementing appropriate security measures for token management.

We hope this article helps you optimize your Azure AD authentication flow. Happy coding!

Enjoy!