WX-1594 Support for private ACR access in Terra (#715)

Co-authored-by: Blair L Murri <[email protected]>

src/CommonUtilities/Models/NodeTask.cs

@@ -46,6 +46,7 @@ public class RuntimeOptions
         public TerraRuntimeOptions? Terra { get; set; }
 
         public string? NodeManagedIdentityResourceId { get; set; }
+        public string? AcrPullManagedIdentityResourceId { get; set; }
 
         public StorageTargetLocation? StorageEventSink { get; set; }
 

src/Tes.ApiClients.Tests/TerraApiStubData.cs

@@ -10,6 +10,7 @@ public class TerraApiStubData
 {
     public const string LandingZoneApiHost = "https://landingzone.host";
     public const string WsmApiHost = "https://wsm.host";
+    public const string SamApiHost = "https://sam.host";
     public const string ResourceGroup = "mrg-terra-dev-previ-20191228";
     public const string WorkspaceAccountName = "lzaccount1";
     public const string SasToken = "SASTOKENSTUB=";
@@ -18,8 +19,12 @@ public class TerraApiStubData
     public const string WorkspaceStorageContainerName = $"sc-{WorkspaceIdValue}";
     public const string WsmGetSasResponseStorageUrl = $"https://{WorkspaceAccountName}.blob.core.windows.net/{WorkspaceStorageContainerName}";
 
+    public const string TerraPetName = "pet-2674060218359759651b0";
+
+    public Guid TenantId { get; } = Guid.NewGuid();
     public Guid LandingZoneId { get; } = Guid.NewGuid();
     public Guid SubscriptionId { get; } = Guid.NewGuid();
+    public Guid AcrPullIdentitySamResourceId { get; } = Guid.NewGuid();
     public Guid ContainerResourceId { get; } = Guid.NewGuid();
     public Guid WorkspaceId { get; } = Guid.Parse(WorkspaceIdValue);
 
@@ -28,6 +33,8 @@ public class TerraApiStubData
     public string BatchAccountId =>
         $"/subscriptions/{SubscriptionId}/resourceGroups/{ResourceGroup}/providers/Microsoft.Batch/batchAccounts/{BatchAccountName}";
 
+    public string ManagedIdentityObjectId =>
+        $"/subscriptions/{SubscriptionId}/resourceGroups/{ResourceGroup}/providers/Microsoft.ManagedIdentity/userAssignedIdentities/{TerraPetName}";
     public string PoolId => "poolId";
 
     public Guid GetWorkspaceIdFromContainerName(string containerName)
@@ -279,6 +286,27 @@ public string GetResourceQuotaApiResponseInJson()
 }}";
     }
 
+    public string GetSamActionManagedIdentityApiResponseInJson()
+    {
+        return $@"{{
+  ""id"": {{
+    ""resourceId"": {{
+      ""resourceTypeName"": ""private_azure_container_registry"",
+      ""resourceId"": ""{AcrPullIdentitySamResourceId}""
+    }},
+    ""action"": ""pull_image"",
+    ""billingProfileId"": ""{AcrPullIdentitySamResourceId}""
+  }},
+  ""objectId"": ""{ManagedIdentityObjectId}"",
+  ""displayName"": ""my nice action identity"",
+  ""managedResourceGroupCoordinates"": {{
+    ""tenantId"": ""{TenantId}"",
+    ""subscriptionId"": ""{SubscriptionId}"",
+    ""managedResourceGroupName"": ""{ResourceGroup}""
+  }}
+}}";
+    }
+
     public ApiCreateBatchPoolRequest GetApiCreateBatchPoolRequest()
     {
         return new ApiCreateBatchPoolRequest()

src/Tes.ApiClients.Tests/TerraSamApiClientTests.cs

@@ -0,0 +1,90 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Azure.Core;
+using CommonUtilities;
+using Microsoft.Extensions.Logging.Abstractions;
+using Moq;
+using Tes.ApiClients.Models.Terra;
+
+namespace Tes.ApiClients.Tests
+{
+    [TestClass, TestCategory("Unit")]
+    public class TerraSamApiClientTests
+    {
+        private TerraSamApiClient terraSamApiClient = null!;
+        private Mock<TokenCredential> tokenCredential = null!;
+        private Mock<CachingRetryPolicyBuilder> cacheAndRetryBuilder = null!;
+        private Lazy<Mock<CachingRetryHandler.CachingAsyncRetryHandlerPolicy<HttpResponseMessage>>> cacheAndRetryHandler = null!;
+        private TerraApiStubData terraApiStubData = null!;
+        private AzureEnvironmentConfig azureEnvironmentConfig = null!;
+        private TimeSpan cacheTTL = TimeSpan.FromMinutes(1);
+
+        [TestInitialize]
+        public void SetUp()
+        {
+            terraApiStubData = new TerraApiStubData();
+            tokenCredential = new Mock<TokenCredential>();
+            cacheAndRetryBuilder = new Mock<CachingRetryPolicyBuilder>();
+            var cache = new Mock<Microsoft.Extensions.Caching.Memory.IMemoryCache>();
+            cache.Setup(c => c.CreateEntry(It.IsAny<object>())).Returns(new Mock<Microsoft.Extensions.Caching.Memory.ICacheEntry>().Object);
+            cacheAndRetryBuilder.SetupGet(c => c.AppCache).Returns(cache.Object);
+            cacheAndRetryHandler = new(TestServices.RetryHandlersHelpers.GetCachingAsyncRetryPolicyMock(cacheAndRetryBuilder, c => c.DefaultRetryHttpResponseMessagePolicyBuilder()));
+            azureEnvironmentConfig = ExpensiveObjectTestUtility.AzureCloudConfig.AzureEnvironmentConfig!;
+
+            terraSamApiClient = new TerraSamApiClient(TerraApiStubData.SamApiHost, tokenCredential.Object,
+                cacheAndRetryBuilder.Object, azureEnvironmentConfig, NullLogger<TerraSamApiClient>.Instance);
+        }
+
+        [TestMethod]
+        public async Task GetActionManagedIdentityAsync_ValidRequest_ReturnsPayload()
+        {
+            cacheAndRetryHandler.Value.Setup(c => c.ExecuteWithRetryConversionAndCachingAsync(
+                    It.IsAny<string>(),
+                    It.IsAny<Func<CancellationToken, Task<HttpResponseMessage>>>(),
+                    It.IsAny<Func<HttpResponseMessage, CancellationToken, Task<SamActionManagedIdentityApiResponse>>>(),
+                    It.IsAny<DateTimeOffset>(),
+                    It.IsAny<CancellationToken>(),
+                    It.IsAny<string>()))
+                .ReturnsAsync(System.Text.Json.JsonSerializer.Deserialize<SamActionManagedIdentityApiResponse>(terraApiStubData.GetSamActionManagedIdentityApiResponseInJson())!);
+
+            var apiResponse = await terraSamApiClient.GetActionManagedIdentityForACRPullAsync(terraApiStubData.AcrPullIdentitySamResourceId, cacheTTL, CancellationToken.None);
+
+            Assert.IsNotNull(apiResponse);
+            Assert.IsTrue(!string.IsNullOrEmpty(apiResponse.ObjectId));
+            Assert.IsTrue(apiResponse.ObjectId.Contains(TerraApiStubData.TerraPetName));
+        }
+
+        [TestMethod]
+        public async Task GetActionManagedIdentityAsync_ValidRequest_Returns404()
+        {
+            cacheAndRetryHandler.Value.Setup(c => c.ExecuteWithRetryConversionAndCachingAsync(
+                    It.IsAny<string>(),
+                    It.IsAny<Func<CancellationToken, Task<HttpResponseMessage>>>(),
+                    It.IsAny<Func<HttpResponseMessage, CancellationToken, Task<SamActionManagedIdentityApiResponse>>>(),
+                    It.IsAny<DateTimeOffset>(),
+                    It.IsAny<CancellationToken>(),
+                    It.IsAny<string>()))
+                .Throws(new HttpRequestException(null, null, System.Net.HttpStatusCode.NotFound));
+
+            var apiResponse = await terraSamApiClient.GetActionManagedIdentityForACRPullAsync(terraApiStubData.AcrPullIdentitySamResourceId, cacheTTL, CancellationToken.None);
+
+            Assert.IsNull(apiResponse);
+        }
+
+        [TestMethod]
+        public async Task GetActionManagedIdentityAsync_ValidRequest_Returns500()
+        {
+            cacheAndRetryHandler.Value.Setup(c => c.ExecuteWithRetryConversionAndCachingAsync(
+                    It.IsAny<string>(),
+                    It.IsAny<Func<CancellationToken, Task<HttpResponseMessage>>>(),
+                    It.IsAny<Func<HttpResponseMessage, CancellationToken, Task<SamActionManagedIdentityApiResponse>>>(),
+                    It.IsAny<DateTimeOffset>(),
+                    It.IsAny<CancellationToken>(),
+                    It.IsAny<string>()))
+                .Throws(new HttpRequestException(null, null, System.Net.HttpStatusCode.BadGateway));
+
+            await Assert.ThrowsExceptionAsync<HttpRequestException>(async () => await terraSamApiClient.GetActionManagedIdentityForACRPullAsync(terraApiStubData.AcrPullIdentitySamResourceId, cacheTTL, CancellationToken.None));
+        }
+    }
+}

src/Tes.ApiClients/HttpApiClient.cs

@@ -162,6 +162,32 @@ protected async Task<TResponse> HttpGetRequestWithCachingAndRetryPolicyAsync<TRe
             (m, ct) => GetApiResponseContentAsync(m, typeInfo, ct), cancellationToken))!;
         }
 
+        /// <summary>
+        /// Checks the cache and if the request was not found, sends the GET request with a retry policy.
+        /// If the GET request is successful, adds it to the cache with the specified TTL.
+        /// </summary>
+        /// <param name="requestUrl"></param>
+        /// <param name="typeInfo">JSON serialization-related metadata.</param>
+        /// <param name="cacheTTL">Time after which a newly-added entry will expire from the cache.</param>
+        /// <param name="cancellationToken">A <see cref="CancellationToken"/> for controlling the lifetime of the asynchronous operation.</param>
+        /// <param name="setAuthorizationHeader">If true, the authentication header is set with an authentication token.</param>
+        /// <typeparam name="TResponse">Response's content deserialization type.</typeparam>
+        /// <returns></returns>
+        protected async Task<TResponse> HttpGetRequestWithExpirableCachingAndRetryPolicyAsync<TResponse>(Uri requestUrl,
+            JsonTypeInfo<TResponse> typeInfo, TimeSpan cacheTTL, CancellationToken cancellationToken, bool setAuthorizationHeader = false)
+        {
+            var cacheKey = await ToCacheKeyAsync(requestUrl, setAuthorizationHeader, cancellationToken);
+
+            return (await cachingRetryHandler.ExecuteWithRetryConversionAndCachingAsync(cacheKey, async ct =>
+            {
+                //request must be recreated in every retry.
+                var httpRequest = await CreateGetHttpRequest(requestUrl, setAuthorizationHeader, ct);
+
+                return await HttpClient.SendAsync(httpRequest, ct);
+            },
+            (m, ct) => GetApiResponseContentAsync(m, typeInfo, ct), DateTimeOffset.Now + cacheTTL, cancellationToken))!;
+        }
+
         /// <summary>
         /// Get request with retry policy
         /// </summary>
@@ -223,7 +249,7 @@ private async Task<HttpRequestMessage> CreateGetHttpRequest(Uri requestUrl, bool
         }
 
         /// <summary>
-        /// Sends an Http request to the URL and deserializes the body response to the specified type 
+        /// Sends an Http request to the URL and deserializes the body response to the specified type
         /// </summary>
         /// <param name="httpRequestFactory">Factory that creates new http requests, in the event of retry the factory is called again
         /// and must be idempotent.</param>

src/Tes.ApiClients/Models/Terra/SamActionManagedIdentityApiResponse.cs

@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Text.Json.Serialization;
+
+namespace Tes.ApiClients.Models.Terra
+{
+    public class SamActionManagedIdentityApiResponse
+    {
+        [JsonPropertyName("id")]
+        public ActionManagedIdentityId actionManagedIdentityId { get; set; }
+
+        [JsonPropertyName("displayName")]
+        public string DisplayName { get; set; }
+
+        [JsonPropertyName("managedResourceGroupCoordinates")]
+        public ManagedResourceGroupCoordinates managedResourceGroupCoordinates { get; set; }
+
+        [JsonPropertyName("objectId")]
+        public string ObjectId { get; set; }
+
+    }
+
+    public class ActionManagedIdentityId
+    {
+        [JsonPropertyName("resourceId")]
+        public FullyQualifiedResourceId ResourceId { get; set; }
+
+        [JsonPropertyName("action")]
+        public string Action { get; set; }
+
+        [JsonPropertyName("billingProfileId")]
+        public Guid BillingProfileId { get; set; }
+    }
+
+    public class FullyQualifiedResourceId
+    {
+        [JsonPropertyName("resourceTypeName")]
+        public string ResourceTypeName { get; set; }
+
+        [JsonPropertyName("resourceId")]
+        public string ResourceId { get; set; }
+    }
+
+    public class ManagedResourceGroupCoordinates
+    {
+        [JsonPropertyName("tenantId")]
+        public Guid TenantId { get; set; }
+
+        [JsonPropertyName("subscriptionId")]
+        public Guid SubscriptionId { get; set; }
+
+        [JsonPropertyName("managedResourceGroupName")]
+        public string ManagedResourceGroupName { get; set; }
+    }
+
+    [JsonSerializable(typeof(SamActionManagedIdentityApiResponse))]
+    public partial class SamActionManagedIdentityApiResponseContext : JsonSerializerContext
+    { }
+
+}

src/Tes.ApiClients/TerraSamApiClient.cs

@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Azure.Core;
+using CommonUtilities;
+using Microsoft.Extensions.Caching.Memory;
+using Microsoft.Extensions.Logging;
+using Tes.ApiClients.Models.Terra;
+
+namespace Tes.ApiClients
+{
+    /// <summary>
+    /// Terra Sam api client
+    /// Sam manages authorization and IAM functionality
+    /// </summary>
+    public class TerraSamApiClient : TerraApiClient
+    {
+        private const string SamApiSegments = @"/api/azure/v1";
+
+        private static readonly IMemoryCache SharedMemoryCache = new MemoryCache(new MemoryCacheOptions());
+
+        /// <summary>
+        /// Constructor of TerraSamApiClient
+        /// </summary>
+        /// <param name="apiUrl">Sam API host</param>
+        /// <param name="tokenCredential"></param>
+        /// <param name="cachingRetryHandler"></param>
+        /// <param name="azureCloudIdentityConfig"></param>
+        /// <param name="logger"></param>
+        public TerraSamApiClient(string apiUrl, TokenCredential tokenCredential, CachingRetryPolicyBuilder cachingRetryHandler,
+            AzureEnvironmentConfig azureCloudIdentityConfig, ILogger<TerraSamApiClient> logger) : base(apiUrl, tokenCredential, cachingRetryHandler, azureCloudIdentityConfig, logger)
+        { }
+
+        public static TerraSamApiClient CreateTerraSamApiClient(string apiUrl, TokenCredential tokenCredential, AzureEnvironmentConfig azureCloudIdentityConfig)
+        {
+            return CreateTerraApiClient<TerraSamApiClient>(apiUrl, SharedMemoryCache, tokenCredential, azureCloudIdentityConfig);
+        }
+
+        /// <summary>
+        /// Protected parameter-less constructor
+        /// </summary>
+        protected TerraSamApiClient() { }
+
+        public virtual async Task<SamActionManagedIdentityApiResponse> GetActionManagedIdentityForACRPullAsync(Guid resourceId, TimeSpan cacheTTL, CancellationToken cancellationToken)
+        {
+            return await GetActionManagedIdentityAsync("private_azure_container_registry", resourceId, "pull_image", cacheTTL, cancellationToken);
+        }
+
+        private async Task<SamActionManagedIdentityApiResponse> GetActionManagedIdentityAsync(string resourceType, Guid resourceId, string action, TimeSpan cacheTTL, CancellationToken cancellationToken)
+        {
+            ArgumentNullException.ThrowIfNull(resourceId);
+
+            var url = GetSamActionManagedIdentityUrl(resourceType, resourceId, action);
+
+            Logger.LogInformation(@"Fetching action managed identity from Sam for {resourceId}", resourceId);
+
+            try
+            {
+                return await HttpGetRequestWithExpirableCachingAndRetryPolicyAsync(url,
+                    SamActionManagedIdentityApiResponseContext.Default.SamActionManagedIdentityApiResponse, cacheTTL, cancellationToken, setAuthorizationHeader: true);
+            }
+            catch (HttpRequestException e)
+            {
+                // Sam will return a 404 if there is no action identity that matches the query,
+                // or if we don't have access to it.
+                if (e.StatusCode == System.Net.HttpStatusCode.NotFound)
+                {
+                    return null;
+                }
+                else
+                {
+                    throw;
+                }
+
+            }
+        }
+
+        public virtual Uri GetSamActionManagedIdentityUrl(string resourceType, Guid resourceId, string action)
+        {
+            var apiRequestUrl = $"{ApiUrl.TrimEnd('/')}{SamApiSegments}/actionManagedIdentity/{resourceType}/{resourceId}/{action}";
+
+            var uriBuilder = new UriBuilder(apiRequestUrl);
+
+            return uriBuilder.Uri;
+        }
+    }
+}

src/Tes.Runner.Test/Docker/ContainerRegistryAuthorizationManagerTests.cs

@@ -32,7 +32,7 @@ public void SetUp()
 
         mockCredentialsManager = new Mock<CredentialsManager>();
         mockCredentials = new Mock<TokenCredential>();
-        mockCredentialsManager.Setup(c => c.GetTokenCredential(It.IsAny<RuntimeOptions>(), It.IsAny<string>()))
+        mockCredentialsManager.Setup(c => c.GetAcrPullTokenCredential(It.IsAny<RuntimeOptions>(), It.IsAny<string>()))
             .Returns(mockCredentials.Object);
 
 

src/Tes.Runner.Test/Docker/DockerExecutorTests.cs

@@ -35,7 +35,7 @@ public void SetUp()
             dockerClientMock.Setup(d => d.Volumes).Returns(dockerVolumeMock.Object);
             dockerClient = dockerClientMock.Object;
             var credentialsManager = new Mock<CredentialsManager>();
-            credentialsManager.Setup(m => m.GetTokenCredential(It.IsAny<RuntimeOptions>(), It.IsAny<string>()))
+            credentialsManager.Setup(m => m.GetAcrPullTokenCredential(It.IsAny<RuntimeOptions>(), It.IsAny<string>()))
                 .Throws(new IdentityUnavailableException());
             containerRegistryAuthorizationManager = new(credentialsManager.Object);
         }