fix(Storage) Use default URLSession configuration for non-multipart uploads #1542

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

@@ -52,7 +52,7 @@ extension AWSS3StorageService {
 
         request.setHTTPRequestHeaders(transferTask: transferTask)
 
-        let downloadTask = urlSession.downloadTask(with: request)
+        let downloadTask = backgroundUrlSession.downloadTask(with: request)
         transferTask.sessionTask = downloadTask
 
         // log task identifier?

AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Service/Storage/AWSS3StorageService+UploadBehavior.swift

@@ -66,7 +66,7 @@ extension AWSS3StorageService {
 
         request.setHTTPRequestHeaders(transferTask: transferTask)
 
-        let uploadTask = urlSession.uploadTask(with: request, fromFile: fileURL)
+        let uploadTask = foregroundUrlSession.uploadTask(with: request, fromFile: fileURL)
         transferTask.sessionTask = uploadTask
 
         // log task identifier?

AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Service/Storage/AWSS3StorageService.swift

@@ -11,6 +11,11 @@ import AWSS3
 import Amplify
 import AWSPluginsCore
 
+/// Represents a concrete implementation of the
+/// [AWSS3StorageServiceBehaviour](x-source-tag://AWSS3StorageServiceBehaviour)
+/// protocol.
+///
+/// - Tag: AWSS3StorageService
 class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
 
     // resettable values
@@ -24,9 +29,10 @@ class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
     var s3Client: S3Client!
 
     let storageConfiguration: StorageConfiguration
-    let sessionConfiguration: URLSessionConfiguration
-    var delegateQueue: OperationQueue?
-    var urlSession: URLSession
+    let backgroundSessionController: StorageServiceSessionController
+    var backgroundUrlSession: URLSession { backgroundSessionController.session }
+    let foregroundSessionController: StorageServiceSessionController
+    var foregroundUrlSession: URLSession { foregroundSessionController.session }
     let storageTransferDatabase: StorageTransferDatabase
     let fileSystem: FileSystem
 
@@ -76,7 +82,7 @@ class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
         self.init(authService: authService,
                   storageConfiguration: storageConfiguration,
                   storageTransferDatabase: storageTransferDatabase,
-                  sessionConfiguration: _sessionConfiguration,
+                  backgroundSessionConfiguration: _sessionConfiguration,
                   s3Client: s3Client,
                   preSignedURLBuilder: preSignedURLBuilder,
                   awsS3: awsS3,
@@ -87,7 +93,8 @@ class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
          storageConfiguration: StorageConfiguration = .default,
          storageTransferDatabase: StorageTransferDatabase = .default,
          fileSystem: FileSystem = .default,
-         sessionConfiguration: URLSessionConfiguration,
+         backgroundSessionConfiguration: URLSessionConfiguration,
+         foregroundSessionConfiguration: URLSessionConfiguration = .default,
          delegateQueue: OperationQueue? = nil,
          logger: Logger = storageLogger,
          s3Client: S3Client,
@@ -97,11 +104,15 @@ class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
         self.storageConfiguration = storageConfiguration
         self.storageTransferDatabase = storageTransferDatabase
         self.fileSystem = fileSystem
-        self.sessionConfiguration = sessionConfiguration
 
-        let delegate = StorageServiceSessionDelegate(identifier: storageConfiguration.sessionIdentifier, logger: logger)
-        self.delegateQueue = delegateQueue
-        self.urlSession = URLSession(configuration: sessionConfiguration, delegate: delegate, delegateQueue: delegateQueue)
+        self.backgroundSessionController = StorageServiceSessionController(identifier: storageConfiguration.sessionIdentifier,
+                                                                           configuration: backgroundSessionConfiguration,
+                                                                           logger: logger,
+                                                                           delegateQueue: delegateQueue)
+        self.foregroundSessionController = StorageServiceSessionController(identifier: storageConfiguration.sessionIdentifier,
+                                                                           configuration: foregroundSessionConfiguration,
+                                                                           logger: logger,
+                                                                           delegateQueue: delegateQueue)
 
         self.logger = logger
         self.s3Client = s3Client
@@ -111,24 +122,32 @@ class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
 
         StorageBackgroundEventsRegistry.register(identifier: identifier)
 
-        delegate.storageService = self
-
-        storageTransferDatabase.recover(urlSession: urlSession) { [weak self] result in
+        self.backgroundSessionController.delegate = self
+        self.foregroundSessionController.delegate = self
+        storageTransferDatabase.recover(urlSession: backgroundUrlSession) { [weak self] result in
             guard let self = self else { fatalError() }
-            switch result {
-            case .success(let pairs):
-                logger.info("Recovery completed: [pairs = \(pairs.count)]")
-                self.processTransferTaskPairs(pairs: pairs)
-            case .failure(let error):
-                logger.error(error: error)
-            }
+            self.didRecover(tasks: result)
+        }
+        storageTransferDatabase.recover(urlSession: foregroundUrlSession) { [weak self] result in
+            guard let self = self else { fatalError() }
+            self.didRecover(tasks: result)
         }
     }
 
     deinit {
         StorageBackgroundEventsRegistry.unregister(identifier: identifier)
     }
 
+    private func didRecover(tasks result: Result<StorageTransferTaskPairs, Error>) {
+        switch result {
+        case .success(let pairs):
+            logger.info("Recovery completed: [pairs = \(pairs.count)]")
+            self.processTransferTaskPairs(pairs: pairs)
+        case .failure(let error):
+            logger.error(error: error)
+        }
+    }
+
     func reset() {
         authService = nil
         preSignedURLBuilder = nil
@@ -139,11 +158,6 @@ class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
         multipartUploadSessions.removeAll()
     }
 
-    func resetURLSession() {
-        let delegate = StorageServiceSessionDelegate(identifier: storageConfiguration.sessionIdentifier, logger: logger)
-        self.urlSession = URLSession(configuration: sessionConfiguration, delegate: delegate, delegateQueue: delegateQueue)
-    }
-
     func attachEventHandlers(onUpload: AWSS3StorageServiceBehaviour.StorageServiceUploadEventHandler? = nil,
                              onDownload: AWSS3StorageServiceBehaviour.StorageServiceDownloadEventHandler? = nil,
                              onMultipartUpload: AWSS3StorageServiceBehaviour.StorageServiceMultiPartUploadEventHandler? = nil) {
@@ -282,3 +296,5 @@ class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
     }
 
 }
+
+extension AWSS3StorageService: StorageServiceSessionControllerDelegate {}

AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageMultipartUploadClient.swift

@@ -106,7 +106,7 @@ class DefaultStorageMultipartUploadClient: StorageMultipartUploadClient {
             request.setValue(userAgent, forHTTPHeaderField: "User-Agent")
              */
 
-            let uploadTask = serviceProxy.urlSession.uploadTask(with: request, fromFile: partialFileURL)
+            let uploadTask = serviceProxy.backgroundUrlSession.uploadTask(with: request, fromFile: partialFileURL)
             subTask.sessionTask = uploadTask
             subTask.uploadPart = multipartUpload.part(for: partNumber)
 
@@ -180,7 +180,7 @@ class DefaultStorageMultipartUploadClient: StorageMultipartUploadClient {
     func cancelUploadTasks(taskIdentifiers: [TaskIdentifier], done: @escaping () -> Void) {
         guard let service = serviceProxy else { return }
         service.unregister(taskIdentifiers: taskIdentifiers)
-        service.urlSession.getActiveTasks { tasks in
+        service.backgroundUrlSession.getActiveTasks { tasks in
             for task in tasks {
                 if taskIdentifiers.contains(task.taskIdentifier) {
                     task.cancel()

AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageServiceProxy.swift

@@ -11,7 +11,8 @@ import Amplify
 protocol StorageServiceProxy: AnyObject {
     var preSignedURLBuilder: AWSS3PreSignedURLBuilderBehavior! { get }
     var awsS3: AWSS3Behavior! { get }
-    var urlSession: URLSession { get }
+    var backgroundUrlSession: URLSession { get }
+    var foregroundUrlSession: URLSession { get }
 
     func register(task: StorageTransferTask)
     func unregister(task: StorageTransferTask)

AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageServiceSessionDelegate.swift → AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageServiceSessionController.swift

@@ -10,16 +10,45 @@ import Foundation
 import Amplify
 import AWSPluginsCore
 
-// MARK: - StorageServiceSessionDelegate -
+/// - Tag: StorageServiceSessionControllerDelegate
+protocol StorageServiceSessionControllerDelegate: AnyObject {
+    var identifier: String { get }
+    func unregister(task: StorageTransferTask)
+    func findTask(taskIdentifier: TaskIdentifier) -> StorageTransferTask?
+    func findMultipartUploadSession(uploadId: UploadID) -> StorageMultipartUploadSession?
+    func completeDownload(taskIdentifier: TaskIdentifier, sourceURL: URL)
+}
+
+/// Represents glue code between a `URLSession` and its
+/// [StorageServiceSessionControllerDelegate](x-source-tag://StorageServiceSessionControllerDelegate)
+/// which at the time of this writing is a
+/// [AWSS3StorageService](x-source-tag://AWSS3StorageService).
+///
+/// - Tag: StorageServiceSessionController
+class StorageServiceSessionController: NSObject {
+
+    weak var delegate: StorageServiceSessionControllerDelegate?
 
-class StorageServiceSessionDelegate: NSObject {
     let identifier: String
+    let configuration: URLSessionConfiguration
+    let delegateQueue: OperationQueue?
     let logger: Logger
-    weak var storageService: AWSS3StorageService?
+    var session: URLSession = .shared
 
-    init(identifier: String, logger: Logger = storageLogger) {
+    init(identifier: String,
+         configuration: URLSessionConfiguration,
+         logger: Logger = storageLogger,
+         delegateQueue: OperationQueue? = nil) {
         self.identifier = identifier
+        self.configuration = configuration
         self.logger = logger
+        self.delegateQueue = delegateQueue
+        super.init()
+        self.resetURLSession()
+    }
+
+    private func resetURLSession() {
+        self.session = URLSession(configuration: configuration, delegate: self, delegateQueue: delegateQueue)
     }
 
     // Set a Symbolic Breakpoint in Xcode to monitor these messages
@@ -32,7 +61,7 @@ class StorageServiceSessionDelegate: NSObject {
     }
 
     private func findTransferTask(for taskIdentifier: TaskIdentifier) -> StorageTransferTask? {
-        guard let storageService = storageService,
+        guard let storageService = delegate,
               let transferTask = storageService.findTask(taskIdentifier: taskIdentifier) else {
                   logger.debug("Did not find transfer task: \(taskIdentifier)")
                   return nil
@@ -48,12 +77,12 @@ public extension Notification.Name {
 
 // MARK: - URLSessionDelegate -
 
-extension StorageServiceSessionDelegate: URLSessionDelegate {
+extension StorageServiceSessionController: URLSessionDelegate {
 
     func urlSessionDidFinishEvents(forBackgroundURLSession session: URLSession) {
         logURLSessionActivity("Session did finish background events")
 
-        if let identifier = storageService?.identifier,
+        if let identifier = delegate?.identifier,
            let continuation = StorageBackgroundEventsRegistry.getContinuation(for: identifier) {
             // Must be run on main thread as covered by Apple Developer docs.
             Task { @MainActor in
@@ -74,13 +103,13 @@ extension StorageServiceSessionDelegate: URLSessionDelegate {
         NotificationCenter.default.post(name: Notification.Name.StorageURLSessionDidBecomeInvalidNotification, object: session)
 
         // Reset URLSession since the current one has become invalid.
-        storageService?.resetURLSession()
+        resetURLSession()
     }
 }
 
 // MARK: - URLSessionTaskDelegate -
 
-extension StorageServiceSessionDelegate: URLSessionTaskDelegate {
+extension StorageServiceSessionController: URLSessionTaskDelegate {
 
     func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
         if let error = error {
@@ -106,7 +135,7 @@ extension StorageServiceSessionDelegate: URLSessionTaskDelegate {
             logURLSessionActivity("Session task did complete: \(task.taskIdentifier)")
         }
 
-        guard let storageService = storageService,
+        guard let storageService = delegate,
               let transferTask = findTransferTask(for: task.taskIdentifier) else {
                   logURLSessionActivity("Session task not handled: \(task.taskIdentifier)")
                   return
@@ -155,7 +184,7 @@ extension StorageServiceSessionDelegate: URLSessionTaskDelegate {
     func urlSession(_ session: URLSession, task: URLSessionTask, didSendBodyData bytesSent: Int64, totalBytesSent: Int64, totalBytesExpectedToSend: Int64) {
         logURLSessionActivity("Session task update: [bytesSent: \(bytesSent)], [totalBytesSent: \(totalBytesSent)], [totalBytesExpectedToSend: \(totalBytesExpectedToSend)]")
 
-        guard let storageService = storageService,
+        guard let storageService = delegate,
               let transferTask = findTransferTask(for: task.taskIdentifier) else { return }
 
         switch transferTask.transferType {
@@ -180,7 +209,7 @@ extension StorageServiceSessionDelegate: URLSessionTaskDelegate {
 
 // MARK: - URLSessionDownloadDelegate -
 
-extension StorageServiceSessionDelegate: URLSessionDownloadDelegate {
+extension StorageServiceSessionController: URLSessionDownloadDelegate {
 
     func urlSession(_ session: URLSession, downloadTask: URLSessionDownloadTask, didWriteData bytesWritten: Int64, totalBytesWritten: Int64, totalBytesExpectedToWrite: Int64) {
         logURLSessionActivity("Session download task [\(downloadTask.taskIdentifier)] did write [\(bytesWritten)], [totalBytesWritten \(totalBytesWritten)], [totalBytesExpectedToWrite: \(totalBytesExpectedToWrite)]")
@@ -195,7 +224,7 @@ extension StorageServiceSessionDelegate: URLSessionDownloadDelegate {
     func urlSession(_ session: URLSession, downloadTask: URLSessionDownloadTask, didFinishDownloadingTo location: URL) {
         logURLSessionActivity("Session download task [\(downloadTask.taskIdentifier)] did finish downloading to \(location.path)")
 
-        guard let storageService = storageService,
+        guard let storageService = delegate,
               let transferTask = findTransferTask(for: downloadTask.taskIdentifier) else { return }
 
         let response = StorageTransferResponse(task: downloadTask, error: nil, transferTask: transferTask)

AmplifyPlugins/Storage/Tests/AWSS3StoragePluginTests/Mocks/MockAWSAuthServiceBehavior.swift

@@ -0,0 +1,53 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+@testable import AWSPluginsCore
+@testable import AWSS3StoragePlugin
+
+import AWSClientRuntime
+import AWSS3
+import Amplify
+import Foundation
+
+/// Test-friendly implementation of a
+/// [AWSAuthServiceBehavior](x-source-tag://AWSAuthServiceBehavior) protocol.
+///
+/// - Tag: MockAWSAuthServiceBehavior
+final class MockAWSAuthServiceBehavior {
+    var interactions: [String] = []
+    var credentialsProvider = MockCredentialsProvider()
+    var tokenClaimsByToken: [String: [String : AnyObject]] = [:]
+    var identityID = UUID().uuidString
+    var userPoolAccessToken = UUID().uuidString
+}
+
+extension MockAWSAuthServiceBehavior: AWSAuthServiceBehavior {
+
+    func getCredentialsProvider() -> CredentialsProvider {
+        interactions.append(#function)
+        return credentialsProvider
+    }
+
+    func getTokenClaims(tokenString: String) -> Result<[String : AnyObject], AuthError> {
+        interactions.append(#function)
+        if let claims = tokenClaimsByToken[tokenString] {
+            return .success(claims)
+        }
+        return .failure(.unknown(tokenString))
+    }
+
+    func getIdentityID() async throws -> String {
+        interactions.append(#function)
+        return identityID
+    }
+
+    func getUserPoolAccessToken() async throws -> String {
+        interactions.append(#function)
+        return userPoolAccessToken
+    }
+
+}

AmplifyPlugins/Storage/Tests/AWSS3StoragePluginTests/Mocks/MockAWSS3PreSignedURLBuilder.swift

@@ -5,15 +5,28 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-/*
-import Foundation
 @testable import AWSS3StoragePlugin
+
 import AWSS3
+import Foundation
 
-public class MockAWSS3PreSignedURLBuilder: AWSS3PreSignedURLBuilderBehavior {
-    public func getPreSignedURL(_ getPreSignedURLRequest: AWSS3GetPreSignedURLRequest) -> AWSTask<NSURL> {
-        return AWSTask()
-    }
+/// Test-friendly implementation of a
+/// [AWSS3PreSignedURLBuilderBehavior](x-source-tag://AWSS3PreSignedURLBuilderBehavior)
+/// protocol.
+///
+/// - Tag: MockAWSS3PreSignedURLBuilder
+final class MockAWSS3PreSignedURLBuilder {
+    var interactions: [String] = []
+    var defaultURL = URL(fileURLWithPath: NSTemporaryDirectory().appendingPathComponent(UUID().uuidString))
+    var preSignedURLs: [String: URL] = [:]
+}
 
+extension MockAWSS3PreSignedURLBuilder: AWSS3PreSignedURLBuilderBehavior {
+    func getPreSignedURL(key: String, signingOperation: AWSS3SigningOperation, expires: Int64?) async throws -> URL {
+        interactions.append(#function)
+        if let url = preSignedURLs[key] {
+            return url
+        }
+        return defaultURL
+    }
 }
-*/