Skip to content

TokenCredentialCache: Adds a fallback mechanism to AAD scope override. #5364

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kirankumarkolli merged 15 commits into from
Aug 26, 2025
Merged

TokenCredentialCache: Adds a fallback mechanism to AAD scope override. #5364

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
5b957a2
Add fallback mechanism for scope override.
aavasthy Aug 19, 2025
ba50af9
Merge branch 'master' into users/aavasthy/fallback_aadscope
aavasthy Aug 19, 2025
a7d3c64
Fix formatting.
aavasthy Aug 19, 2025
889fd4d
Update the fallback logic
aavasthy Aug 21, 2025
3c8d3d3
Fux formatting.
aavasthy Aug 21, 2025
90029c0
Code cleanup.
aavasthy Aug 21, 2025
0e709ad
Code cleanup.
aavasthy Aug 22, 2025
76e7d9e
Merge branch 'master' into users/aavasthy/fallback_aadscope
aavasthy Aug 22, 2025
c48e1fb
Added updated tests for cosmos scope provider.
aavasthy Aug 22, 2025
74b0311
Code cleanup.
aavasthy Aug 23, 2025
4f5a017
Code cleanup
aavasthy Aug 25, 2025
481e76b
nit updates based on review.
aavasthy Aug 25, 2025
afcf4bd
Update logging for fallback.
aavasthy Aug 25, 2025
19c526f
Update logging for AAD fallback.
aavasthy Aug 25, 2025
5548870
Update logging for AAD fallback.
aavasthy Aug 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Code cleanup.
  • Loading branch information
@aavasthy
aavasthy committed Aug 22, 2025
commit 0e709adddf27f77b6d8f6639e979fb243cb0b5b2
10 changes: 6 additions & 4 deletions Microsoft.Azure.Cosmos/src/Authorization/CosmosScopeProvider.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ public TokenRequestContext GetTokenRequestContext()
return new TokenRequestContext(new[] { this.currentScope });
}

public bool TryFallback(Exception ex)
public bool TryFallback(Exception exception)
{
// If override scope is set, never fallback
if (!string.IsNullOrEmpty(this.overrideScope))
Expand All @@ -40,16 +40,18 @@ public bool TryFallback(Exception ex)
}

// If already using fallback scope, do not fallback again
if (this.currentScope == AadDefaultScope)
if (this.currentScope == CosmosScopeProvider.AadDefaultScope)
{
return false;
}

if (ex.InnerException?.Message.Contains(AadInvalidScopeErrorMessage) == true)
#pragma warning disable CDX1003 // DontUseExceptionToString
if (exception.ToString().Contains(CosmosScopeProvider.AadInvalidScopeErrorMessage) == true)
{
this.currentScope = AadDefaultScope;
this.currentScope = CosmosScopeProvider.AadDefaultScope;
return true;
}
#pragma warning restore CDX1003 // DontUseExceptionToString

return false;
}
Expand Down
20 changes: 11 additions & 9 deletions Microsoft.Azure.Cosmos/src/Authorization/TokenCredentialCache.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ internal sealed class TokenCredentialCache : IDisposable
private Task<AccessToken>? currentRefreshOperation = null;
private AccessToken? cachedAccessToken = null;
private bool isBackgroundTaskRunning = false;
private bool isDisposed = false;
private bool isDisposed = false;

internal TokenCredentialCache(
TokenCredential tokenCredential,
Expand Down Expand Up @@ -163,11 +163,13 @@ private async Task<AccessToken> GetNewTokenAsync(

private async ValueTask<AccessToken> RefreshCachedTokenWithRetryHelperAsync(
ITrace trace)
{
{
try
{
Exception? lastException = null;
const int totalRetryCount = 2;
const int totalRetryCount = 2;
TokenRequestContext tokenRequestContext = default;

for (int retry = 0; retry < totalRetryCount; retry++)
{
if (this.cancellationToken.IsCancellationRequested)
Expand All @@ -182,10 +184,10 @@ private async ValueTask<AccessToken> RefreshCachedTokenWithRetryHelperAsync(
name: nameof(this.RefreshCachedTokenWithRetryHelperAsync),
component: TraceComponent.Authorization,
level: Tracing.TraceLevel.Info))
{
{
try
{
TokenRequestContext tokenRequestContext = this.scopeProvider.GetTokenRequestContext();
tokenRequestContext = this.scopeProvider.GetTokenRequestContext();

this.cachedAccessToken = await this.tokenCredential.GetTokenAsync(
requestContext: tokenRequestContext,
Expand Down Expand Up @@ -214,13 +216,13 @@ private async ValueTask<AccessToken> RefreshCachedTokenWithRetryHelperAsync(
return this.cachedAccessToken.Value;
}
catch (RequestFailedException requestFailedException)
{
{
lastException = requestFailedException;
getTokenTrace.AddDatum(
$"RequestFailedException at {DateTime.UtcNow.ToString(CultureInfo.InvariantCulture)}",
requestFailedException.Message);

DefaultTrace.TraceError($"TokenCredential.GetToken() failed with RequestFailedException. scope = {string.Join(";", this.scopeProvider.GetTokenRequestContext().Scopes)}, retry = {retry}, Exception = {lastException.Message}");
DefaultTrace.TraceError($"TokenCredential.GetToken() failed with RequestFailedException. scope = {string.Join(";", tokenRequestContext.Scopes ?? Array.Empty<string>())}, retry = {retry}, Exception = {lastException.Message}");

// Don't retry on auth failures
if (requestFailedException.Status == (int)HttpStatusCode.Unauthorized ||
Expand All @@ -244,7 +246,7 @@ private async ValueTask<AccessToken> RefreshCachedTokenWithRetryHelperAsync(
operationCancelled.Message);

DefaultTrace.TraceError(
$"TokenCredential.GetTokenAsync() failed. scope = {string.Join(";", this.scopeProvider.GetTokenRequestContext().Scopes)}, retry = {retry}, Exception = {lastException.Message}");
$"TokenCredential.GetTokenAsync() failed. scope = {string.Join(";", tokenRequestContext.Scopes ?? Array.Empty<string>())}, retry = {retry}, Exception = {lastException.Message}");

throw CosmosExceptionFactory.CreateRequestTimeoutException(
message: ClientResources.FailedToGetAadToken,
Expand All @@ -263,7 +265,7 @@ private async ValueTask<AccessToken> RefreshCachedTokenWithRetryHelperAsync(
exception.Message);

DefaultTrace.TraceError(
$"TokenCredential.GetTokenAsync() failed. scope = {string.Join(";", this.scopeProvider.GetTokenRequestContext().Scopes)}, retry = {retry}, Exception = {lastException.Message}");
$"TokenCredential.GetTokenAsync() failed. scope = {string.Join(";", tokenRequestContext.Scopes ?? Array.Empty<string>())}, retry = {retry}, Exception = {lastException.Message}");

// Fallback logic
if (this.scopeProvider.TryFallback(exception))
Expand Down