Fix concurrency issue with challenge auth policy

sdk/keyvault/Azure.Security.KeyVault.Certificates/CHANGELOG.md

@@ -1,5 +1,11 @@
 # Release History
 
+## 4.0.2 (2020-03-18)
+
+### Fixed
+
+- Fixed concurrency issue in our challenge-based authentication policy ([#9737](https://github.com/Azure/azure-sdk-for-net/issues/9737))
+
 ## 4.0.1 (2020-03-03)
 
 ### Fixed

sdk/keyvault/Azure.Security.KeyVault.Certificates/src/Azure.Security.KeyVault.Certificates.csproj

@@ -3,7 +3,7 @@
   <PropertyGroup>
     <Description>This is the Microsoft Azure Key Vault Certificates client library</Description>
     <AssemblyTitle>Microsoft Azure.Security.KeyVault.Certificates client library</AssemblyTitle>
-    <Version>4.0.1</Version>
+    <Version>4.0.2</Version>
     <PackageTags>Microsoft Azure Key Vault Certificates;$(PackageCommonTags)</PackageTags>
     <TargetFrameworks>$(RequiredTargetFrameworks)</TargetFrameworks>
     <NoWarn>$(NoWarn);3021</NoWarn>

sdk/keyvault/Azure.Security.KeyVault.Keys/CHANGELOG.md

@@ -1,5 +1,11 @@
 # Release History
 
+## 4.0.3 (2020-03-18)
+
+### Fixed
+
+- Fixed concurrency issue in our challenge-based authentication policy ([#9737](https://github.com/Azure/azure-sdk-for-net/issues/9737))
+
 ## 4.0.2 (2020-03-03)
 
 ### Fixed

sdk/keyvault/Azure.Security.KeyVault.Keys/src/Azure.Security.KeyVault.Keys.csproj

@@ -3,7 +3,7 @@
   <PropertyGroup>
     <Description>This is the Microsoft Azure Key Vault Keys client library</Description>
     <AssemblyTitle>Microsoft Azure.Security.KeyVault.Keys client library</AssemblyTitle>
-    <Version>4.0.2</Version>
+    <Version>4.0.3</Version>
     <PackageTags>Microsoft Azure Key Vault Keys;$(PackageCommonTags)</PackageTags>
 
     <TargetFrameworks>$(RequiredTargetFrameworks)</TargetFrameworks>

sdk/keyvault/Azure.Security.KeyVault.Secrets/CHANGELOG.md

@@ -1,5 +1,11 @@
 # Release History
 
+## 4.0.3 (2020-03-18)
+
+### Fixed
+
+- Fixed concurrency issue in our challenge-based authentication policy ([#9737](https://github.com/Azure/azure-sdk-for-net/issues/9737))
+
 ## 4.0.2 (2020-03-03)
 
 ### Fixed

sdk/keyvault/Azure.Security.KeyVault.Secrets/src/Azure.Security.KeyVault.Secrets.csproj

@@ -3,7 +3,7 @@
   <PropertyGroup>
     <Description>This is the Microsoft Azure Key Vault Secrets client library</Description>
     <AssemblyTitle>Microsoft Azure.Security.KeyVault.Secrets client library</AssemblyTitle>
-    <Version>4.0.2</Version>
+    <Version>4.0.3</Version>
     <PackageTags>Microsoft Azure Key Vault Secrets;$(PackageCommonTags)</PackageTags>
 
     <TargetFrameworks>$(RequiredTargetFrameworks)</TargetFrameworks>

sdk/keyvault/Azure.Security.KeyVault.Shared/src/ChallengeBasedAuthenticationPolicy.cs

@@ -12,12 +12,11 @@ namespace Azure.Security.KeyVault
     internal class ChallengeBasedAuthenticationPolicy : HttpPipelinePolicy
     {
         private const string BearerChallengePrefix = "Bearer ";
+
         private readonly TokenCredential _credential;
 
         private AuthenticationChallenge _challenge = null;
-
         private string _headerValue;
-
         private DateTimeOffset _refreshOn;
 
         public ChallengeBasedAuthenticationPolicy(TokenCredential credential)
@@ -45,18 +44,18 @@ private async ValueTask ProcessCoreAsync(HttpMessage message, ReadOnlyMemory<Htt
             RequestContent originalContent = message.Request.Content;
 
             // if this policy doesn't have _challenge cached try to get it from the static challenge cache
-            _challenge ??= AuthenticationChallenge.GetChallenge(message);
+            AuthenticationChallenge challenge = _challenge ?? AuthenticationChallenge.GetChallenge(message);
 
             // if we still don't have the challenge for the endpoint
             // remove the content from the request and send without authentication to get the challenge
-            if (_challenge == null)
+            if (challenge == null)
             {
                 message.Request.Content = null;
             }
             // otherwise if we already know the challenge authenticate the request
             else
             {
-                await AuthenticateRequestAsync(message, async).ConfigureAwait(false);
+                await AuthenticateRequestAsync(message, async, challenge).ConfigureAwait(false);
             }
 
             if (async)
@@ -75,16 +74,18 @@ private async ValueTask ProcessCoreAsync(HttpMessage message, ReadOnlyMemory<Htt
                 message.Request.Content = originalContent;
 
                 // update the cached challenge
-                var challenge = AuthenticationChallenge.GetChallenge(message);
+                challenge = AuthenticationChallenge.GetChallenge(message);
 
-                // if a challenge was returned and its different from the cached _challenge
-                if (challenge != null && !challenge.Equals(_challenge))
+                if (challenge != null)
                 {
-                    // update the cached challenge
-                    _challenge = challenge;
+                    // update the cached challenge if not yet set or different from the current challenge (e.g. moved tenants)
+                    if (_challenge == null || !challenge.Equals(_challenge))
+                    {
+                        _challenge = challenge;
+                    }
 
                     // authenticate the request and resend
-                    await AuthenticateRequestAsync(message, async).ConfigureAwait(false);
+                    await AuthenticateRequestAsync(message, async, challenge).ConfigureAwait(false);
 
                     if (async)
                     {
@@ -98,13 +99,13 @@ private async ValueTask ProcessCoreAsync(HttpMessage message, ReadOnlyMemory<Htt
             }
         }
 
-        private async Task AuthenticateRequestAsync(HttpMessage message, bool async)
+        private async Task AuthenticateRequestAsync(HttpMessage message, bool async, AuthenticationChallenge challenge)
         {
-            if (DateTimeOffset.UtcNow >= _refreshOn)
+            if (_headerValue is null || DateTimeOffset.UtcNow >= _refreshOn)
             {
                 AccessToken token = async ?
-                        await _credential.GetTokenAsync(new TokenRequestContext(_challenge.Scopes), message.CancellationToken).ConfigureAwait(false) :
-                        _credential.GetToken(new TokenRequestContext(_challenge.Scopes), message.CancellationToken);
+                        await _credential.GetTokenAsync(new TokenRequestContext(challenge.Scopes), message.CancellationToken).ConfigureAwait(false) :
+                        _credential.GetToken(new TokenRequestContext(challenge.Scopes), message.CancellationToken);
 
                 _headerValue = "Bearer " + token.Token;
                 _refreshOn = token.ExpiresOn - TimeSpan.FromMinutes(2);
@@ -115,8 +116,8 @@ await _credential.GetTokenAsync(new TokenRequestContext(_challenge.Scopes), mess
 
         internal class AuthenticationChallenge
         {
-            private static readonly Dictionary<string, AuthenticationChallenge> _cache = new Dictionary<string, AuthenticationChallenge>();
-            private static readonly object _cacheLock = new object();
+            private static readonly Dictionary<string, AuthenticationChallenge> s_cache = new Dictionary<string, AuthenticationChallenge>();
+            private static readonly object s_cacheLock = new object();
 
             private AuthenticationChallenge(string scope)
             {
@@ -127,7 +128,7 @@ private AuthenticationChallenge(string scope)
 
             public override bool Equals(object obj)
             {
-                if (base.Equals(obj))
+                if (ReferenceEquals(this, obj))
                 {
                     return true;
                 }
@@ -161,18 +162,20 @@ public static AuthenticationChallenge GetChallenge(HttpMessage message)
                     // if the challenge is non-null cache it
                     if (challenge != null)
                     {
-                        lock (_cacheLock)
+                        string authority = GetRequestAuthority(message.Request);
+                        lock (s_cacheLock)
                         {
-                            _cache[GetRequestAuthority(message.Request)] = challenge;
+                            s_cache[authority] = challenge;
                         }
                     }
                 }
                 else
                 {
                     // try to get the challenge from the cache
-                    lock (_cacheLock)
+                    string authority = GetRequestAuthority(message.Request);
+                    lock (s_cacheLock)
                     {
-                        _cache.TryGetValue(GetRequestAuthority(message.Request), out challenge);
+                        s_cache.TryGetValue(authority, out challenge);
                     }
                 }
 
@@ -182,9 +185,9 @@ public static AuthenticationChallenge GetChallenge(HttpMessage message)
             internal static void ClearCache()
             {
                 // try to get the challenge from the cache
-                lock (_cacheLock)
+                lock (s_cacheLock)
                 {
-                    _cache.Clear();
+                    s_cache.Clear();
                 }
             }