[release/8.0] Hook which rewrites the container registry in annotatio…

…ns. (#3583)

* Hook which rewrites the container registry in annotations.

* Change to using options instead of extension method.

---------

Co-authored-by: Mitch Denny <midenn@microsoft.com>

src/Aspire.Hosting/DistributedApplicationBuilder.cs

@@ -157,6 +157,12 @@ public DistributedApplicationBuilder(DistributedApplicationOptions options)
         _innerBuilder.Services.AddLifecycleHook<Http2TransportMutationHook>();
         _innerBuilder.Services.AddKeyedSingleton<IDistributedApplicationPublisher, ManifestPublisher>("manifest");
 
+        // Overwrite registry if override specified in options
+        if (!string.IsNullOrEmpty(options.ContainerRegistryOverride))
+        {
+            _innerBuilder.Services.AddLifecycleHook<ContainerRegistryHook>();
+        }
+
         _innerBuilder.Services.AddSingleton(ExecutionContext);
         LogBuilderConstructed(this);
     }

src/Aspire.Hosting/DistributedApplicationOptions.cs

@@ -24,6 +24,12 @@ public DistributedApplicationOptions()
         _projectDirectoryLazy = new(ResolveProjectDirectory);
     }
 
+    /// <summary>
+    /// When containers are used, use this value instead to override the container registry
+    /// that is specified.
+    /// </summary>
+    public string? ContainerRegistryOverride { get; set; }
+
     /// <summary>
     /// The command line arguments.
     /// </summary>

src/Aspire.Hosting/Lifecycle/ContainerRegistryCheckHook.cs

@@ -0,0 +1,25 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Aspire.Hosting.ApplicationModel;
+using Aspire.Hosting.Lifecycle;
+
+namespace Aspire.Hosting;
+
+internal class ContainerRegistryHook(DistributedApplicationOptions options) : IDistributedApplicationLifecycleHook
+{
+    public Task BeforeStartAsync(DistributedApplicationModel appModel, CancellationToken cancellationToken = default)
+    {
+
+        var resourcesWithContainerImages = appModel.Resources
+                                                   .SelectMany(r => r.Annotations.OfType<ContainerImageAnnotation>()
+                                                                                 .Select(cia => new { Resource = r, Annotation = cia }));
+
+        foreach (var resourceWithContainerImage in resourcesWithContainerImages)
+        {
+            resourceWithContainerImage.Annotation.Registry = options.ContainerRegistryOverride;
+        }
+
+        return Task.CompletedTask;
+    }
+}

tests/Aspire.Hosting.Tests/ManifestGenerationTests.cs

@@ -86,6 +86,37 @@ public void EnsureWorkerProjectDoesNotGetBindingsGenerated()
         Assert.False(workerA.TryGetProperty("bindings", out _));
     }
 
+    [Fact]
+    public async Task WithContainerRegistryUpdatesContainerImageAnnotationsDuringPublish()
+    {
+        var builder = DistributedApplication.CreateBuilder(new DistributedApplicationOptions
+        {
+            Args = GetManifestArgs(),
+            ContainerRegistryOverride = "myprivateregistry.company.com"
+        });
+
+        var redis = builder.AddRedis("redis");
+        builder.Build().Run();
+
+        var redisManifest = await ManifestUtils.GetManifest(redis.Resource);
+        var expectedManifest = $$"""
+            {
+              "type": "container.v0",
+              "connectionString": "{redis.bindings.tcp.host}:{redis.bindings.tcp.port}",
+              "image": "myprivateregistry.company.com/{{RedisContainerImageTags.Image}}:{{RedisContainerImageTags.Tag}}",
+              "bindings": {
+                "tcp": {
+                  "scheme": "tcp",
+                  "protocol": "tcp",
+                  "transport": "tcp",
+                  "targetPort": 6379
+                }
+              }
+            }
+            """;
+        Assert.Equal(expectedManifest, redisManifest.ToString());
+    }
+
     [Fact]
     public void EnsureExecutablesWithDockerfileProduceDockerfilev0Manifest()
     {