Add thread safety for KeyVaultChallengeBasedAuthenticationPolicy (#5189)

* Add thread safety for `KeyVaultChallengeBasedAuthenticationPolicy`

* Copy constructor

* Add comment

---------

Co-authored-by: Anton Kolesnyk <antkmsft@users.noreply.github.com>

sdk/core/azure-core/inc/azure/core/http/policies/policy.hpp

@@ -570,6 +570,7 @@ namespace Azure { namespace Core { namespace Http { namespace Policies {
 
       std::unique_ptr<HttpPolicy> Clone() const override
       {
+        // Can't use std::make_shared here because copy constructor is not public.
         return std::unique_ptr<HttpPolicy>(new BearerTokenAuthenticationPolicy(*this));
       }
 

sdk/keyvault/azure-security-keyvault-administration/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- [[#4754]](https://github.com/Azure/azure-sdk-for-cpp/issues/4754) Thread safety for authentication policy.
+
 ### Other Changes
 
 ## 4.0.0-beta.3 (2023-04-11)

sdk/keyvault/azure-security-keyvault-certificates/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- [[#4754]](https://github.com/Azure/azure-sdk-for-cpp/issues/4754) Thread safety for authentication policy.
+
 ### Other Changes
 
 ## 4.2.0 (2023-05-09)

sdk/keyvault/azure-security-keyvault-keys/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- [[#4754]](https://github.com/Azure/azure-sdk-for-cpp/issues/4754) Thread safety for authentication policy.
+
 ### Other Changes
 
 - Fixed GCC 13 compilation error. (A community contribution, courtesy of _[adamdebreceni](https://github.com/adamdebreceni)_)

sdk/keyvault/azure-security-keyvault-secrets/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- [[#4754]](https://github.com/Azure/azure-sdk-for-cpp/issues/4754) Thread safety for authentication policy.
+
 ### Other Changes
 
 ## 4.2.0 (2023-05-09)

sdk/keyvault/azure-security-keyvault-shared/inc/azure/keyvault/shared/keyvault_challenge_based_auth.hpp

@@ -12,6 +12,8 @@
 #include <azure/core/http/policies/policy.hpp>
 #include <azure/core/internal/credentials/authorization_challenge_parser.hpp>
 
+#include <mutex>
+#include <shared_mutex>
 #include <stdexcept>
 
 namespace Azure { namespace Security { namespace KeyVault { namespace _internal {
@@ -23,6 +25,16 @@ namespace Azure { namespace Security { namespace KeyVault { namespace _internal
       : public Core::Http::Policies::_internal::BearerTokenAuthenticationPolicy {
   private:
     mutable Core::Credentials::TokenRequestContext m_tokenRequestContext;
+    mutable std::shared_timed_mutex m_tokenRequestContextMutex;
+
+  protected:
+    KeyVaultChallengeBasedAuthenticationPolicy(
+        KeyVaultChallengeBasedAuthenticationPolicy const& other)
+        : BearerTokenAuthenticationPolicy(other)
+    {
+      std::shared_lock<std::shared_timed_mutex> readLock(other.m_tokenRequestContextMutex);
+      m_tokenRequestContext = other.m_tokenRequestContext;
+    }
 
   public:
     explicit KeyVaultChallengeBasedAuthenticationPolicy(
@@ -35,7 +47,9 @@ namespace Azure { namespace Security { namespace KeyVault { namespace _internal
 
     std::unique_ptr<HttpPolicy> Clone() const override
     {
-      return std::make_unique<KeyVaultChallengeBasedAuthenticationPolicy>(*this);
+      // Can't use std::make_shared here because copy constructor is not public.
+      return std::unique_ptr<KeyVaultChallengeBasedAuthenticationPolicy>(
+          new KeyVaultChallengeBasedAuthenticationPolicy(*this));
     }
 
   private:
@@ -44,7 +58,11 @@ namespace Azure { namespace Security { namespace KeyVault { namespace _internal
         Core::Http::Policies::NextHttpPolicy& nextPolicy,
         Core::Context const& context) const override
     {
-      AuthenticateAndAuthorizeRequest(request, m_tokenRequestContext, context);
+      {
+        std::shared_lock<std::shared_timed_mutex> readLock(m_tokenRequestContextMutex);
+        AuthenticateAndAuthorizeRequest(request, m_tokenRequestContext, context);
+      }
+
       return nextPolicy.Send(request, context);
     }
 
@@ -62,8 +80,11 @@ namespace Azure { namespace Security { namespace KeyVault { namespace _internal
       ValidateChallengeResponse(scope, request.GetUrl().GetHost());
 
       auto const tenantId = GetTenantId(GetAuthorization(challenge));
-      m_tokenRequestContext.TenantId = tenantId;
-      m_tokenRequestContext.Scopes = {scope};
+      {
+        std::unique_lock<std::shared_timed_mutex> writeLock(m_tokenRequestContextMutex);
+        m_tokenRequestContext.TenantId = tenantId;
+        m_tokenRequestContext.Scopes = {scope};
+      }
 
       AuthenticateAndAuthorizeRequest(request, m_tokenRequestContext, context);
       return true;