diff --git a/Sources/OmarchyAppleInstaller/VerifiedArtifactStager.swift b/Sources/OmarchyAppleInstaller/VerifiedArtifactStager.swift index 3411589..a8d1c58 100644 --- a/Sources/OmarchyAppleInstaller/VerifiedArtifactStager.swift +++ b/Sources/OmarchyAppleInstaller/VerifiedArtifactStager.swift @@ -231,7 +231,8 @@ public enum ArtifactStageError: Error, Equatable, Sendable { /// Multi-part payloads are downloaded part by part, then streamed together into /// a single pending file whose SHA-256 is computed in the same pass. Assembly /// therefore needs about twice the payload size on disk until the parts are -/// deleted, and interrupted transfers resume at part granularity. +/// deleted. Within one run a dropped connection resumes at byte granularity +/// (see `RangedArtifactDownloader`); across runs, at part granularity. public struct VerifiedArtifactStager: Sendable { /// Disk room required while a payload is assembled from parts (payload plus /// the pending concatenation) before the parts are deleted. @@ -257,7 +258,7 @@ public struct VerifiedArtifactStager: Sendable { private let bundledEngineDirectory: URL? public init() { - downloader = ProgressReportingArtifactDownloader() + downloader = RangedArtifactDownloader() promoter = AtomicArtifactFilePromoter() bundledEngineDirectory = Bundle.main.resourceURL?.appendingPathComponent("Engine/artifacts") } @@ -1009,3 +1010,340 @@ final class ProgressReportingArtifactDownloader: NSObject, ArtifactDownloading, transfer.continuation.resume(with: result) } } + +/// Downloads an artifact as parallel HTTP byte ranges written in place, each +/// over its own session, so one slow or dropped connection neither stalls nor +/// restarts the whole file. A dropped range resumes from its last byte. A +/// server that ignores Range falls back to the single-stream downloader. The +/// stager still checks the size and SHA-256 of the result. +final class RangedArtifactDownloader: NSObject, ArtifactDownloading, URLSessionDataDelegate, + @unchecked Sendable +{ + struct Tuning: Sendable { + var rangeBytes: UInt64 = 64 * 1_048_576 + var connections = 4 + var attemptsWithoutProgress = 4 + var retryDelay: Duration = .seconds(1) + } + + private enum RangeFailure: Error { + case ignored + } + + private struct Attempt { + let written: UInt64 + let failure: (any Error)? + } + + private struct Transfer { + let continuation: CheckedContinuation + let job: Job + let range: ClosedRange + var written: UInt64 = 0 + var failure: (any Error)? + } + + private let configuration: URLSessionConfiguration + private let tuning: Tuning + private let fallback: any ArtifactDownloading + private let lock = NSLock() + private var transfers = [ObjectIdentifier: Transfer]() + + init( + configuration: URLSessionConfiguration = .ephemeral, + tuning: Tuning = Tuning(), + fallback: (any ArtifactDownloading)? = nil + ) { + self.configuration = configuration + self.tuning = tuning + self.fallback = + fallback ?? ProgressReportingArtifactDownloader(configuration: configuration) + super.init() + } + + func download(from sourceURL: URL) async throws -> URL { + try await fallback.download(from: sourceURL) + } + + func download( + from sourceURL: URL, + expectedSizeBytes: UInt64, + onBytes: ArtifactByteProgressHandler? + ) async throws -> URL { + guard expectedSizeBytes > 0, expectedSizeBytes != UInt64.max else { + return try await fallback.download( + from: sourceURL, expectedSizeBytes: expectedSizeBytes, onBytes: onBytes) + } + try Task.checkCancellation() + + let output = FileManager.default.temporaryDirectory.appendingPathComponent( + "omarchy-artifact-download-\(UUID().uuidString.lowercased())", + isDirectory: false + ) + let descriptor = Darwin.open(output.path, O_WRONLY | O_CREAT | O_EXCL, 0o600) + guard descriptor >= 0 else { + throw POSIXError(POSIXErrorCode(rawValue: Darwin.errno) ?? .EIO) + } + var keepOutput = false + defer { + Darwin.close(descriptor) + if !keepOutput { + try? FileManager.default.removeItem(at: output) + } + } + guard Darwin.ftruncate(descriptor, off_t(expectedSizeBytes)) == 0 else { + throw POSIXError(POSIXErrorCode(rawValue: Darwin.errno) ?? .EIO) + } + + let job = Job( + descriptor: descriptor, + sizeBytes: expectedSizeBytes, + rangeBytes: max(1, tuning.rangeBytes), + onBytes: onBytes + ) + do { + try await withThrowingTaskGroup(of: Void.self) { group in + for _ in 0.. 0 ? 1 : failuresWithoutProgress + 1 + guard failuresWithoutProgress < tuning.attemptsWithoutProgress else { + throw failure + } + try await Task.sleep(for: tuning.retryDelay) + } + } + } + + private func fetch( + _ range: ClosedRange, + of url: URL, + for job: Job, + in session: URLSession + ) async -> Attempt { + var request = URLRequest(url: url) + request.setValue( + "bytes=\(range.lowerBound)-\(range.upperBound)", forHTTPHeaderField: "Range") + let task = session.dataTask(with: request) + let key = ObjectIdentifier(task) + + return await withTaskCancellationHandler { + await withCheckedContinuation { continuation in + lock.withLock { + transfers[key] = Transfer(continuation: continuation, job: job, range: range) + } + task.resume() + } + } onCancel: { + task.cancel() + } + } + + func urlSession( + _ session: URLSession, + dataTask: URLSessionDataTask, + didReceive response: URLResponse, + completionHandler: @escaping @Sendable (URLSession.ResponseDisposition) -> Void + ) { + let key = ObjectIdentifier(dataTask) + let failure: (any Error)? = lock.withLock { + guard var transfer = transfers[key] else { + return nil + } + transfer.failure = Self.rangeFailure( + response, range: transfer.range, sizeBytes: transfer.job.sizeBytes) + transfers[key] = transfer + return transfer.failure + } + completionHandler(failure == nil ? .allow : .cancel) + } + + func urlSession( + _ session: URLSession, + dataTask: URLSessionDataTask, + didReceive data: Data + ) { + let key = ObjectIdentifier(dataTask) + guard let transfer = lock.withLock({ transfers[key] }), transfer.failure == nil else { + return + } + let length = UInt64(transfer.range.count) + let incoming = UInt64(data.count) + guard transfer.written + incoming <= length else { + record( + ArtifactStageError.sizeMismatch(expected: length, actual: transfer.written + incoming), + for: key) + dataTask.cancel() + return + } + let offset = transfer.range.lowerBound + transfer.written + if let failure = transfer.job.write(data, at: offset) { + record(failure, for: key) + dataTask.cancel() + return + } + lock.withLock { transfers[key]?.written += incoming } + transfer.job.received(incoming) + } + + func urlSession( + _ session: URLSession, + task: URLSessionTask, + didCompleteWithError error: (any Error)? + ) { + let key = ObjectIdentifier(task) + guard let transfer = lock.withLock({ transfers.removeValue(forKey: key) }) else { + return + } + let failure: (any Error)? + if let recorded = transfer.failure { + failure = recorded + } else if let error { + failure = (error as? URLError)?.code == .cancelled ? CancellationError() : error + } else if transfer.written != UInt64(transfer.range.count) { + failure = URLError(.networkConnectionLost) + } else { + failure = nil + } + transfer.continuation.resume(returning: Attempt(written: transfer.written, failure: failure)) + } + + private func record(_ failure: any Error, for key: ObjectIdentifier) { + lock.withLock { transfers[key]?.failure = failure } + } + + /// Accepts only a 206 for exactly the requested bytes of an object of the + /// pinned size. A 200 means the server ignored Range. + private static func rangeFailure( + _ response: URLResponse, + range: ClosedRange, + sizeBytes: UInt64 + ) -> (any Error)? { + guard let http = response as? HTTPURLResponse else { + return ArtifactStageError.unexpectedHTTPStatus(0) + } + if http.statusCode == 200 { + return RangeFailure.ignored + } + guard http.statusCode == 206 else { + return ArtifactStageError.unexpectedHTTPStatus(http.statusCode) + } + // Content-Range: bytes -/ + let fields = (http.value(forHTTPHeaderField: "Content-Range") ?? "") + .split(whereSeparator: { " -/".contains($0) }) + guard fields.count == 4, fields[0] == "bytes", + let first = UInt64(fields[1]), let last = UInt64(fields[2]), + let total = UInt64(fields[3]) + else { + return RangeFailure.ignored + } + guard total == sizeBytes else { + return ArtifactStageError.sizeMismatch(expected: sizeBytes, actual: total) + } + return first == range.lowerBound && last == range.upperBound ? nil : RangeFailure.ignored + } + + /// One download: hands out ranges, writes bytes in place, and meters progress. + private final class Job: @unchecked Sendable { + private static let reportIntervalSeconds: TimeInterval = 0.25 + + let sizeBytes: UInt64 + let rangeCount: Int + private let descriptor: Int32 + private let rangeBytes: UInt64 + private let onBytes: ArtifactByteProgressHandler? + private let lock = NSLock() + private var nextRange = 0 + private var receivedBytes: UInt64 = 0 + private var lastReportedAt: Date? + + init( + descriptor: Int32, + sizeBytes: UInt64, + rangeBytes: UInt64, + onBytes: ArtifactByteProgressHandler? + ) { + self.descriptor = descriptor + self.sizeBytes = sizeBytes + self.rangeBytes = rangeBytes + self.onBytes = onBytes + rangeCount = Int((sizeBytes + rangeBytes - 1) / rangeBytes) + } + + func claim() -> ClosedRange? { + lock.withLock { + guard nextRange < rangeCount else { + return nil + } + let first = UInt64(nextRange) * rangeBytes + nextRange += 1 + return first...(min(first + rangeBytes, sizeBytes) - 1) + } + } + + func write(_ data: Data, at offset: UInt64) -> (any Error)? { + data.withUnsafeBytes { buffer -> (any Error)? in + var done = 0 + while done < buffer.count { + let count = Darwin.pwrite( + descriptor, buffer.baseAddress! + done, buffer.count - done, + off_t(offset) + off_t(done)) + guard count > 0 else { + return POSIXError(POSIXErrorCode(rawValue: Darwin.errno) ?? .EIO) + } + done += count + } + return nil + } + } + + func received(_ count: UInt64) { + guard let onBytes else { + return + } + let report: UInt64? = lock.withLock { + receivedBytes += count + let now = Date() + let due = + receivedBytes == sizeBytes + || lastReportedAt.map { now.timeIntervalSince($0) >= Self.reportIntervalSeconds } + ?? true + guard due else { + return nil + } + lastReportedAt = now + return receivedBytes + } + if let report { + onBytes(report) + } + } + } +} diff --git a/Tests/OmarchyAppleInstallerTrustCoreTests/VerifiedArtifactStagerTests.swift b/Tests/OmarchyAppleInstallerTrustCoreTests/VerifiedArtifactStagerTests.swift index 0c56fe3..af0385d 100644 --- a/Tests/OmarchyAppleInstallerTrustCoreTests/VerifiedArtifactStagerTests.swift +++ b/Tests/OmarchyAppleInstallerTrustCoreTests/VerifiedArtifactStagerTests.swift @@ -615,6 +615,65 @@ final class VerifiedArtifactStagerTests: XCTestCase { } } + func testRangedDownloadResumesADroppedRangeFromItsLastByte() async throws { + let payload = rangeFixture(count: 300_000) + let server = try RangeHTTPServer(payload: payload, drops: 1) + let url = try await server.start() + defer { server.stop() } + let downloader = RangedArtifactDownloader( + tuning: .init(rangeBytes: 1_000_000, connections: 1, retryDelay: .zero)) + + let file = try await downloader.download( + from: url, expectedSizeBytes: UInt64(payload.count), onBytes: nil) + defer { try? FileManager.default.removeItem(at: file) } + + XCTAssertEqual(try Data(contentsOf: file), payload) + XCTAssertEqual(server.rangeHeaders, ["bytes=0-299999", "bytes=150000-299999"]) + } + + func testRangesRunInParallelOverSeparateConnectionsAndProgressReachesTheTotal() async throws { + let payload = rangeFixture(count: 1_048_576) + let server = try RangeHTTPServer(payload: payload, responseDelay: .milliseconds(50)) + let url = try await server.start() + defer { server.stop() } + let downloader = RangedArtifactDownloader( + tuning: .init(rangeBytes: 65_536, connections: 4, retryDelay: .zero)) + let reported = ByteRecorder() + + let file = try await downloader.download( + from: url, expectedSizeBytes: UInt64(payload.count), onBytes: reported.handler) + defer { try? FileManager.default.removeItem(at: file) } + + XCTAssertEqual(try Data(contentsOf: file), payload) + let expected = (0..<16).map { "bytes=\($0 * 65_536)-\($0 * 65_536 + 65_535)" } + XCTAssertEqual(Set(server.rangeHeaders.compactMap { $0 }), Set(expected)) + XCTAssertEqual(server.rangeHeaders.count, 16) + XCTAssertGreaterThanOrEqual(server.peakActiveRequests, 2) + XCTAssertGreaterThanOrEqual(server.connectionCount, 2) + XCTAssertEqual(reported.values.last, UInt64(payload.count)) + XCTAssertEqual(reported.values, reported.values.sorted()) + } + + func testServerThatIgnoresRangeFallsBackToOneStream() async throws { + let payload = rangeFixture(count: 200_000) + let server = try RangeHTTPServer(payload: payload, honorsRange: false) + let url = try await server.start() + defer { server.stop() } + let downloader = RangedArtifactDownloader( + tuning: .init(rangeBytes: 65_536, connections: 2, retryDelay: .zero)) + + let file = try await downloader.download( + from: url, expectedSizeBytes: UInt64(payload.count), onBytes: nil) + defer { try? FileManager.default.removeItem(at: file) } + + XCTAssertEqual(try Data(contentsOf: file), payload) + XCTAssertTrue(server.rangeHeaders.contains(nil), "the fallback asks for the whole file") + } + + private func rangeFixture(count: Int) -> Data { + Data((0..> 8)) }) + } + private func descriptor( source: String = "https://example.com/installer.tar.gz", fileName: String = "installer.tar.gz", @@ -846,3 +905,164 @@ private final class SlowHTTPServer: @unchecked Sendable { pending?.resume(with: result) } } + +private final class ByteRecorder: @unchecked Sendable { + private let lock = NSLock() + private var storage = [UInt64]() + + var values: [UInt64] { lock.withLock { storage } } + + var handler: ArtifactByteProgressHandler { + { [self] bytes in lock.withLock { storage.append(bytes) } } + } +} + +/// Serves one payload over HTTP/1.1 with keep-alive. Honors `Range: bytes=a-b` +/// unless told not to, can cut the first `drops` responses off halfway (with a +/// clean FIN), and records every request's Range header. +private final class RangeHTTPServer: @unchecked Sendable { + private let payload: Data + private let honorsRange: Bool + private let responseDelay: DispatchTimeInterval + private let listener: NWListener + private let lock = NSLock() + private var connections: [NWConnection] = [] + private var dropsLeft: Int + private var activeRequests = 0 + private var requests = [String?]() + private var peak = 0 + private var startContinuation: CheckedContinuation? + + init( + payload: Data, + honorsRange: Bool = true, + drops: Int = 0, + responseDelay: DispatchTimeInterval = .never + ) throws { + self.payload = payload + self.honorsRange = honorsRange + self.dropsLeft = drops + self.responseDelay = responseDelay + listener = try NWListener(using: .tcp, on: .any) + } + + var rangeHeaders: [String?] { lock.withLock { requests } } + var peakActiveRequests: Int { lock.withLock { peak } } + var connectionCount: Int { lock.withLock { connections.count } } + + func start() async throws -> URL { + try await withCheckedThrowingContinuation { continuation in + lock.withLock { startContinuation = continuation } + listener.stateUpdateHandler = { [weak self] state in + guard let self else { return } + switch state { + case .ready: + guard let port = self.listener.port else { return } + self.resumeStart( + .success(URL(string: "http://127.0.0.1:\(port.rawValue)/payload.bin")!)) + case .failed(let error): + self.resumeStart(.failure(error)) + default: + break + } + } + listener.newConnectionHandler = { [weak self] connection in + guard let self else { return } + self.lock.withLock { self.connections.append(connection) } + connection.start(queue: .global()) + self.receiveRequest(on: connection, buffer: Data()) + } + listener.start(queue: .global()) + } + } + + func stop() { + let pending = lock.withLock { () -> [NWConnection] in + let current = connections + connections.removeAll() + return current + } + for connection in pending { + connection.cancel() + } + listener.cancel() + } + + private func receiveRequest(on connection: NWConnection, buffer: Data) { + connection.receive(minimumIncompleteLength: 1, maximumLength: 65_536) { + [weak self] data, _, isComplete, error in + guard let self else { return } + var buffer = buffer + if let data { + buffer.append(data) + } + if let end = buffer.range(of: Data("\r\n\r\n".utf8)) { + self.respond(to: String(decoding: buffer[.. Bool in + requests.append(range) + activeRequests += 1 + peak = max(peak, activeRequests) + guard dropsLeft > 0 else { return false } + dropsLeft -= 1 + return true + } + let head = Data("HTTP/1.1 \(status)\r\nContent-Length: \(body.count)\r\n\(extra)\r\n".utf8) + let sent = drop ? body.prefix(body.count / 2) : body + let send = { [weak self] in + connection.send( + content: head + sent, + completion: .contentProcessed { _ in + guard let self else { return } + self.lock.withLock { self.activeRequests -= 1 } + if drop { + // Close only after the client has read the half body, like a link + // that dies mid-transfer. + DispatchQueue.global().asyncAfter(deadline: .now() + .milliseconds(200)) { + connection.send( + content: nil, contentContext: .finalMessage, isComplete: true, + completion: .contentProcessed { _ in }) + } + } else { + self.receiveRequest(on: connection, buffer: Data()) + } + }) + } + if responseDelay == .never { + send() + } else { + DispatchQueue.global().asyncAfter(deadline: .now() + responseDelay, execute: send) + } + } + + private func resumeStart(_ result: Result) { + let pending = lock.withLock { () -> CheckedContinuation? in + let value = startContinuation + startContinuation = nil + return value + } + pending?.resume(with: result) + } +}