Merge pull request #9378 from cormacpayne/token-fix

Fix bug with incorrect URL being used in some cases for Functions calls

Author: markcowl

src/Accounts/Accounts/ChangeLog.md

@@ -18,6 +18,8 @@
         - Additional information about change #1
 -->
 ## Upcoming Release
+* Fix bug with incorrect URL being used in some cases for Functions calls
+    - More information here: https://github.com/Azure/azure-powershell/issues/8983
 
 ## Version 1.5.2
 * Update Authentication Library to fix ADFS issues with username/password auth

src/Accounts/Authentication.Test/AuthenticationFactoryTests.cs

@@ -399,5 +399,100 @@ public void AppServiceManagedIdentity()
             Assert.Equal(expectedAccessToken, accessToken);
             Assert.Equal(expectedExpiresOn, msat.ExpiresOn);
         }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void AppServiceManagedIdentityWithDataPlane()
+        {
+            AzureSessionInitializer.InitializeAzureSession();
+            var tenant = Guid.NewGuid().ToString();
+            var userId = Guid.NewGuid().ToString();
+            var environment = AzureEnvironment.PublicEnvironments["AzureCloud"];
+            var account = new AzureAccount
+            {
+                Id = userId,
+                Type = AzureAccount.AccountType.ManagedService
+            };
+            const string resource = @"https://vault.azure.com/";
+            const string endpoint = @"http://127.0.0.1:41217/MSI/token/";
+            var expectedUri = $"{endpoint}?resource={resource}&api-version=2017-09-01";
+            account.SetProperty(AzureAccount.Property.MSILoginUri, endpoint);
+            account.SetProperty(AzureAccount.Property.MSILoginSecret, @"bar");
+            const string expectedAccessToken = "foo";
+            var expectedExpiresOn = DateTimeOffset.Parse("1/23/2019 7:15:42 AM +00:00");
+            var responses = new Dictionary<string, ManagedServiceAppServiceTokenInfo>(StringComparer.OrdinalIgnoreCase)
+            {
+                {
+                    expectedUri,
+                    new ManagedServiceAppServiceTokenInfo()
+                    {
+                        AccessToken = expectedAccessToken,
+                        ExpiresOn = expectedExpiresOn,
+                        Resource = resource,
+                        TokenType = "Bearer",
+                    }
+                }
+            };
+            AzureSession.Instance.RegisterComponent(HttpClientOperationsFactory.Name, () => TestHttpOperationsFactory.Create(responses, _output), true);
+            var msat = new ManagedServiceAppServiceAccessToken(account, environment, environment.GetEndpoint(resource) ?? resource, tenant);
+            Assert.Equal(expectedUri, msat.RequestUris.Peek());
+            var accessToken = msat.AccessToken;
+            Assert.Equal(expectedAccessToken, accessToken);
+            Assert.Equal(expectedExpiresOn, msat.ExpiresOn);
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void AppServiceManagedIdentityWithServiceManagement()
+        {
+            AzureSessionInitializer.InitializeAzureSession();
+            var tenant = Guid.NewGuid().ToString();
+            var userId = Guid.NewGuid().ToString();
+            var environment = AzureEnvironment.PublicEnvironments["AzureCloud"];
+            var account = new AzureAccount
+            {
+                Id = userId,
+                Type = AzureAccount.AccountType.ManagedService
+            };
+            const string resource = @"https://management.azure.com/";
+            const string serviceManagementResource = @"https://management.core.windows.net/";
+            const string endpoint = @"http://127.0.0.1:41217/MSI/token/";
+            var expectedUri = $"{endpoint}?resource={resource}&api-version=2017-09-01";
+            account.SetProperty(AzureAccount.Property.MSILoginUri, endpoint);
+            account.SetProperty(AzureAccount.Property.MSILoginSecret, @"bar");
+            const string expectedAccessToken = "foo";
+            var expectedExpiresOn = DateTimeOffset.Parse("1/23/2019 7:15:42 AM +00:00");
+            var responses = new Dictionary<string, ManagedServiceAppServiceTokenInfo>(StringComparer.OrdinalIgnoreCase)
+            {
+                {
+                    expectedUri,
+                    new ManagedServiceAppServiceTokenInfo()
+                    {
+                        AccessToken = expectedAccessToken,
+                        ExpiresOn = expectedExpiresOn,
+                        Resource = resource,
+                        TokenType = "Bearer",
+                    }
+                }
+            };
+            AzureSession.Instance.RegisterComponent(HttpClientOperationsFactory.Name, () => TestHttpOperationsFactory.Create(responses, _output), true);
+            var msat = new ManagedServiceAppServiceAccessToken(account, environment, GetFunctionsResourceId(serviceManagementResource, environment), tenant);
+            Assert.Equal(expectedUri, msat.RequestUris.Peek());
+            var accessToken = msat.AccessToken;
+            Assert.Equal(expectedAccessToken, accessToken);
+            Assert.Equal(expectedExpiresOn, msat.ExpiresOn);
+        }
+        private string GetFunctionsResourceId(string resourceIdOrEndpointName, IAzureEnvironment environment)
+        {
+            var resourceId = environment.GetEndpoint(resourceIdOrEndpointName) ?? resourceIdOrEndpointName;
+            if (string.Equals(
+                environment.GetEndpoint(AzureEnvironment.Endpoint.ActiveDirectoryServiceEndpointResourceId),
+                resourceId, StringComparison.OrdinalIgnoreCase))
+            {
+                resourceId = environment.GetEndpoint(AzureEnvironment.Endpoint.ResourceManager);
+            }
+
+            return resourceId;
+        }
     }
 }

src/Accounts/Authentication/Authentication/ManagedServiceAppServiceAccessToken.cs

@@ -24,6 +24,11 @@ public ManagedServiceAppServiceAccessToken(IAzureAccount account, IAzureEnvironm
         {
         }
 
+        public ManagedServiceAppServiceAccessToken(IAzureAccount account, IAzureEnvironment environment, string resourceId, string tenant = "Common")
+            : base(account, environment, resourceId, tenant)
+        {
+        }
+
         protected override IEnumerable<string> BuildTokenUri(string baseUri, IAzureAccount account, IdentityType identityType,
             string resourceId)
         {

src/Accounts/Authentication/Factories/AuthenticationFactory.cs

@@ -427,7 +427,7 @@ private IAccessToken GetManagedServiceToken(IAzureAccount account, IAzureEnviron
 
             if (account.IsPropertySet(AuthenticationFactory.AppServiceManagedIdentityFlag))
             {
-                return new ManagedServiceAppServiceAccessToken(account, environment, tenant);
+                return new ManagedServiceAppServiceAccessToken(account, environment, GetFunctionsResourceId(resourceId, environment), tenant);
             }
 
             return new ManagedServiceAccessToken(account, environment, GetResourceId(resourceId, environment), tenant);
@@ -438,6 +438,19 @@ private string GetResourceId(string resourceIdorEndpointName, IAzureEnvironment
             return environment.GetEndpoint(resourceIdorEndpointName) ?? resourceIdorEndpointName;
         }
 
+        private string GetFunctionsResourceId(string resourceIdOrEndpointName, IAzureEnvironment environment)
+        {
+            var resourceId = environment.GetEndpoint(resourceIdOrEndpointName) ?? resourceIdOrEndpointName;
+            if (string.Equals(
+                environment.GetEndpoint(AzureEnvironment.Endpoint.ActiveDirectoryServiceEndpointResourceId),
+                resourceId, StringComparison.OrdinalIgnoreCase))
+            {
+                resourceId = environment.GetEndpoint(AzureEnvironment.Endpoint.ResourceManager);
+            }
+
+            return resourceId;
+        }
+
         private AdalConfiguration GetAdalConfiguration(IAzureEnvironment environment, string tenantId,
             string resourceId, TokenCache tokenCache)
         {