[Storage] Fix fetcherServiceMap potential race (#13446)

FirebaseStorage/Sources/Internal/StorageFetcherService.swift

@@ -23,44 +23,43 @@ import Foundation
 /// Manage Storage's fetcherService
 @available(iOS 13, tvOS 13, macOS 10.15, macCatalyst 13, watchOS 7, *)
 actor StorageFetcherService {
+  static let shared = StorageFetcherService()
+
   private var _fetcherService: GTMSessionFetcherService?
 
-  func fetcherService(_ storage: Storage) -> GTMSessionFetcherService {
+  func service(_ storage: Storage) async -> GTMSessionFetcherService {
     if let _fetcherService {
       return _fetcherService
     }
     let app = storage.app
-    if StorageFetcherService.fetcherServiceMap[app.name] == nil {
-      StorageFetcherService.fetcherServiceMap[app.name] = [:]
-    }
-    var fetcherService = StorageFetcherService.fetcherServiceMap[app.name]?[storage.storageBucket]
-    if fetcherService == nil {
-      fetcherService = GTMSessionFetcherService()
-      fetcherService?.isRetryEnabled = true
-      fetcherService?.retryBlock = retryWhenOffline
-      fetcherService?.allowLocalhostRequest = true
-      fetcherService?.maxRetryInterval = storage.maxOperationRetryInterval
-      fetcherService?.testBlock = testBlock
+    if let fetcherService = getFromMap(appName: app.name, bucket: storage.storageBucket) {
+      return fetcherService
+    } else {
+      let fetcherService = GTMSessionFetcherService()
+      fetcherService.isRetryEnabled = true
+      fetcherService.retryBlock = retryWhenOffline
+      fetcherService.allowLocalhostRequest = true
+      fetcherService.maxRetryInterval = storage.maxOperationRetryInterval
+      fetcherService.testBlock = testBlock
       let authorizer = StorageTokenAuthorizer(
         googleAppID: app.options.googleAppID,
         callbackQueue: storage.callbackQueue,
         authProvider: storage.auth,
         appCheck: storage.appCheck
       )
-      fetcherService?.authorizer = authorizer
-      StorageFetcherService.fetcherServiceMap[app.name]?[storage.storageBucket] = fetcherService
-    }
-    if storage.usesEmulator {
-      fetcherService?.allowLocalhostRequest = true
-      fetcherService?.allowedInsecureSchemes = ["http"]
+      fetcherService.authorizer = authorizer
+      if storage.usesEmulator {
+        fetcherService.allowLocalhostRequest = true
+        fetcherService.allowedInsecureSchemes = ["http"]
+      }
+      setMap(appName: app.name, bucket: storage.storageBucket, fetcher: fetcherService)
+      return fetcherService
     }
-    _fetcherService = fetcherService
-    return fetcherService!
   }
 
   /// Update the testBlock for unit testing. Save it as a property since this may be called before
   /// fetcherService is initialized.
-  func updateTestBlock(_ block: @escaping GTMSessionFetcherTestBlock) {
+  func updateTestBlock(_ block: GTMSessionFetcherTestBlock?) {
     testBlock = block
     if let _fetcherService {
       _fetcherService.testBlock = testBlock
@@ -69,9 +68,6 @@ actor StorageFetcherService {
 
   private var testBlock: GTMSessionFetcherTestBlock?
 
-  /// Map of apps to a dictionary of buckets to GTMSessionFetcherService.
-  private static var fetcherServiceMap: [String: [String: GTMSessionFetcherService]] = [:]
-
   private var retryWhenOffline: GTMSessionFetcherRetryBlock = {
     (suggestedWillRetry: Bool,
      error: Error?,
@@ -84,4 +80,15 @@ actor StorageFetcherService {
     }
     response(shouldRetry)
   }
+
+  /// Map of apps to a dictionary of buckets to GTMSessionFetcherService.
+  private var fetcherServiceMap: [String: [String: GTMSessionFetcherService]] = [:]
+
+  private func getFromMap(appName: String, bucket: String) -> GTMSessionFetcherService? {
+    return fetcherServiceMap[appName]?[bucket]
+  }
+
+  private func setMap(appName: String, bucket: String, fetcher: GTMSessionFetcherService) {
+    fetcherServiceMap[appName, default: [:]][bucket] = fetcher
+  }
 }

FirebaseStorage/Sources/Internal/StorageInternalTask.swift

@@ -38,9 +38,7 @@ class StorageInternalTask: StorageTask {
     dispatchQueue.async { [self] in
       self.state = .queueing
       Task {
-        let fetcherService = await reference.storage.fetcherService
-          .fetcherService(reference.storage)
-
+        let fetcherService = await StorageFetcherService.shared.service(reference.storage)
         var request = request ?? self.baseRequest
         request.httpMethod = httpMethod
         request.timeoutInterval = self.reference.storage.maxOperationRetryTime

FirebaseStorage/Sources/Storage.swift

@@ -266,8 +266,6 @@ import FirebaseCore
     }
   }
 
-  let fetcherService = StorageFetcherService()
-
   let dispatchQueue: DispatchQueue
 
   init(app: FirebaseApp, bucket: String) {

FirebaseStorage/Sources/StorageDownloadTask.swift

@@ -121,7 +121,7 @@ open class StorageDownloadTask: StorageObservableTask, StorageTaskManagement {
       fetcher = GTMSessionFetcher(downloadResumeData: resumeData)
       fetcher.comment = "Resuming DownloadTask"
     } else {
-      let fetcherService = await reference.storage.fetcherService.fetcherService(reference.storage)
+      let fetcherService = await StorageFetcherService.shared.service(reference.storage)
 
       fetcher = fetcherService.fetcher(with: request)
       fetcher.comment = "Starting DownloadTask"

FirebaseStorage/Sources/StorageUploadTask.swift

@@ -59,8 +59,7 @@ import Foundation
       let bodyData = try? JSONSerialization.data(withJSONObject: dataRepresentation)
 
       Task {
-        let fetcherService = await reference.storage.fetcherService
-          .fetcherService(reference.storage)
+        let fetcherService = await StorageFetcherService.shared.service(reference.storage)
         var request = self.baseRequest
         request.httpMethod = "POST"
         request.timeoutInterval = self.reference.storage.maxUploadRetryTime

FirebaseStorage/Tests/Unit/StorageAuthorizerTests.swift

@@ -23,7 +23,6 @@ class StorageAuthorizerTests: StorageTestHelpers {
   var appCheckTokenSuccess: FIRAppCheckTokenResultFake!
   var appCheckTokenError: FIRAppCheckTokenResultFake!
   var fetcher: GTMSessionFetcher!
-  var fetcherService: GTMSessionFetcherService!
   var auth: FIRAuthInteropFake!
   var appCheck: FIRAppCheckFake!
 
@@ -52,7 +51,6 @@ class StorageAuthorizerTests: StorageTestHelpers {
 
   override func tearDown() {
     fetcher = nil
-    fetcherService = nil
     auth = nil
     appCheck = nil
     appCheckTokenSuccess = nil

FirebaseStorage/Tests/Unit/StorageDeleteTests.swift

@@ -19,29 +19,9 @@ import XCTest
 
 @available(iOS 13, tvOS 13, macOS 10.15, macCatalyst 13, watchOS 7, *)
 class StorageDeleteTests: StorageTestHelpers {
-  var fetcherService: GTMSessionFetcherService?
-  var dispatchQueue: DispatchQueue?
-
-  override func setUp() {
-    super.setUp()
-    fetcherService = GTMSessionFetcherService()
-    fetcherService?.authorizer = StorageTokenAuthorizer(
-      googleAppID: "dummyAppID",
-      authProvider: nil,
-      appCheck: nil
-    )
-    dispatchQueue = DispatchQueue(label: "Test dispatch queue")
-  }
-
-  override func tearDown() {
-    fetcherService = nil
-    super.tearDown()
-  }
-
-  func testFetcherConfiguration() {
-    let expectation = self.expectation(description: #function)
-    fetcherService!.testBlock = { (fetcher: GTMSessionFetcher!,
-                                   response: GTMSessionFetcherTestResponse) in
+  func testFetcherConfiguration() async {
+    let testBlock = { (fetcher: GTMSessionFetcher!,
+                       response: GTMSessionFetcherTestResponse) in
         XCTAssertEqual(fetcher.request?.url, self.objectURL())
         XCTAssertEqual(fetcher.request?.httpMethod, "DELETE")
         let httpResponse = HTTPURLResponse(
@@ -52,67 +32,46 @@ class StorageDeleteTests: StorageTestHelpers {
         )
         response(httpResponse, nil, nil)
     }
+    await StorageFetcherService.shared.updateTestBlock(testBlock)
     let path = objectPath()
     let ref = StorageReference(storage: storage(), path: path)
-    StorageDeleteTask.deleteTask(
-      reference: ref,
-      queue: dispatchQueue!.self
-    ) { _, error in
-      expectation.fulfill()
+    do {
+      let _ = try await ref.delete()
+    } catch {
+      // All testing is in test block.
     }
-    waitForExpectation(test: self)
   }
 
-  func testSuccessfulFetch() {
-    let expectation = self.expectation(description: #function)
-    fetcherService!.testBlock = { (fetcher: GTMSessionFetcher!,
-                                   response: GTMSessionFetcherTestResponse) in
-        XCTAssertEqual(fetcher.request?.url, self.objectURL())
-        XCTAssertEqual(fetcher.request?.httpMethod, "DELETE")
-        let httpResponse = HTTPURLResponse(
-          url: (fetcher.request?.url)!,
-          statusCode: 200,
-          httpVersion: "HTTP/1.1",
-          headerFields: nil
-        )
-        response(httpResponse, nil, nil)
-    }
+  func testSuccessfulFetch() async {
+    await StorageFetcherService.shared.updateTestBlock(successBlock())
     let path = objectPath()
     let ref = StorageReference(storage: storage(), path: path)
-    StorageDeleteTask.deleteTask(
-      reference: ref,
-      queue: dispatchQueue!.self
-    ) { _, error in
-      expectation.fulfill()
+    do {
+      let _ = try await ref.delete()
+    } catch {
+      // All testing is in test block.
     }
-    waitForExpectation(test: self)
   }
 
-  func testSuccessfulFetchWithEmulator() {
-    let expectation = self.expectation(description: #function)
+  func testSuccessfulFetchWithEmulator() async {
     let storage = self.storage()
     storage.useEmulator(withHost: "localhost", port: 8080)
-    fetcherService?.allowLocalhostRequest = true
-
-    fetcherService!
-      .testBlock = successBlock(
-        withURL: URL(string: "http://localhost:8080/v0/b/bucket/o/object")!
-      )
-
+    let testBlock = successBlock(
+      withURL: URL(string: "http://localhost:8080/v0/b/bucket/o/object")!
+    )
+    await StorageFetcherService.shared.updateTestBlock(successBlock())
     let path = objectPath()
     let ref = StorageReference(storage: storage, path: path)
-    StorageDeleteTask.deleteTask(
-      reference: ref,
-      queue: dispatchQueue!.self
-    ) { _, error in
-      expectation.fulfill()
+    do {
+      let _ = try await ref.delete()
+    } catch {
+      // All testing is in test block.
     }
-    waitForExpectation(test: self)
   }
 
   func testUnsuccessfulFetchUnauthenticated() async {
     let storage = storage()
-    await storage.fetcherService.updateTestBlock(unauthenticatedBlock())
+    await StorageFetcherService.shared.updateTestBlock(unauthenticatedBlock())
     let path = objectPath()
     let ref = StorageReference(storage: storage, path: path)
     do {
@@ -124,7 +83,7 @@ class StorageDeleteTests: StorageTestHelpers {
 
   func testUnsuccessfulFetchUnauthorized() async {
     let storage = storage()
-    await storage.fetcherService.updateTestBlock(unauthorizedBlock())
+    await StorageFetcherService.shared.updateTestBlock(unauthorizedBlock())
     let path = objectPath()
     let ref = StorageReference(storage: storage, path: path)
     do {
@@ -136,7 +95,7 @@ class StorageDeleteTests: StorageTestHelpers {
 
   func testUnsuccessfulFetchObjectDoesntExist() async {
     let storage = storage()
-    await storage.fetcherService.updateTestBlock(notFoundBlock())
+    await StorageFetcherService.shared.updateTestBlock(notFoundBlock())
     let path = objectPath()
     let ref = StorageReference(storage: storage, path: path)
     do {