From f36321cdbbedc94f30ecbe1e097c2cb5a0034579 Mon Sep 17 00:00:00 2001 From: Jacob Fu <141651335+FuJacob@users.noreply.github.com> Date: Mon, 1 Jun 2026 22:52:18 -0700 Subject: [PATCH] Harden model download integrity (pre-release audit, wave 2) Every curated catalog model ships with nil expectedSizeBytes/sha256, so those validators no-op and the only integrity gate was the HTTP 2xx check. A body the server truncated while ending the transfer cleanly (HTTP/2 reset, proxy close) was therefore promoted as a complete model and later failed to load with an opaque llama.cpp error. - Add ModelFileValidator.validateCompleteness, comparing the staged file's size to the server's declared Content-Length (no-op when unknown). This catches truncation without needing catalog metadata. - Clean up the rescued temp file on every pre-staging failure (cancel, non-2xx, staging-move) and the staged file on every validation/promotion failure; both previously leaked into the temp/runtime directory. - Promote atomically with replaceItemAt when a model already exists, so a failure mid-promotion can never destroy the existing good model before the replacement lands (the old delete-then-move could leave nothing installed). --- .../Utilities/ModelDownloadManager.swift | 41 ++++++++++++++++--- .../Utilities/ModelFileValidator.swift | 19 +++++++++ CotabbyTests/ModelFileValidatorTests.swift | 23 +++++++++++ 3 files changed, 77 insertions(+), 6 deletions(-) diff --git a/Cotabby/Services/Utilities/ModelDownloadManager.swift b/Cotabby/Services/Utilities/ModelDownloadManager.swift index 9ced7a26..742450d9 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 2b4c2663..a4e57320 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 e288805f..d8a097c7 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 {