spb/zyquo-local Public MIT

1//2//  FileTransfer.swift3//  Zyquo Local4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation1011/// One HTTP transfer of one file into a `.partial` on disk, delegate-backed12/// for chunked throughput. Supports Range resume (server 206), transparent13/// full restarts (server 200), cancellation, and strips the Authorization14/// header when redirected off huggingface.co to the signed CDN.15final class FileTransfer: NSObject, URLSessionDataDelegate, @unchecked Sendable {16    enum TransferError: LocalizedError {17        case badStatus(Int)1819        var errorDescription: String? {20            switch self {21            case .badStatus(let code): "Server returned HTTP \(code)."22            }23        }24    }2526    private let partial: URL27    /// Cumulative bytes on disk, throttled (~10 Hz).28    private let onProgress: @Sendable (Int64) -> Void2930    private var handle: FileHandle?31    private var received: Int64 = 032    private var status = 033    private var rejectedByStatus = false34    private var continuation: CheckedContinuation<Int64, Error>?35    private var lastReport = ContinuousClock.now3637    init(partial: URL, onProgress: @escaping @Sendable (Int64) -> Void) {38        self.partial = partial39        self.onProgress = onProgress40    }4142    /// Runs the transfer starting at `offset` (0 = fresh). Returns final byte43    /// count on disk. Throws CancellationError when the task is cancelled.44    func run(request: URLRequest, offset: Int64) async throws -> Int64 {45        received = offset46        status = 047        rejectedByStatus = false4849        let config = URLSessionConfiguration.default50        config.timeoutIntervalForRequest = 6051        config.networkServiceType = .responsiveData52        let session = URLSession(configuration: config, delegate: self, delegateQueue: nil)53        defer { session.finishTasksAndInvalidate() }5455        var request = request56        if offset > 0 {57            request.setValue("bytes=\(offset)-", forHTTPHeaderField: "Range")58        }59        let task = session.dataTask(with: request)6061        return try await withTaskCancellationHandler {62            try await withCheckedThrowingContinuation { c in63                continuation = c64                task.resume()65            }66        } onCancel: {67            task.cancel()68        }69    }7071    // MARK: - URLSessionDataDelegate7273    func urlSession(74        _ session: URLSession, task: URLSessionTask,75        willPerformHTTPRedirection response: HTTPURLResponse,76        newRequest request: URLRequest,77        completionHandler: @escaping (URLRequest?) -> Void78    ) {79        var request = request80        if request.url?.host != task.originalRequest?.url?.host {81            request.setValue(nil, forHTTPHeaderField: "Authorization")82        }83        completionHandler(request)84    }8586    func urlSession(87        _ session: URLSession, dataTask: URLSessionDataTask,88        didReceive response: URLResponse,89        completionHandler: @escaping (URLSession.ResponseDisposition) -> Void90    ) {91        guard let http = response as? HTTPURLResponse else {92            rejectedByStatus = true93            completionHandler(.cancel)94            return95        }96        status = http.statusCode97        switch http.statusCode {98        case 200:99            // Fresh body (or the server ignored our Range): restart the file.100            received = 0101            FileManager.default.createFile(atPath: partial.path, contents: nil)102            handle = try? FileHandle(forWritingTo: partial)103            completionHandler(handle == nil ? .cancel : .allow)104        case 206:105            handle = try? FileHandle(forWritingTo: partial)106            _ = try? handle?.seekToEnd()107            completionHandler(handle == nil ? .cancel : .allow)108        default:109            rejectedByStatus = true110            completionHandler(.cancel)111        }112    }113114    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {115        guard let handle else { return }116        do {117            try handle.write(contentsOf: data)118            received += Int64(data.count)119        } catch {120            dataTask.cancel()121            return122        }123        let now = ContinuousClock.now124        if lastReport.duration(to: now) > .milliseconds(100) {125            lastReport = now126            onProgress(received)127        }128    }129130    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {131        try? handle?.close()132        handle = nil133        onProgress(received)134        let c = continuation135        continuation = nil136        if rejectedByStatus {137            c?.resume(throwing: TransferError.badStatus(status))138        } else if let error {139            if (error as? URLError)?.code == .cancelled {140                c?.resume(throwing: CancellationError())141            } else {142                c?.resume(throwing: error)143            }144        } else {145            c?.resume(returning: received)146        }147    }148}149