Restrict Azure Storage large-message resolution to the current task hub

src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs

@@ -2176,6 +2176,7 @@ public Task PurgeOrchestrationHistoryAsync(DateTime thresholdDateTimeUtc, Orches
         /// such as the input or output status fields, and for which the blob URI was stored instead.
         /// </summary>
         /// <param name="blobUri">The URI of the blob.</param>
+        /// <exception cref="ArgumentException">The URI is not in this task hub's large-message container.</exception>
         public Task<string> DownloadBlobAsync(string blobUri)
         {
             return this.messageManager.DownloadAndDecompressAsBytesAsync(new Uri(blobUri));

src/DurableTask.AzureStorage/EntityTrackingStoreQueries.cs

@@ -242,7 +242,7 @@ bool OrchestrationIsRunning(OrchestrationStatus? status)
             {
                 // first, retrieve the entity scheduler state (= input of the orchestration state), possibly from blob storage.
                 string serializedSchedulerState;
-                if (MessageManager.TryGetLargeMessageReference(state.Input, out Uri blobUrl))
+                if (this.messageManager.TryGetLargeMessageReference(state.Input, out Uri blobUrl))
                 {
                     serializedSchedulerState = await this.messageManager.DownloadAndDecompressAsBytesAsync(blobUrl);
                 }

src/DurableTask.AzureStorage/MessageManager.cs

@@ -25,6 +25,7 @@ namespace DurableTask.AzureStorage
     using System.Threading;
     using System.Threading.Tasks;
     using Azure;
+    using Azure.Storage.Blobs;
     using Azure.Storage.Blobs.Models;
     using Azure.Storage.Queues.Models;
     using DurableTask.AzureStorage.Storage;
@@ -123,32 +124,71 @@ public async Task<string> SerializeMessageDataAsync(MessageData messageData, Can
         }
 
         /// <summary>
-        /// If the "message" of an orchestration state is actually a URI retrieves actual message from blob storage.
+        /// If the message is a reference in this task hub's large-message container, retrieve it from blob storage.
         /// Otherwise returns the message as is.
         /// </summary>
-        /// <param name="message">The message to be fetched if it is a url.</param>
+        /// <param name="message">The message to be fetched if it is a large-message reference.</param>
         /// <param name="cancellationToken">A token used for canceling the operation.</param>
         /// <returns>Actual string representation of message.</returns>
         public async Task<string> FetchLargeMessageIfNecessary(string message, CancellationToken cancellationToken = default)
         {
-            if (TryGetLargeMessageReference(message, out Uri blobUrl))
+            if (this.TryGetLargeMessageReference(message, out _, out Blob blob))
             {
-                return await this.DownloadAndDecompressAsBytesAsync(blobUrl, cancellationToken);
+                return await DownloadAndDecompressAsBytesAsync(blob, cancellationToken);
             }
             else
             {
                 return message;
             }
         }
 
-        internal static bool TryGetLargeMessageReference(string messagePayload, out Uri blobUrl)
+        internal bool TryGetLargeMessageReference(string messagePayload, out Uri blobUrl)
         {
-            if (Uri.IsWellFormedUriString(messagePayload, UriKind.Absolute))
+            return this.TryGetLargeMessageReference(messagePayload, out blobUrl, out _);
+        }
+
+        private bool TryGetLargeMessageReference(string messagePayload, out Uri blobUrl, out Blob blob)
+        {
+            if (Uri.IsWellFormedUriString(messagePayload, UriKind.Absolute) &&
+                Uri.TryCreate(messagePayload, UriKind.Absolute, out Uri candidate))
             {
-                return Uri.TryCreate(messagePayload, UriKind.Absolute, out blobUrl);
+                Uri containerUri = this.blobContainer.GetBlobContainerUri();
+                if (IsInLargeMessageContainer(candidate, containerUri))
+                {
+                    Blob candidateBlob = this.blobContainer.GetBlobReference(new BlobUriBuilder(candidate).BlobName);
+
+                    // SDK decoding and URI reconstruction must not change the allowed storage scope.
+                    if (IsInLargeMessageContainer(candidateBlob.Uri, containerUri))
+                    {
+                        blobUrl = candidate;
+                        blob = candidateBlob;
+                        return true;
+                    }
+                }
             }
 
             blobUrl = null;
+            blob = null;
+            return false;
+        }
+
+        private static bool IsInLargeMessageContainer(Uri candidate, Uri containerUri)
+        {
+            string containerPath = containerUri.AbsolutePath.TrimEnd('/') + "/";
+            if (candidate.Scheme == containerUri.Scheme &&
+                string.Equals(candidate.Host, containerUri.Host, StringComparison.OrdinalIgnoreCase) &&
+                candidate.Port == containerUri.Port &&
+                candidate.UserInfo.Length == 0 &&
+                candidate.Query.Length == 0 &&
+                candidate.Fragment.Length == 0 &&
+                candidate.AbsolutePath.StartsWith(containerPath, StringComparison.Ordinal) &&
+                candidate.AbsolutePath.Length > containerPath.Length)
+            {
+                var blobUri = new BlobUriBuilder(candidate);
+                var container = new BlobUriBuilder(containerUri);
+                return blobUri.BlobContainerName == container.BlobContainerName && !string.IsNullOrEmpty(blobUri.BlobName);
+            }
+
             return false;
         }
 
@@ -249,7 +289,11 @@ public async Task<string> DownloadAndDecompressAsBytesAsync(string blobName, Can
 
         public Task<string> DownloadAndDecompressAsBytesAsync(Uri blobUri, CancellationToken cancellationToken = default)
         {
-            Blob blob = this.azureStorageClient.GetBlobReference(blobUri);
+            if (blobUri == null || !this.TryGetLargeMessageReference(blobUri.OriginalString, out _, out Blob blob))
+            {
+                throw new ArgumentException("Blob must belong to this task hub's large-message container.", nameof(blobUri));
+            }
+
             return DownloadAndDecompressAsBytesAsync(blob, cancellationToken);
         }
 

src/DurableTask.AzureStorage/Tracking/AzureTableTrackingStore.cs

@@ -489,7 +489,7 @@ async Task<OrchestrationState> ConvertFromAsync(OrchestrationInstanceStatus orch
 
             if (this.settings.FetchLargeMessageDataEnabled)
             {
-                if (MessageManager.TryGetLargeMessageReference(orchestrationState.Input, out Uri blobUrl))
+                if (this.messageManager.TryGetLargeMessageReference(orchestrationState.Input, out Uri blobUrl))
                 {
                     string json = await this.messageManager.DownloadAndDecompressAsBytesAsync(blobUrl, cancellationToken);
 

test/DurableTask.AzureStorage.Tests/AzureStorageScenarioTests.cs

@@ -31,6 +31,7 @@ namespace DurableTask.AzureStorage.Tests
     using System.Collections.Generic;
     using System.Diagnostics;
     using System.IO;
+    using System.IO.Compression;
     using System.Linq;
     using System.Net;
     using System.Reflection;
@@ -1295,6 +1296,104 @@ public async Task TerminateOrchestration(bool enableExtendedSessions)
             }
         }
 
+        [DataTestMethod]
+        [DataRow("terminate")]
+        [DataRow("suspend")]
+        [DataRow("resume")]
+        public async Task ManagementReason_PrivateBlobUrl_RemainsLiteral(string operation)
+        {
+            string connectionString = TestHelpers.GetTestStorageAccountConnectionString();
+            var container = new BlobServiceClient(connectionString)
+                .GetBlobContainerClient("private-reason-" + Guid.NewGuid().ToString("N"));
+            await container.CreateAsync(PublicAccessType.None);
+
+            try
+            {
+                var blob = container.GetBlobClient("private-backup.json.gz");
+                string canary = "synthetic-private-value-" + Guid.NewGuid().ToString("N");
+                using (var buffer = new MemoryStream())
+                {
+                    using (var gzip = new GZipStream(buffer, CompressionLevel.Optimal, leaveOpen: true))
+                    {
+                        byte[] payload = Encoding.UTF8.GetBytes(JsonConvert.SerializeObject(canary));
+                        await gzip.WriteAsync(payload, 0, payload.Length);
+                    }
+
+                    buffer.Position = 0;
+                    await blob.UploadAsync(buffer);
+                }
+
+                var anonymousBlob = new BlobClient(blob.Uri);
+                var denied = await Assert.ThrowsExceptionAsync<Azure.RequestFailedException>(async () =>
+                {
+                    using BlobDownloadStreamingResult response = (await anonymousBlob.DownloadStreamingAsync()).Value;
+                });
+                Assert.IsTrue(
+                    denied.Status == 401 || denied.Status == 403 || denied.Status == 404,
+                    $"Unexpected anonymous read failure status: {denied.Status}.");
+
+                using TestOrchestrationHost host = TestHelpers.GetTestOrchestrationHost(
+                    enableExtendedSessions: false,
+                    modifySettingsAction: settings =>
+                    {
+                        settings.StorageAccountClientProvider = new StorageAccountClientProvider(connectionString);
+                        settings.TaskHubName = "ReasonTest" + Guid.NewGuid().ToString("N");
+                    });
+                await host.StartAsync();
+                try
+                {
+                    var client = await host.StartOrchestrationAsync(typeof(Orchestrations.Counter), 0);
+                    await client.WaitForStartupAsync(StandardTimeout);
+
+                    await client.SuspendAsync("literal-control");
+                    var control = await client.WaitForStatusChange(StandardTimeout, OrchestrationStatus.Suspended);
+                    Assert.AreEqual("literal-control", control.Output);
+
+                    if (operation != "resume")
+                    {
+                        await client.ResumeAsync("literal-control");
+                        await client.WaitForStatusChange(StandardTimeout, OrchestrationStatus.Running);
+                    }
+
+                    OrchestrationStatus expectedStatus;
+                    if (operation == "terminate")
+                    {
+                        await client.TerminateAsync(blob.Uri.AbsoluteUri);
+                        expectedStatus = OrchestrationStatus.Terminated;
+                    }
+                    else if (operation == "suspend")
+                    {
+                        await client.SuspendAsync(blob.Uri.AbsoluteUri);
+                        expectedStatus = OrchestrationStatus.Suspended;
+                    }
+                    else
+                    {
+                        await client.ResumeAsync(blob.Uri.AbsoluteUri);
+                        expectedStatus = OrchestrationStatus.Running;
+                    }
+
+                    var status = await client.WaitForStatusChange(StandardTimeout, expectedStatus);
+                    Assert.AreEqual(blob.Uri.AbsoluteUri, status.Output,
+                        "Management reasons must not be replaced with private blob contents.");
+
+                    if (expectedStatus != OrchestrationStatus.Terminated)
+                    {
+                        await client.TerminateAsync("cleanup");
+                        await client.WaitForCompletionAsync(StandardTimeout);
+                    }
+                }
+                finally
+                {
+                    await host.StopAsync();
+                    await host.service.DeleteAsync();
+                }
+            }
+            finally
+            {
+                await container.DeleteAsync();
+            }
+        }
+
         /// <summary>
         /// End-to-end test which validates the Suspend-Resume functionality.
         /// </summary>