fix(storage): delete file if key not found on download (#652)

Author: lawmicha

AmplifyPlugins/Storage/AWSS3StoragePlugin/Service/Storage/AWSS3StorageService+DownloadBehavior.swift

@@ -19,8 +19,9 @@ extension AWSS3StorageService {
         let downloadTaskCreatedHandler = AWSS3StorageService.makeDownloadTaskCreatedHandler(onEvent: onEvent)
         let expression = AWSS3TransferUtilityDownloadExpression()
         expression.progressBlock = AWSS3StorageService.makeOnDownloadProgressHandler(onEvent: onEvent)
-        let onDownloadCompletedHandler = AWSS3StorageService.makeDownloadCompletedHandler(onEvent: onEvent,
-                                                                                          serviceKey: serviceKey)
+        let onDownloadCompletedHandler = AWSS3StorageService.makeDownloadCompletedHandler(fileURL: fileURL,
+                                                                                          serviceKey: serviceKey,
+                                                                                          onEvent: onEvent)
 
         if let fileURL = fileURL {
             transferUtility.download(to: fileURL,
@@ -74,25 +75,27 @@ extension AWSS3StorageService {
     }
 
     private static func makeDownloadCompletedHandler(
-        onEvent: @escaping StorageServiceDownloadEventHandler,
-        serviceKey: String) -> AWSS3TransferUtilityDownloadCompletionHandlerBlock {
+        fileURL: URL? = nil,
+        serviceKey: String,
+        onEvent: @escaping StorageServiceDownloadEventHandler) -> AWSS3TransferUtilityDownloadCompletionHandlerBlock {
 
         let block: AWSS3TransferUtilityDownloadCompletionHandlerBlock = { task, location, data, error in
             guard let response = task.response else {
                 onEvent(StorageEvent.failed(StorageError.unknown("Missing HTTP Response")))
                 return
             }
 
-            let storageError = StorageErrorHelper.mapHttpResponseCode(statusCode: response.statusCode,
-                                                                      serviceKey: serviceKey)
-            guard storageError == nil else {
-                onEvent(StorageEvent.failed(storageError!))
+            if let storageError = StorageErrorHelper.mapHttpResponseCode(statusCode: response.statusCode,
+                                                                         serviceKey: serviceKey) {
+                deleteFileIfKeyNotFound(storageError: storageError, fileURL: fileURL)
+                onEvent(StorageEvent.failed(storageError))
                 return
             }
 
             guard error == nil else {
                 let error = error! as NSError
                 let storageError = StorageErrorHelper.mapTransferUtilityError(error)
+                deleteFileIfKeyNotFound(storageError: storageError, fileURL: fileURL)
                 onEvent(StorageEvent.failed(storageError))
                 return
             }
@@ -102,4 +105,16 @@ extension AWSS3StorageService {
 
         return block
     }
+
+    // TransferUtility saves the error response at the file location when the key cannot be found in S3
+    // This is an best-effort attempt to ensure that the file is removed
+    private static func deleteFileIfKeyNotFound(storageError: StorageError, fileURL: URL?) {
+        if case .keyNotFound = storageError, let fileURL = fileURL {
+            do {
+                try FileManager.default.removeItem(at: fileURL)
+            } catch {
+                Amplify.Logging.error(error: error)
+            }
+        }
+    }
 }

AmplifyPlugins/Storage/AWSS3StoragePluginFunctionalTests/AWSS3StoragePluginAccessLevelTests.swift

@@ -341,9 +341,9 @@ class AWSS3StoragePluginAccessLevelTests: AWSS3StoragePluginTestBase {
         let signInInvoked = expectation(description: "sign in completed")
         _ = Amplify.Auth.signIn(username: username, password: password) { event in
             switch event {
-            case .completed:
+            case .success:
                 signInInvoked.fulfill()
-            case .failed(let error):
+            case .failure(let error):
                 XCTFail("Failed to Sign in user \(error)")
             default:
                 XCTFail("Unexpected event")
@@ -357,7 +357,7 @@ class AWSS3StoragePluginAccessLevelTests: AWSS3StoragePluginTestBase {
         var resultOptional: String?
         _ = Amplify.Auth.fetchAuthSession(listener: { event in
             switch event {
-            case .completed(let authSession):
+            case .success(let authSession):
                 guard let cognitoAuthSession = authSession as? AuthCognitoIdentityProvider else {
                     XCTFail("Could not get auth session as AuthCognitoIdentityProvider")
                     return
@@ -369,7 +369,7 @@ class AWSS3StoragePluginAccessLevelTests: AWSS3StoragePluginTestBase {
                 case .failure(let error):
                     XCTFail("Failed to get auth session \(error)")
                 }
-            case .failed(let error):
+            case .failure(let error):
                 XCTFail("Failed to get auth session \(error)")
             default:
                 XCTFail("Unexpected event")
@@ -387,9 +387,9 @@ class AWSS3StoragePluginAccessLevelTests: AWSS3StoragePluginTestBase {
         let signOutCompleted = expectation(description: "sign out completed")
         Amplify.Auth.signOut { event in
             switch event {
-            case .completed:
+            case .success:
                 signOutCompleted.fulfill()
-            case .failed(let error):
+            case .failure(let error):
                 XCTFail("Could not sign out user \(error)")
             default:
                 XCTFail("Unexpected event")

AmplifyPlugins/Storage/AWSS3StoragePluginFunctionalTests/AWSS3StoragePluginNegativeTests.swift

@@ -43,6 +43,41 @@ class AWSS3StoragePluginNegativeTests: AWSS3StoragePluginTestBase {
         waitForExpectations(timeout: TestCommonConstants.networkTimeout)
     }
 
+    /// Given: Object does not exist in storage
+    /// When: Call the downloadFile API with path to local file
+    /// Then: Download fails and local file should not exist
+    func testDownloadToFileNonexistentKey() {
+        let key = UUID().uuidString
+        let expectedKey = "public/" + key
+        let filePath = NSTemporaryDirectory() + key + ".tmp"
+        let fileURL = URL(fileURLWithPath: filePath)
+        let failInvoked = expectation(description: "Failed is invoked")
+        let operation = Amplify.Storage.downloadFile(
+            key: key,
+            local: fileURL,
+            progressListener: nil) { event in
+                switch event {
+                case .success:
+                    XCTFail("Should not have completed successfully")
+                case .failure(let error):
+                    guard case let .keyNotFound(key, _, _, _) = error else {
+                        XCTFail("Should have been validation error")
+                        return
+                    }
+
+                    if FileManager.default.fileExists(atPath: fileURL.path) {
+                        XCTFail("local file should not exist")
+                    }
+
+                    XCTAssertEqual(key, expectedKey)
+                    failInvoked.fulfill()
+                }
+        }
+
+        XCTAssertNotNil(operation)
+        wait(for: [failInvoked], timeout: TestCommonConstants.networkTimeout)
+    }
+
     /// Given: A path to file that does not exist
     /// When: Upload the file
     /// Then: The operation fails with StorageError.missingLocalFile