Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Registry: limit the text output on unexpected status code #981

Merged
merged 4 commits into from
Dec 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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