OCI: always read channel until end (#284)

* Use URLSession.dataTask() with delegate instead of URLSession.bytes()

* Use URLSession.shared instead of creating a new one each time

Sources/tart/Fetcher.swift

@@ -0,0 +1,46 @@
+import Foundation
+import AsyncAlgorithms
+
+class Fetcher: NSObject, URLSessionTaskDelegate, URLSessionDelegate, URLSessionDataDelegate {
+    let responseCh = AsyncThrowingChannel<URLResponse, Error>()
+    let dataCh = AsyncThrowingChannel<Data, Error>()
+
+    func fetch(_ request: URLRequest) async throws -> (AsyncThrowingChannel<Data, Error>, URLResponse) {
+        let task = URLSession.shared.dataTask(with: request)
+        task.delegate = self
+        task.resume()
+
+        // Wait for the response and only then return
+        var iter = responseCh.makeAsyncIterator()
+        let response = try await iter.next()!
+
+        return (dataCh, response)
+    }
+
+    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse) async -> URLSession.ResponseDisposition {
+        await responseCh.send(response)
+
+        return .allow
+    }
+
+    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
+        let sema = DispatchSemaphore(value: 0)
+
+        Task {
+            await dataCh.send(data)
+            sema.signal()
+        }
+
+        sema.wait()
+    }
+
+    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
+        if let error = error {
+            // Premature termination
+            responseCh.fail(error)
+            dataCh.fail(error)
+        } else {
+            dataCh.finish()
+        }
+    }
+}

Sources/tart/OCI/Registry.swift

@@ -2,8 +2,6 @@ import Foundation
 import Algorithms
 import AsyncAlgorithms
 
-let chunkSizeBytes = 1 * 1024 * 1024
-
 enum RegistryError: Error {
   case UnexpectedHTTPStatusCode(when: String, code: Int, details: String = "")
   case MissingLocationHeader
@@ -31,11 +29,11 @@ extension Data {
   }
 }
 
-extension URLSession.AsyncBytes {
+extension AsyncThrowingChannel<Data, Error> {
   func asData() async throws -> Data {
     var result = Data()
 
-    for try await chunk in chunks(ofCount: chunkSizeBytes) {
+    for try await chunk in self {
       result += chunk
     }
 
@@ -214,8 +212,8 @@ class Registry {
         parameters: lastChunk ? ["digest": digest] : [:],
         body: chunk
       )
-      let expectedStatus = lastChunk ? HTTPCode.Created.rawValue : HTTPCode.Accepted.rawValue
+      // always accept both statuses since AWS ECR is not following specification
-      if response.statusCode != expectedStatus {
+      if response.statusCode != HTTPCode.Created.rawValue && response.statusCode != HTTPCode.Accepted.rawValue {
         throw RegistryError.UnexpectedHTTPStatusCode(when: "streaming blob to \(uploadLocation)",
           code: response.statusCode, details: data.asText())
       }
@@ -228,14 +226,14 @@ class Registry {
   }
 
   public func pullBlob(_ digest: String, handler: (Data) throws -> Void) async throws {
-    let (bytes, response) = try await bytesRequest(.GET, endpointURL("\(namespace)/blobs/\(digest)"))
+    let (channel, response) = try await channelRequest(.GET, endpointURL("\(namespace)/blobs/\(digest)"))
     if response.statusCode != HTTPCode.Ok.rawValue {
-      let body = try await bytes.asData().asText()
+      let body = try await channel.asData().asText()
       throw RegistryError.UnexpectedHTTPStatusCode(when: "pulling blob", code: response.statusCode,
         details: body)
     }
 
-    for try await part in bytes.chunks(ofCount: chunkSizeBytes) {
+    for try await part in channel {
       try Task.checkCancellation()
 
       try handler(Data(part))
@@ -256,20 +254,20 @@ class Registry {
           body: Data? = nil,
           doAuth: Bool = true
   ) async throws -> (Data, HTTPURLResponse) {
-    let (bytes, response) = try await bytesRequest(method, urlComponents,
+    let (channel, response) = try await channelRequest(method, urlComponents,
             headers: headers, parameters: parameters, body: body, doAuth: doAuth)
 
-    return (try await bytes.asData(), response)
+    return (try await channel.asData(), response)
   }
 
-  private func bytesRequest(
+  private func channelRequest(
     _ method: HTTPMethod,
     _ urlComponents: URLComponents,
     headers: Dictionary<String, String> = Dictionary(),
     parameters: Dictionary<String, String> = Dictionary(),
     body: Data? = nil,
     doAuth: Bool = true
-  ) async throws -> (URLSession.AsyncBytes, HTTPURLResponse) {
+  ) async throws -> (AsyncThrowingChannel<Data, Error>, HTTPURLResponse) {
     var urlComponents = urlComponents
 
     if urlComponents.queryItems == nil && !parameters.isEmpty {
@@ -294,14 +292,15 @@ class Registry {
       currentAuthToken = nil
     }
 
-    var (bytes, response) = try await authAwareRequest(request: request)
+    var (channel, response) = try await authAwareRequest(request: request)
 
     if doAuth && response.statusCode == HTTPCode.Unauthorized.rawValue {
+      _ = try await channel.asData()
       try await auth(response: response)
-      (bytes, response) = try await authAwareRequest(request: request)
+      (channel, response) = try await authAwareRequest(request: request)
     }
 
-    return (bytes, response)
+    return (channel, response)
   }
 
   private func auth(response: HTTPURLResponse) async throws {
@@ -373,7 +372,7 @@ class Registry {
     return nil
   }
 
-  private func authAwareRequest(request: URLRequest) async throws -> (URLSession.AsyncBytes, HTTPURLResponse) {
+  private func authAwareRequest(request: URLRequest) async throws -> (AsyncThrowingChannel<Data, Error>, HTTPURLResponse) {
     var request = request
 
     if let token = currentAuthToken {
@@ -381,8 +380,8 @@ class Registry {
       request.addValue(value, forHTTPHeaderField: name)
     }
 
-    let (bytes, response) = try await URLSession.shared.bytes(for: request)
+    let (channel, response) = try await Fetcher().fetch(request)
 
-    return (bytes, response as! HTTPURLResponse)
+    return (channel, response as! HTTPURLResponse)
   }
 }

Sources/tart/VM.swift

@@ -63,7 +63,7 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
 
   static func retrieveIPSW(remoteURL: URL) async throws -> URL {
     // Check if we already have this IPSW in cache
-    let (bytes, response) = try await URLSession.shared.bytes(from: remoteURL)
+    let (channel, response) = try await Fetcher().fetch(URLRequest(url: remoteURL))
 
     if let hash = (response as! HTTPURLResponse).value(forHTTPHeaderField: "x-amz-meta-digest-sha256") {
       let ipswLocation = try IPSWCache().locationFor(fileName: "sha256:\(hash).ipsw")
@@ -90,7 +90,7 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
     let fileHandle = try FileHandle(forWritingTo: temporaryLocation)
     let digest = Digest()
 
-    for try await chunk in bytes.chunks(ofCount: chunkSizeBytes) {
+    for try await chunk in channel {
       let chunkAsData = Data(chunk)
       fileHandle.write(chunkAsData)
       digest.update(chunkAsData)