diff --git a/Cotabby/Services/Utilities/ModelDownloadManager.swift b/Cotabby/Services/Utilities/ModelDownloadManager.swift index 9ced7a2..742450d 100644 --- a/Cotabby/Services/Utilities/ModelDownloadManager.swift +++ b/Cotabby/Services/Utilities/ModelDownloadManager.swift @@ -311,32 +311,61 @@ final class ModelDownloadManager: ObservableObject { } } let downloadResult = try await delegate.download(from: url) - try Task.checkCancellation() - try validate(response: downloadResult.response) - let fileManager = FileManager.default + let temporaryURL = downloadResult.temporaryURL + + // The rescued temp file is ours now. Any failure before it is moved into staging (a user + // cancel, a non-2xx response — exactly when HuggingFace returns a small HTML error page, or + // the move itself failing) must remove it, or it leaks in the temp directory. + do { + try Task.checkCancellation() + try validate(response: downloadResult.response) + } catch { + try? fileManager.removeItem(at: temporaryURL) + throw error + } + let stagingURL = runtimeDirectoryURL.appendingPathComponent( "\(model.filename).staging-\(UUID().uuidString)", isDirectory: false ) - try fileManager.moveItem(at: downloadResult.temporaryURL, to: stagingURL) + do { + try fileManager.moveItem(at: temporaryURL, to: stagingURL) + } catch { + try? fileManager.removeItem(at: temporaryURL) + throw error + } + // From here the staged file must be removed on any validation OR promotion failure. do { + try ModelFileValidator.validateCompleteness( + of: stagingURL, declaredContentLength: downloadResult.response.expectedContentLength + ) try ModelFileValidator.validateSize( of: stagingURL, expectedBytes: model.expectedSizeBytes ) try ModelFileValidator.validateSHA256( of: stagingURL, expectedSHA256: model.sha256 ) + try Self.promoteStagedFile(at: stagingURL, to: destinationURL, fileManager: fileManager) } catch { try? fileManager.removeItem(at: stagingURL) throw error } + } + /// Promotes a validated staged file into the install location. When a model already exists there, + /// an atomic replace is used so a crash or error mid-promotion can never destroy the existing good + /// model before the replacement is committed (the old delete-then-move could leave nothing + /// installed). `replaceItemAt` removes the staged file as part of the swap. + private static func promoteStagedFile( + at stagingURL: URL, to destinationURL: URL, fileManager: FileManager + ) throws { if fileManager.fileExists(atPath: destinationURL.path) { - try fileManager.removeItem(at: destinationURL) + _ = try fileManager.replaceItemAt(destinationURL, withItemAt: stagingURL) + } else { + try fileManager.moveItem(at: stagingURL, to: destinationURL) } - try fileManager.moveItem(at: stagingURL, to: destinationURL) } private func validate(response: URLResponse) throws { diff --git a/Cotabby/Services/Utilities/ModelFileValidator.swift b/Cotabby/Services/Utilities/ModelFileValidator.swift index 2b4c266..a4e5732 100644 --- a/Cotabby/Services/Utilities/ModelFileValidator.swift +++ b/Cotabby/Services/Utilities/ModelFileValidator.swift @@ -66,6 +66,25 @@ enum ModelFileValidator { } } + /// Throws if the file's byte size differs from the server's declared `Content-Length`. + /// + /// The curated catalog ships with nil `expectedSizeBytes`/`sha256`, so those validators no-op and + /// the only remaining integrity gate is the HTTP status check — which does NOT catch a body the + /// server truncated while ending the transfer "cleanly" (an HTTP/2 stream reset or a proxy closing + /// the connection both surface as a finished task with no error). Comparing the on-disk size to the + /// declared length closes that gap without needing catalog metadata. No-op when the length is + /// unknown (`expectedContentLength <= 0`, i.e. `NSURLSessionTransferSizeUnknown`). + static func validateCompleteness( + of url: URL, + declaredContentLength: Int64, + fileManager: FileManager = .default + ) throws { + guard declaredContentLength > 0 else { + return + } + try validateSize(of: url, expectedBytes: declaredContentLength, fileManager: fileManager) + } + /// Throws if the file's SHA-256 differs from `expectedSHA256` (case-insensitive). /// No-op when `expectedSHA256` is nil. /// diff --git a/CotabbyTests/ModelFileValidatorTests.swift b/CotabbyTests/ModelFileValidatorTests.swift index e288805..d8a097c 100644 --- a/CotabbyTests/ModelFileValidatorTests.swift +++ b/CotabbyTests/ModelFileValidatorTests.swift @@ -59,6 +59,29 @@ final class ModelFileValidatorTests: XCTestCase { } } + // MARK: - validateCompleteness + + func test_validateCompleteness_passesWhenSizeMatchesDeclaredLength() throws { + let url = try makeFixture(contents: Data(repeating: 0xAB, count: 100)) + XCTAssertNoThrow(try ModelFileValidator.validateCompleteness(of: url, declaredContentLength: 100)) + } + + func test_validateCompleteness_throwsOnTruncatedBody() throws { + let url = try makeFixture(contents: Data(repeating: 0xAB, count: 60)) + XCTAssertThrowsError(try ModelFileValidator.validateCompleteness(of: url, declaredContentLength: 100)) { error in + guard case ModelFileValidator.ValidationError.sizeMismatch = error else { + XCTFail("Expected sizeMismatch, got \(error)") + return + } + } + } + + func test_validateCompleteness_noOpsWhenLengthUnknown() throws { + let url = try makeFixture(contents: Data(repeating: 0xAB, count: 60)) + XCTAssertNoThrow(try ModelFileValidator.validateCompleteness(of: url, declaredContentLength: -1)) + XCTAssertNoThrow(try ModelFileValidator.validateCompleteness(of: url, declaredContentLength: 0)) + } + // MARK: - validateSHA256 func test_validateSHA256_passesForKnownChecksum() throws {