Skip to content

Commit

Permalink
Registry: limit the text output on unexpected status code (#981)
Browse files Browse the repository at this point in the history
* Registry: limit the text output on unexpected status code

* pullBlob(): limit channel read-out on error to 4 KiB

* No need to always read channel until end

This was introduced in #284
because we were blocking in "urlSession(_ session: URLSession, dataTask:
URLSessionDataTask, didReceive data: Data)", which we don't do anymore.

* Fetcher.fetch(): remove "progress" argument as we don't need it anymore
  • Loading branch information
edigaryev authored Dec 20, 2024
1 parent 5a8b48a commit 04c6df2
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 22 deletions.
2 changes: 1 addition & 1 deletion Sources/tart/Fetcher.swift
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ fileprivate var urlSession: URLSession = {
}()

class Fetcher {
static func fetch(_ request: URLRequest, viaFile: Bool = false, progress: Progress? = nil) async throws -> (AsyncThrowingStream<Data, Error>, HTTPURLResponse) {
static func fetch(_ request: URLRequest, viaFile: Bool = false) async throws -> (AsyncThrowingStream<Data, Error>, HTTPURLResponse) {
let task = urlSession.dataTask(with: request)

let delegate = Delegate()
Expand Down
31 changes: 21 additions & 10 deletions Sources/tart/OCI/Registry.swift
Original file line number Diff line number Diff line change
Expand Up @@ -29,14 +29,26 @@ extension Data {
func asText() -> String {
String(decoding: self, as: UTF8.self)
}

func asTextPreview(limit: Int = 1000) -> String {
guard count > limit else {
return asText()
}

return "\(asText().prefix(limit))..."
}
}

extension AsyncThrowingStream<Data, Error> {
func asData() async throws -> Data {
func asData(limitBytes: Int64? = nil) async throws -> Data {
var result = Data()

for try await chunk in self {
result += chunk

if let limitBytes, result.count > limitBytes {
return result
}
}

return result
Expand Down Expand Up @@ -160,7 +172,7 @@ class Registry {
body: manifestJSON)
if response.statusCode != HTTPCode.Created.rawValue {
throw RegistryError.UnexpectedHTTPStatusCode(when: "pushing manifest", code: response.statusCode,
details: data.asText())
details: data.asTextPreview())
}

return Digest.hash(manifestJSON)
Expand All @@ -171,7 +183,7 @@ class Registry {
headers: ["Accept": ociManifestMediaType])
if response.statusCode != HTTPCode.Ok.rawValue {
throw RegistryError.UnexpectedHTTPStatusCode(when: "pulling manifest", code: response.statusCode,
details: data.asText())
details: data.asTextPreview())
}

let manifest = try OCIManifest(fromJSON: data)
Expand All @@ -197,7 +209,7 @@ class Registry {
headers: ["Content-Length": "0"])
if postResponse.statusCode != HTTPCode.Accepted.rawValue {
throw RegistryError.UnexpectedHTTPStatusCode(when: "pushing blob (POST)", code: postResponse.statusCode,
details: data.asText())
details: data.asTextPreview())
}

// Figure out where to upload the blob
Expand All @@ -218,7 +230,7 @@ class Registry {
)
if response.statusCode != HTTPCode.Created.rawValue {
throw RegistryError.UnexpectedHTTPStatusCode(when: "pushing blob (PUT) to \(uploadLocation)",
code: response.statusCode, details: data.asText())
code: response.statusCode, details: data.asTextPreview())
}
return digest
}
Expand All @@ -241,7 +253,7 @@ class Registry {
// always accept both statuses since AWS ECR is not following specification
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())
code: response.statusCode, details: data.asTextPreview())
}
uploadedBytes += chunk.count
// Update location for the next chunk
Expand All @@ -260,7 +272,7 @@ class Registry {
case HTTPCode.NotFound.rawValue:
return false
default:
throw RegistryError.UnexpectedHTTPStatusCode(when: "checking blob", code: response.statusCode, details: data.asText())
throw RegistryError.UnexpectedHTTPStatusCode(when: "checking blob", code: response.statusCode, details: data.asTextPreview())
}
}

Expand All @@ -279,7 +291,7 @@ class Registry {

let (channel, response) = try await channelRequest(.GET, endpointURL("\(namespace)/blobs/\(digest)"), headers: headers, viaFile: true)
if response.statusCode != expectedStatusCode.rawValue {
let body = try await channel.asData().asText()
let body = try await channel.asData(limitBytes: 4096).asTextPreview()
throw RegistryError.UnexpectedHTTPStatusCode(when: "pulling blob", code: response.statusCode,
details: body)
}
Expand Down Expand Up @@ -342,7 +354,6 @@ class Registry {
var (channel, response) = try await authAwareRequest(request: request, viaFile: viaFile, doAuth: doAuth)

if doAuth && response.statusCode == HTTPCode.Unauthorized.rawValue {
_ = try await channel.asData()
try await auth(response: response)
(channel, response) = try await authAwareRequest(request: request, viaFile: viaFile, doAuth: doAuth)
}
Expand Down Expand Up @@ -404,7 +415,7 @@ class Registry {
let (data, response) = try await dataRequest(.GET, authenticateURL, headers: headers, doAuth: false)
if response.statusCode != HTTPCode.Ok.rawValue {
throw RegistryError.AuthFailed(why: "received unexpected HTTP status code \(response.statusCode) "
+ "while retrieving an authentication token", details: data.asText())
+ "while retrieving an authentication token", details: data.asTextPreview())
}

await authenticationKeeper.set(try TokenResponse.parse(fromData: data))
Expand Down
18 changes: 7 additions & 11 deletions Sources/tart/VM.swift
Original file line number Diff line number Diff line change
Expand Up @@ -99,16 +99,13 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
// Download the IPSW
defaultLogger.appendNewLine("Fetching \(remoteURL.lastPathComponent)...")

let downloadProgress = Progress(totalUnitCount: 100)
ProgressObserver(downloadProgress).log(defaultLogger)

let request = URLRequest(url: remoteURL)
let (channel, response) = try await Fetcher.fetch(request, viaFile: true, progress: downloadProgress)
let (channel, response) = try await Fetcher.fetch(request, viaFile: true)

let temporaryLocation = try Config().tartTmpDir.appendingPathComponent(UUID().uuidString + ".ipsw")
defaultLogger.appendNewLine("Computing digest for \(temporaryLocation.path)...")
let digestProgress = Progress(totalUnitCount: response.expectedContentLength)
ProgressObserver(digestProgress).log(defaultLogger)

let progress = Progress(totalUnitCount: response.expectedContentLength)
ProgressObserver(progress).log(defaultLogger)

FileManager.default.createFile(atPath: temporaryLocation.path, contents: nil)
let lock = try FileLock(lockURL: temporaryLocation)
Expand All @@ -118,10 +115,9 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject {
let digest = Digest()

for try await chunk in channel {
let chunkAsData = Data(chunk)
fileHandle.write(chunkAsData)
digest.update(chunkAsData)
digestProgress.completedUnitCount += Int64(chunk.count)
fileHandle.write(chunk)
digest.update(chunk)
progress.completedUnitCount += Int64(chunk.count)
}

try fileHandle.close()
Expand Down

0 comments on commit 04c6df2

Please sign in to comment.