Fix issues pushing images to Amazon ECR

Author: normj

Microsoft.NET.Build.Containers/AmazonECRMessageHandler.cs

@@ -0,0 +1,37 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Microsoft.NET.Build.Containers;
+
+/// <summary>
+/// A delegating handler that handles the special error handling needed for Amazon ECR.
+/// 
+/// When pushing images to ECR if the target container repository does not exist ECR ends
+/// the connection causing an IOException with a generic "The response ended prematurely."
+/// error message. The handler catches the generic error and provides a more informed error
+/// message to let the user know they need to create the repository.
+/// </summary>
+public class AmazonECRMessageHandler : DelegatingHandler
+{
+    public AmazonECRMessageHandler(HttpMessageHandler innerHandler) : base(innerHandler) { }
+
+    protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+    {
+        try
+        {
+            return await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
+        }
+        catch (HttpRequestException e) when (e.InnerException is IOException ioe && ioe.Message.Equals("The response ended prematurely.", StringComparison.OrdinalIgnoreCase))
+        {
+            var message = "Request to Amazon Elastic Container Registry failed prematurely. This is often caused when the target repository does not exist in the registry.";
+            throw new ContainerHttpException(message, request.RequestUri?.ToString(), null);
+        }
+        catch
+        {
+            throw;
+        }
+    }
+}

Microsoft.NET.Build.Containers/Registry.cs

@@ -6,20 +6,58 @@
 using System.Net.Http;
 using System.Net.Http.Headers;
 using System.Net.Http.Json;
+using System.Reflection.Metadata.Ecma335;
 using System.Runtime.ExceptionServices;
 using System.Text;
 using System.Text.Json.Nodes;
 using System.Xml.Linq;
 
 namespace Microsoft.NET.Build.Containers;
 
-public record struct Registry(Uri BaseUri)
+public struct Registry
 {
     private const string DockerManifestV2 = "application/vnd.docker.distribution.manifest.v2+json";
     private const string DockerContainerV1 = "application/vnd.docker.container.image.v1+json";
-    private const int MaxChunkSizeBytes = 1024 * 64;
 
-    private string RegistryName { get; } = BaseUri.Host;
+    private Uri BaseUri { get; init; }
+    private string RegistryName => BaseUri.Host;
+
+    public Registry(Uri baseUri)
+    {
+        BaseUri = baseUri;
+        _client = CreateClient();
+    }
+
+    /// <summary>
+    /// The max chunk size for patch blob uploads. By default the size is 64 KB.
+    /// Amazon Elasic Container Registry (ECR) requires patch chunk size to be 5 MB except for the last chunk.
+    /// </summary>
+    public readonly int MaxChunkSizeBytes => IsAmazonECRRegistry ? 5248080 : 1024 * 64;
+
+    /// <summary>
+    /// Check to see if the registry is for Amazon Elastic Container Registry (ECR).
+    /// </summary>
+    public readonly bool IsAmazonECRRegistry
+    {
+        get
+        {
+            // If this the registry is to public ECR the name will contain "public.ecr.aws".
+            if (RegistryName.Contains("public.ecr.aws"))
+            {
+                return true;
+            }
+
+            // If the registry is to a private ECR the registry will start with an account id which is a 12 digit number and will container either
+            // ".ecr." or ".ecr-" if pushed to a FIPS endpoint.
+            var accountId = RegistryName.Split('.')[0];
+            if ((RegistryName.Contains(".ecr.") || RegistryName.Contains(".ecr-")) && accountId.Length == 12 && long.TryParse(accountId, out _))
+            {
+                return true;
+            }
+
+            return false;
+        }
+    }
 
     public async Task<Image> GetImageManifest(string name, string reference)
     {
@@ -118,7 +156,7 @@ private readonly async Task UploadBlob(string name, string digest, Stream conten
 
         if (pushResponse.StatusCode != HttpStatusCode.Accepted)
         {
-            string errorMessage = $"Failed to upload blob to {pushUri}; recieved {pushResponse.StatusCode} with detail {await pushResponse.Content.ReadAsStringAsync()}";
+            string errorMessage = $"Failed to upload blob to {pushUri}; received {pushResponse.StatusCode} with detail {await pushResponse.Content.ReadAsStringAsync()}";
             throw new ApplicationException(errorMessage);
         }
 
@@ -162,9 +200,10 @@ private readonly async Task UploadBlob(string name, string digest, Stream conten
 
             HttpResponseMessage patchResponse = await client.PatchAsync(patchUri, content);
 
-            if (patchResponse.StatusCode != HttpStatusCode.Accepted)
+            // Fail the upload if the response code is not Accepted (202) or if uploading to Amazon ECR which returns back Created (201).
+            if (!(patchResponse.StatusCode == HttpStatusCode.Accepted || (IsAmazonECRRegistry && patchResponse.StatusCode == HttpStatusCode.Created)))
             {
-                string errorMessage = $"Failed to upload blob to {patchUri}; recieved {patchResponse.StatusCode} with detail {await patchResponse.Content.ReadAsStringAsync()}";
+                string errorMessage = $"Failed to upload blob to {patchUri}; received {patchResponse.StatusCode} with detail {await patchResponse.Content.ReadAsStringAsync()}";
                 throw new ApplicationException(errorMessage);
             }
 
@@ -194,7 +233,7 @@ private readonly async Task UploadBlob(string name, string digest, Stream conten
 
         if (finalizeResponse.StatusCode != HttpStatusCode.Created)
         {
-            string errorMessage = $"Failed to finalize upload to {putUri}; recieved {finalizeResponse.StatusCode} with detail {await finalizeResponse.Content.ReadAsStringAsync()}";
+            string errorMessage = $"Failed to finalize upload to {putUri}; received {finalizeResponse.StatusCode} with detail {await finalizeResponse.Content.ReadAsStringAsync()}";
             throw new ApplicationException(errorMessage);
         }
     }
@@ -211,16 +250,22 @@ private readonly async Task<bool> BlobAlreadyUploaded(string name, string digest
         return false;
     }
 
-    private static HttpClient _client = CreateClient();
+    private HttpClient _client;
 
-    private static HttpClient GetClient()
+    private HttpClient GetClient()
     {
         return _client;
     }
 
-    private static HttpClient CreateClient()
+    private HttpClient CreateClient()
     {
-        var clientHandler = new AuthHandshakeMessageHandler(new SocketsHttpHandler() { PooledConnectionLifetime = TimeSpan.FromMilliseconds(10 /* total guess */) });
+        HttpMessageHandler clientHandler = new AuthHandshakeMessageHandler(new SocketsHttpHandler() { PooledConnectionLifetime = TimeSpan.FromMilliseconds(10 /* total guess */) });
+
+        if(IsAmazonECRRegistry)
+        {
+            clientHandler = new AmazonECRMessageHandler(clientHandler);
+        }
+
         HttpClient client = new(clientHandler);
 
         client.DefaultRequestHeaders.Accept.Clear();
@@ -242,7 +287,9 @@ public async Task Push(Image x, string name, string? tag, string baseName, Actio
 
         HttpClient client = GetClient();
         var reg = this;
-        await Task.WhenAll(x.LayerDescriptors.Select(async descriptor => {
+
+        Func<Descriptor, Task> uploadLayerFunc = async (descriptor) =>
+        {
             string digest = descriptor.Digest;
             logProgressMessage($"Uploading layer {digest} to {reg.RegistryName}");
             if (await reg.BlobAlreadyUploaded(name, digest, client))
@@ -269,7 +316,21 @@ await Task.WhenAll(x.LayerDescriptors.Select(async descriptor => {
                 await reg.Push(Layer.FromDescriptor(descriptor), name, logProgressMessage);
                 logProgressMessage($"Finished uploading layer {digest} to {reg.RegistryName}");
             }
-        }));
+        };
+
+        // Pushing to ECR uses a much larger chunk size. To avoid getting too many socket disconnects trying to do too many
+        // parallel uploads be more conservative and upload one layer at a time.
+        if(IsAmazonECRRegistry)
+        {
+            foreach(var descriptor in x.LayerDescriptors)
+            {
+                await uploadLayerFunc(descriptor);
+            }
+        }
+        else
+        {
+            await Task.WhenAll(x.LayerDescriptors.Select(descriptor => uploadLayerFunc(descriptor)));
+        }
 
         using (MemoryStream stringStream = new MemoryStream(Encoding.UTF8.GetBytes(x.config.ToJsonString())))
         {

Test.Microsoft.NET.Build.Containers.Filesystem/RegistryTests.cs

@@ -19,4 +19,19 @@ public async Task GetFromRegistry()
 
         Assert.IsNotNull(downloadedImage);
     }
+
+    [DataRow("public.ecr.aws", true)]
+    [DataRow("123412341234.dkr.ecr.us-west-2.amazonaws.com", true)]
+    [DataRow("123412341234.dkr.ecr-fips.us-west-2.amazonaws.com", true)]
+    [DataRow("notvalid.dkr.ecr.us-west-2.amazonaws.com", false)]
+    [DataRow("1111.dkr.ecr.us-west-2.amazonaws.com", false)]
+    [DataRow("mcr.microsoft.com", false)]
+    [DataRow("localhost", false)]
+    [DataRow("hub", false)]
+    [TestMethod]
+    public void CheckIfAmazonECR(string registryName, bool isECR)
+    {
+        Registry registry = new Registry(ContainerHelpers.TryExpandRegistryToUri(registryName));
+        Assert.AreEqual(isECR, registry.IsAmazonECRRegistry);
+    }
 }