Successfully don't add SAS token on blob storage URLs with no permissions (#816)

Author: BMurri

src/Tes.Runner.Test/Storage/ArmUrlTransformationStrategyTests.cs

@@ -5,7 +5,6 @@
 using Azure.Storage.Blobs.Models;
 using Azure.Storage.Sas;
 using CommonUtilities;
-using CommonUtilities.Options;
 using Moq;
 using Tes.Runner.Models;
 using Tes.Runner.Storage;
@@ -16,6 +15,7 @@ namespace Tes.Runner.Test.Storage
     public class ArmUrlTransformationStrategyTests
     {
         private Mock<BlobServiceClient> mockBlobServiceClient = null!;
+        private Mock<Runner.Transfer.BlobApiHttpUtils> mockBlobApiHttpUtils = null!;
         private ArmUrlTransformationStrategy armUrlTransformationStrategy = null!;
         private UserDelegationKey userDelegationKey = null!;
         const string StorageAccountName = "foo";
@@ -30,12 +30,42 @@ public void SetUp()
                 AzureEnvironmentConfig = AzureEnvironmentConfig.FromArmEnvironmentEndpoints(CommonUtilities.AzureCloud.AzureCloudConfig.FromKnownCloudNameAsync().Result)
             };
 
-            armUrlTransformationStrategy = new ArmUrlTransformationStrategy(_ => mockBlobServiceClient.Object, options);
+            mockBlobApiHttpUtils = new();
+            mockBlobApiHttpUtils.Setup(x => x.IsEndPointPublic(It.IsAny<Uri>()))
+                .ReturnsAsync(false);
+
+            armUrlTransformationStrategy = new ArmUrlTransformationStrategy(_ => mockBlobServiceClient.Object, options, mockBlobApiHttpUtils.Object);
             userDelegationKey = BlobsModelFactory.UserDelegationKey(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), DateTimeOffset.UtcNow,
                 DateTimeOffset.UtcNow.AddHours(1), "SIGNED_SERVICE", "V1_0", RunnerTestUtils.GenerateRandomTestAzureStorageKey());
             mockBlobServiceClient.Setup(c => c.GetUserDelegationKeyAsync(It.IsAny<DateTimeOffset?>(), It.IsAny<DateTimeOffset>(), It.IsAny<CancellationToken>()))
                 .ReturnsAsync(Azure.Response.FromValue(userDelegationKey, null!));
+        }
+
+        [TestMethod]
+        public async Task TransformUrlWithStrategyAsync_BlobStorageUrlWithOutPermissions_WhenRequestingRead_UrlIsReturnAsIs()
+        {
+            mockBlobApiHttpUtils.Setup(x => x.IsEndPointPublic(It.IsAny<Uri>()))
+                .ReturnsAsync(true);
+
+            var sourceUrl = $"https://{StorageAccountName}.blob.core.windows.net/cont/blob";
+            var transformedUrl = await armUrlTransformationStrategy.TransformUrlWithStrategyAsync(sourceUrl, BlobSasPermissions.Read);
+
+            Assert.AreEqual(1, mockBlobApiHttpUtils.Invocations.Count);
+            Assert.IsNotNull(transformedUrl);
+            var blobUri = new Uri(sourceUrl);
+            Assert.AreEqual(blobUri.AbsoluteUri, transformedUrl.AbsoluteUri);
+        }
+
+        [TestMethod]
+        public async Task TransformUrlWithStrategyAsync_BlobStorageUrlWithOutPermissions_WhenRequestingWrite_IsEndPointPublicIsNotCalled()
+        {
+            mockBlobApiHttpUtils.Setup(x => x.IsEndPointPublic(It.IsAny<Uri>()))
+                .ReturnsAsync(true);
+
+            var sourceUrl = $"https://{StorageAccountName}.blob.core.windows.net/cont/blob";
+            var transformedUrl = await armUrlTransformationStrategy.TransformUrlWithStrategyAsync(sourceUrl, BlobSasPermissions.Create);
 
+            Assert.IsFalse(mockBlobApiHttpUtils.Invocations.Any());
         }
 
         [TestMethod]

src/Tes.Runner/Storage/ArmUrlTransformationStrategy.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Net;
 using Azure.Storage.Blobs;
 using Azure.Storage.Blobs.Models;
 using Azure.Storage.Sas;
@@ -15,22 +16,25 @@ public class ArmUrlTransformationStrategy : IUrlTransformationStrategy
         const string BlobUrlPrefix = ".blob."; // "core.windows.net";
         private const int BlobSasTokenExpirationInHours = 24 * 7; //7 days which is the Azure Batch node runtime;
         const int UserDelegationKeyExpirationInHours = 1;
+        private static readonly Lazy<BlobApiHttpUtils> blobApiHttpUtilsFactory = new(() => new());
 
         private readonly ILogger logger = PipelineLoggerFactory.Create<ArmUrlTransformationStrategy>();
         private readonly Dictionary<string, UserDelegationKey> userDelegationKeyDictionary = [];
         private readonly SemaphoreSlim semaphoreSlim = new(1, 1);
         private readonly Func<Uri, BlobServiceClient> blobServiceClientFactory;
         private readonly RuntimeOptions runtimeOptions;
         private readonly string storageHostSuffix;
+        private readonly BlobApiHttpUtils blobApiHttpUtils;
 
-        public ArmUrlTransformationStrategy(Func<Uri, BlobServiceClient> blobServiceClientFactory, RuntimeOptions runtimeOptions)
+        public ArmUrlTransformationStrategy(Func<Uri, BlobServiceClient> blobServiceClientFactory, RuntimeOptions runtimeOptions, BlobApiHttpUtils? blobApiHttpUtils = default)
         {
             ArgumentNullException.ThrowIfNull(blobServiceClientFactory);
             ArgumentNullException.ThrowIfNull(runtimeOptions);
 
             this.blobServiceClientFactory = blobServiceClientFactory;
             this.runtimeOptions = runtimeOptions;
             storageHostSuffix = BlobUrlPrefix + this.runtimeOptions!.AzureEnvironmentConfig!.StorageUrlSuffix;
+            this.blobApiHttpUtils = blobApiHttpUtils ?? blobApiHttpUtilsFactory.Value;
         }
 
         public async Task<Uri> TransformUrlWithStrategyAsync(string sourceUrl, BlobSasPermissions blobSasPermissions)
@@ -60,9 +64,16 @@ private async Task<Uri> GetStorageUriWithSasTokenAsync(string sourceUrl, BlobSas
         {
             try
             {
-                var blobUrl = new BlobUriBuilder(new Uri(sourceUrl));
+                Uri uri = new(sourceUrl);
+                var blobUrl = new BlobUriBuilder(uri);
                 var blobServiceClient = blobServiceClientFactory(new Uri($"https://{blobUrl.Host}"));
 
+                if ((permissions & ~(BlobSasPermissions.Read | BlobSasPermissions.List | BlobSasPermissions.Execute)) == 0 && await blobApiHttpUtils.IsEndPointPublic(uri))
+                {
+                    logger.LogWarning("The URL provided is not a storage account the managed identity can access. The resolution strategy won't be applied. Host: {Host}", blobUrl.Host);
+                    return uri;
+                }
+
                 var userKey = await GetUserDelegationKeyAsync(blobServiceClient, blobUrl.AccountName);
 
                 var sasBuilder = new BlobSasBuilder()
@@ -73,13 +84,9 @@ private async Task<Uri> GetStorageUriWithSasTokenAsync(string sourceUrl, BlobSas
                 };
 
                 sasBuilder.SetPermissions(permissions);
+                blobUrl.Sas = sasBuilder.ToSasQueryParameters(userKey, blobUrl.AccountName);
 
-                var blobUriWithSas = new BlobUriBuilder(blobUrl.ToUri())
-                {
-                    Sas = sasBuilder.ToSasQueryParameters(userKey, blobUrl.AccountName)
-                };
-
-                return blobUriWithSas.ToUri();
+                return blobUrl.ToUri();
             }
             catch (Exception e)
             {

src/Tes.Runner/Transfer/BlobApiHttpUtils.cs

@@ -93,6 +93,7 @@ public static Uri ParsePutBlockUrl(Uri? baseUri, int ordinal)
     {
         return new Uri($"{baseUri?.AbsoluteUri}&comp=block&blockid={ToBlockId(ordinal)}");
     }
+
     public static Uri ParsePutAppendBlockUrl(Uri? baseUri)
     {
         ArgumentNullException.ThrowIfNull(baseUri);
@@ -180,6 +181,29 @@ public static bool UrlContainsSasToken(string sourceUrl)
 
         return !string.IsNullOrWhiteSpace(blobBuilder.Sas?.Signature);
     }
+
+    public virtual async Task<bool> IsEndPointPublic(Uri uri)
+    {
+        try
+        {
+            using var _ = await ExecuteHttpRequestAsync(() => new HttpRequestMessage(HttpMethod.Head, uri));
+            return true;
+
+        }
+        catch (HttpRequestException e) when (e.StatusCode switch
+        {
+            HttpStatusCode.Unauthorized => true,
+            // Because some servers confuse the difference between 401 & 403
+            HttpStatusCode.Forbidden => true,
+            // Azure Blob Storage returns 409 instead of 401
+            HttpStatusCode.Conflict => true,
+            _ => false,
+        })
+        {
+            return false;
+        }
+    }
+
     private async Task<HttpResponseMessage> ExecuteHttpRequestImplAsync(Func<HttpRequestMessage> request, CancellationToken cancellationToken)
     {
         HttpResponseMessage? response = null;