Harden model download integrity (pre-release audit, wave 2)

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 {

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.
     ///

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 {