ContainerRegistry: Remember initial auth challenge and send credentials eagerly (#30)

### Motivation

Currently, `ContainerRegistry` handles authentication lazily.   First it tries an unauthenticated request.   If the registry responds with a challenge, `ContainerRegistry` gets a suitable credential from `.netrc` or from an authentication service and re-tries the request.

When pushing large blobs to some registries, the client can get stuck in a state where it has received an authentication challenge but is still waiting to finish the blob upload.   The symptoms are that metadata uploads and very small image layers can be uploaded, but pushing larger images stalls and times out, as reported in #17.

### Modifications

When pushing large blobs, the registry's `401` challenge arrives after the client has sent some of the blob data, but before it has completed the upload.   With some registries, `URLSession` gets stuck in this state.   Implementing the `didReceiveResponse` delegate allows the client to see the challenge and cancel the task but in some cases, mainly on Linux, `URLSession` can still get stuck.

A more reliable solution to this problem is to handle challenges eagerly.   Before pushing any data, the client checks whether the registry supports the v2 distribution protocol.  If it is challenged during this check, it will then eagerly send a suitable authentication token with all subsequent requests to the registry.   This avoids the state where the client receives a challenge part-way through uploading a large blob and gets stuck.

### Result

It is possible to upload large images to ECR and Docker Hub from macOS hosts, and to Docker Hub from Linux hosts.   Uploading to ECR from Linux hosts continues to lead to a timeout.

Fixes #17 

### Test Plan

Automated tests continue to pass;  tested manually with several different registries.

Author: euanh

Sources/ContainerRegistry/AuthHandler.swift

@@ -110,6 +110,22 @@ func parseChallenge(_ s: String) throws -> BearerChallenge {
     return res
 }
 
+public enum AuthChallenge: Equatable {
+    case none
+    case basic(String)
+    case bearer(String)
+
+    init(challenge: String) {
+        if challenge.lowercased().starts(with: "basic") {
+            self = .basic(challenge)
+        } else if challenge.lowercased().starts(with: "bearer") {
+            self = .bearer(challenge)
+        } else {
+            self = .none
+        }
+    }
+}
+
 /// AuthHandler manages provides credentials for HTTP requests
 public struct AuthHandler {
     var username: String?
@@ -127,11 +143,9 @@ public struct AuthHandler {
         self.auth = auth
     }
 
-    /// Get locally-configured credentials, such as netrc or username/password, for a request
-    func localCredentials(for request: HTTPRequest) -> String? {
-        guard let requestURL = request.url else { return nil }
-
-        if let netrcEntry = auth?.httpAuthorizationHeader(for: requestURL) { return netrcEntry }
+    /// Get locally-configured credentials, such as netrc or username/password, for a host
+    func localCredentials(forURL url: URL) -> String? {
+        if let netrcEntry = auth?.httpAuthorizationHeader(for: url) { return netrcEntry }
 
         if let username, let password {
             let authorization = Data("\(username):\(password)".utf8).base64EncodedString()
@@ -142,52 +156,35 @@ public struct AuthHandler {
         return nil
     }
 
-    /// Add authorization to an HTTP rquest before it has been sent to a server.
-    /// Currently this function always passes the request back unmodified, to trigger a challenge.
-    /// In future it could provide cached responses from previous challenges.
-    /// - Parameter request: The request to authorize.
-    /// - Returns: The request, with an appropriate authorization header added, or nil if no credentials are available.
-    public func auth(for request: HTTPRequest) -> HTTPRequest? { nil }
-
-    /// Add authorization to an HTTP rquest in response to a challenge from the server.
-    /// - Parameters:
-    ///   - request: The reuqest to authorize.
-    ///   - challenge: The server's challeng.
-    ///   - client: An HTTP client, used to retrieve tokens if necessary.
-    /// - Returns: The request, with an appropriate authorization header added, or nil if no credentials are available.
-    /// - Throws: If an error occurs while retrieving a credential.
-    public func auth(for request: HTTPRequest, withChallenge challenge: String, usingClient client: HTTPClient)
-        async throws -> HTTPRequest?
-    {
-        if challenge.lowercased().starts(with: "basic") {
-            guard let authHeader = localCredentials(for: request) else { return nil }
-            var request = request
-            request.headerFields[.authorization] = authHeader
-            return request
-
-        } else if challenge.lowercased().starts(with: "bearer") {
+    public func auth(
+        registry: URL,
+        repository: String,
+        actions: [String],
+        withScheme scheme: AuthChallenge,
+        usingClient client: HTTPClient
+    ) async throws -> String? {
+        switch scheme {
+        case .none: return nil
+        case .basic: return localCredentials(forURL: registry)
+
+        case .bearer(let challenge):
             // Preemptively offer suitable basic auth credentials to the token server.
             // Instead of challenging, public token servers often return anonymous tokens when no credentials are offered.
             // These tokens allow pull access to public repositories, but attempts to push will fail with 'unauthorized'.
             // There is no obvious prompt for the client to retry with authentication.
-            let parsedChallenge = try parseChallenge(
+            var parsedChallenge = try parseChallenge(
                 challenge.dropFirst("bearer".count).trimmingCharacters(in: .whitespacesAndNewlines)
             )
+            parsedChallenge.scope = ["repository:\(repository):\(actions.joined(separator: ","))"]
             guard let challengeURL = parsedChallenge.url else { return nil }
             var tokenRequest = HTTPRequest(url: challengeURL)
-            if let credentials = localCredentials(for: tokenRequest) {
+            if let credentials = localCredentials(forURL: challengeURL) {
                 tokenRequest.headerFields[.authorization] = credentials
             }
 
             let (data, _) = try await client.executeRequestThrowing(tokenRequest, expectingStatus: .ok)
             let tokenResponse = try JSONDecoder().decode(BearerTokenResponse.self, from: data)
-            var request = request
-            request.headerFields[.authorization] = "Bearer \(tokenResponse.token)"
-            return request
-
-        } else {
-            // No other authentication methods available
-            return nil
+            return "Bearer \(tokenResponse.token)"
         }
     }
 }

Sources/ContainerRegistry/CheckAPI.swift

@@ -12,23 +12,29 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Foundation
+
 public extension RegistryClient {
-    /// Returns a boolean value indicating whether the registry supports v2 of the distribution specification.
-    /// - Returns: `true` if the registry supports the distribution specification, otherwise `false`.
-    func checkAPI() async throws -> Bool {
+    /// Checks whether the registry supports v2 of the distribution specification.
+    /// - Returns: an `true` if the registry supports the distribution specification.
+    /// - Throws: if the registry does not support the distribution specification.
+    static func checkAPI(client: HTTPClient, registryURL: URL) async throws -> AuthChallenge {
         // See https://github.com/opencontainers/distribution-spec/blob/main/spec.md#determining-support
+
         // The registry indicates that it supports the v2 protocol by returning a 200 OK response.
         // Many registries also set `Content-Type: application/json` and return empty JSON objects,
         // but this is not required and some do not.
         // The registry may require authentication on this endpoint.
+
         do {
             // Using the bare HTTP client because this is the only endpoint which does not include a repository path
-            let _ = try await executeRequestThrowing(
-                .get("", url: registryURL.distributionEndpoint),
-                expectingStatus: .ok,
-                decodingErrors: [.unauthorized, .notFound]
+            // and to avoid RegistryClient's auth handling
+            let _ = try await client.executeRequestThrowing(
+                .get(registryURL.distributionEndpoint, withAuthorization: nil),
+                expectingStatus: .ok
             )
-            return true
-        } catch HTTPClientError.unexpectedStatusCode(status: .notFound, _, _) { return false }
+            return .none
+
+        } catch HTTPClientError.authenticationChallenge(let challenge, _, _) { return .init(challenge: challenge) }
     }
 }

Sources/ContainerRegistry/HTTPClient.swift

@@ -159,4 +159,13 @@ extension HTTPRequest {
         //    https://developer.apple.com/forums/thread/89811
         if let authorization { headerFields[.authorization] = authorization }
     }
+
+    static func get(
+        _ url: URL,
+        accepting: [String] = [],
+        contentType: String? = nil,
+        withAuthorization authorization: String? = nil
+    ) -> HTTPRequest {
+        .init(method: .get, url: url, accepting: accepting, contentType: contentType, withAuthorization: authorization)
+    }
 }

Sources/ContainerRegistry/RegistryClient.swift

@@ -50,6 +50,7 @@ public struct RegistryClient {
 
     /// Authentication handler
     var auth: AuthHandler?
+    var authChallenge: AuthChallenge
 
     var encoder: JSONEncoder
     var decoder: JSONDecoder
@@ -92,7 +93,7 @@ public struct RegistryClient {
         self.decoder = decoder ?? JSONDecoder()
 
         // Verify that we can talk to the registry
-        _ = try await checkAPI()
+        self.authChallenge = try await RegistryClient.checkAPI(client: self.client, registryURL: self.registryURL)
     }
 
     /// Creates a new RegistryClient, constructing a suitable URLSession-based client.
@@ -142,15 +143,14 @@ extension RegistryClient {
         var method: HTTPRequest.Method  // HTTP method
         var repository: String  // Repository path on the registry
         var destination: Destination  // Destination of the operation: can be a subpath or remote URL
+        var actions: [String]  // Actions required by this operation
         var accepting: [String] = []  // Acceptable response types
         var contentType: String? = nil  // Request data type
 
         func url(relativeTo registry: URL) -> URL {
             switch destination {
             case .url(let url): return url
-            case .subpath(let path):
-                let subpath = registry.distributionEndpoint(forRepository: repository, andEndpoint: path)
-                return subpath
+            case .subpath(let path): return registry.distributionEndpoint(forRepository: repository, andEndpoint: path)
             }
         }
 
@@ -166,6 +166,7 @@ extension RegistryClient {
                 method: .get,
                 repository: repository,
                 destination: .subpath(path),
+                actions: ["pull"],
                 accepting: accepting,
                 contentType: contentType
             )
@@ -182,6 +183,7 @@ extension RegistryClient {
                 method: .get,
                 repository: repository,
                 destination: .url(url),
+                actions: ["pull"],
                 accepting: accepting,
                 contentType: contentType
             )
@@ -198,6 +200,7 @@ extension RegistryClient {
                 method: .head,
                 repository: repository,
                 destination: .subpath(path),
+                actions: ["pull"],
                 accepting: accepting,
                 contentType: contentType
             )
@@ -215,6 +218,7 @@ extension RegistryClient {
                 method: .put,
                 repository: repository,
                 destination: .url(url),
+                actions: ["push", "pull"],
                 accepting: accepting,
                 contentType: contentType
             )
@@ -231,6 +235,7 @@ extension RegistryClient {
                 method: .put,
                 repository: repository,
                 destination: .subpath(path),
+                actions: ["push", "pull"],
                 accepting: accepting,
                 contentType: contentType
             )
@@ -247,6 +252,7 @@ extension RegistryClient {
                 method: .post,
                 repository: repository,
                 destination: .subpath(path),
+                actions: ["push", "pull"],
                 accepting: accepting,
                 contentType: contentType
             )
@@ -268,22 +274,25 @@ extension RegistryClient {
         expectingStatus success: HTTPResponse.Status = .ok,
         decodingErrors errors: [HTTPResponse.Status]
     ) async throws -> (data: Data, response: HTTPResponse) {
+        let authorization = try await auth?
+            .auth(
+                registry: registryURL,
+                repository: operation.repository,
+                actions: operation.actions,
+                withScheme: authChallenge,
+                usingClient: client
+            )
+
         let request = HTTPRequest(
             method: operation.method,
             url: operation.url(relativeTo: registryURL),
             accepting: operation.accepting,
-            contentType: operation.contentType
+            contentType: operation.contentType,
+            withAuthorization: authorization
         )
 
         do {
-            let authenticatedRequest = auth?.auth(for: request) ?? request
-            return try await client.executeRequestThrowing(authenticatedRequest, expectingStatus: success)
-        } catch HTTPClientError.authenticationChallenge(let challenge, let request, let response) {
-            guard
-                let authenticatedRequest = try await auth?
-                    .auth(for: request, withChallenge: challenge, usingClient: client)
-            else { throw HTTPClientError.unauthorized(request: request, response: response) }
-            return try await client.executeRequestThrowing(authenticatedRequest, expectingStatus: success)
+            return try await client.executeRequestThrowing(request, expectingStatus: success)
         } catch HTTPClientError.unexpectedStatusCode(let status, _, let .some(responseData))
             where errors.contains(status)
         {
@@ -330,30 +339,25 @@ extension RegistryClient {
         expectingStatus success: HTTPResponse.Status,
         decodingErrors errors: [HTTPResponse.Status]
     ) async throws -> (data: Data, response: HTTPResponse) {
+        let authorization = try await auth?
+            .auth(
+                registry: registryURL,
+                repository: operation.repository,
+                actions: operation.actions,
+                withScheme: authChallenge,
+                usingClient: client
+            )
+
         let request = HTTPRequest(
             method: operation.method,
             url: operation.url(relativeTo: registryURL),
             accepting: operation.accepting,
-            contentType: operation.contentType
+            contentType: operation.contentType,
+            withAuthorization: authorization
         )
 
         do {
-            let authenticatedRequest = auth?.auth(for: request) ?? request
-            return try await client.executeRequestThrowing(
-                authenticatedRequest,
-                uploading: payload,
-                expectingStatus: success
-            )
-        } catch HTTPClientError.authenticationChallenge(let challenge, let request, let response) {
-            guard
-                let authenticatedRequest = try await auth?
-                    .auth(for: request, withChallenge: challenge, usingClient: client)
-            else { throw HTTPClientError.unauthorized(request: request, response: response) }
-            return try await client.executeRequestThrowing(
-                authenticatedRequest,
-                uploading: payload,
-                expectingStatus: success
-            )
+            return try await client.executeRequestThrowing(request, uploading: payload, expectingStatus: success)
         } catch HTTPClientError.unexpectedStatusCode(let status, _, let .some(responseData))
             where errors.contains(status)
         {

Tests/ContainerRegistryTests/AuthTests.swift

@@ -16,7 +16,7 @@ import Foundation
 import Basics
 import XCTest
 
-class AuthTests: XCTestCase {
+class AuthTests: XCTestCase, @unchecked Sendable {
     // SwiftPM's NetrcAuthorizationProvider does not throw an error if the .netrc file
     // does not exist.   For simplicity the local vendored version does the same.
     func testNonexistentNetrc() async throws {

Tests/ContainerRegistryTests/SmokeTests.swift

@@ -16,7 +16,7 @@ import Foundation
 import ContainerRegistry
 import XCTest
 
-class SmokeTests: XCTestCase {
+class SmokeTests: XCTestCase, @unchecked Sendable {
     // These are basic tests to exercise the main registry operations.
     // The tests assume that a fresh, empty registry instance is available at
     // http://$REGISTRY_HOST:$REGISTRY_PORT
@@ -31,11 +31,6 @@ class SmokeTests: XCTestCase {
         client = try await RegistryClient(registry: "\(registryHost):\(registryPort)", insecure: true)
     }
 
-    func testCheckAPI() async throws {
-        let supported = try await client.checkAPI()
-        XCTAssert(supported)
-    }
-
     func testGetTags() async throws {
         let repository = "testgettags"