Add oAuthEndpoint to _oAuthClients cache key (#6298) (#6318)

* Add oAuthEndpoint to _oAuthClients cache key * Add CreateOAuthClientWithDifferentEndpoints test Co-authored-by: Craig Jensen <crjens@hotmail.com> Co-authored-by: craigjensen <cjensen@microsoft.com> Co-authored-by: Craig Jensen <crjens@hotmail.com>
microsoft · May 3, 2022 · 43be5b0 · 43be5b0
1 parent 87f9413
commit 43be5b0
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 3 deletions.
diff --git a/libraries/Microsoft.Bot.Builder/BotFrameworkAdapter.cs b/libraries/Microsoft.Bot.Builder/BotFrameworkAdapter.cs
@@ -63,7 +63,7 @@ public class BotFrameworkAdapter : BotAdapter, IAdapterIntegration, IExtendedUse
         private readonly ConcurrentDictionary<string, ConnectorClient> _connectorClients = new ConcurrentDictionary<string, ConnectorClient>();
 
         // Cache for OAuthClient to speed up OAuth operations
-        // _oAuthClients is a cache using [appId + oAuthCredentialAppId]
+        // _oAuthClients is a cache using [appId + oAuthCredentialAppId + oAuthEndpoint]
         private readonly ConcurrentDictionary<string, OAuthClient> _oAuthClients = new ConcurrentDictionary<string, OAuthClient>();
 
         /// <summary>
@@ -1395,7 +1395,6 @@ protected virtual async Task<OAuthClient> CreateOAuthApiClientAsync(ITurnContext
 
             var appId = GetBotAppId(turnContext);
 
-            var clientKey = $"{appId}:{oAuthAppCredentials?.MicrosoftAppId}";
             var oAuthScope = GetBotFrameworkOAuthScope();
 
             var appCredentials = oAuthAppCredentials ?? await GetAppCredentialsAsync(appId, oAuthScope).ConfigureAwait(false);
@@ -1407,6 +1406,8 @@ protected virtual async Task<OAuthClient> CreateOAuthApiClientAsync(ITurnContext
                 OAuthClientConfig.EmulateOAuthCards = true;
             }
 
+            var oAuthEndpoint = OAuthClientConfig.OAuthEndpoint;
+            var clientKey = $"{appId}:{oAuthAppCredentials?.MicrosoftAppId}:{oAuthEndpoint}";
             var oAuthClient = _oAuthClients.GetOrAdd(clientKey, (key) =>
             {
                 OAuthClient oAuthClientInner;
@@ -1418,7 +1419,7 @@ protected virtual async Task<OAuthClient> CreateOAuthApiClientAsync(ITurnContext
                 }
                 else
                 {
-                    oAuthClientInner = new OAuthClient(new Uri(OAuthClientConfig.OAuthEndpoint), appCredentials);
+                    oAuthClientInner = new OAuthClient(new Uri(oAuthEndpoint), appCredentials);
                 }
 
                 return oAuthClientInner;

diff --git a/tests/Microsoft.Bot.Builder.Tests/BotFrameworkAdapterOAuthTests.cs b/tests/Microsoft.Bot.Builder.Tests/BotFrameworkAdapterOAuthTests.cs
@@ -20,6 +20,47 @@ namespace Microsoft.Bot.Builder.Tests
 {
     public class BotFrameworkAdapterOAuthTests
     {
+        [Fact]
+        public async Task CreateOAuthClientWithDifferentEndpoints()
+        {
+            var mockAppCredentials = new MockAppCredentials();
+            var mockCredentialProvider = new Mock<ICredentialProvider>();
+            var adapter = new MockAdapter(mockCredentialProvider.Object, () => new TokenResponse());
+
+            var claimsIdentity = new ClaimsIdentity();
+            claimsIdentity.AddClaim(new Claim(AuthenticationConstants.AudienceClaim, "AppId"));
+
+            var turnStateCollection = new TurnContextStateCollection();
+            turnStateCollection.Add(BotAdapter.BotIdentityKey, claimsIdentity);
+
+            var turnContext = new Mock<ITurnContext>();
+            turnContext.SetupGet(x => x.Activity).Returns(new Activity());
+            turnContext.SetupGet(x => x.TurnState).Returns(turnStateCollection);
+
+            var oauthEndpoint1 = "https://foo.com";
+            OAuthClientConfig.OAuthEndpoint = oauthEndpoint1;
+            var client = await adapter.CallCreateOAuthApiClientAsync(turnContext.Object, mockAppCredentials);
+            Assert.NotNull(client);
+            Assert.Equal(new Uri(oauthEndpoint1), client.BaseUri);
+            Assert.Same(mockAppCredentials, client.Credentials);
+
+            // client2 should come from the cache so it should be the same object
+            var client2 = await adapter.CallCreateOAuthApiClientAsync(turnContext.Object, mockAppCredentials);
+            Assert.NotNull(client2);
+            Assert.Same(client, client2);  
+            Assert.Equal(new Uri(oauthEndpoint1), client2.BaseUri);
+            Assert.Same(mockAppCredentials, client2.Credentials);
+
+            // Changing the OAuthEndpoint should result in a different OAuthClient
+            var oauthEndpoint2 = "https://bar.com";
+            OAuthClientConfig.OAuthEndpoint = oauthEndpoint2;
+            var client3 = await adapter.CallCreateOAuthApiClientAsync(turnContext.Object, mockAppCredentials);
+            Assert.NotNull(client3);
+            Assert.NotSame(client3, client);
+            Assert.Equal(new Uri(oauthEndpoint2), client3.BaseUri);
+            Assert.Same(mockAppCredentials, client3.Credentials);
+        }
+
         [Fact]
         public async Task BotFrameworkAdapterFindsOAuthCards()
         {
@@ -345,6 +386,11 @@ public override Task<TokenResponse> GetUserTokenAsync(ITurnContext turnContext,
             {
                 return Task.FromResult(_getUserTokenAction());
             }
+
+            public Task<OAuthClient> CallCreateOAuthApiClientAsync(ITurnContext turnContext, AppCredentials oAuthAppCredentials)
+            {
+                return CreateOAuthApiClientAsync(turnContext, oAuthAppCredentials);
+            }
         }
 
         private class MockLogger : ILogger<BotFrameworkHttpAdapterBase>

diff --git a/tests/Microsoft.Bot.Builder.Tests/MockAppCredentials.cs b/tests/Microsoft.Bot.Builder.Tests/MockAppCredentials.cs
@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Net.Http;
+using Microsoft.Bot.Connector.Authentication;
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.Bot.Builder.Tests
+{
+    public class MockAppCredentials : AppCredentials
+    {
+        public MockAppCredentials(string channelAuthTenant = null, HttpClient customHttpClient = null, ILogger logger = null)
+            : base(channelAuthTenant, customHttpClient, logger)
+        {
+        }
+
+        protected override Lazy<AdalAuthenticator> BuildAuthenticator()
+        {
+            return new Lazy<AdalAuthenticator>();
+        }
+    }
+}