AuthNotificationManager from continuation to AsyncStream (#14232)

FirebaseAuth/Sources/Swift/SystemService/AuthNotificationManager.swift

@@ -56,8 +56,7 @@
     /// Only tests should access this property.
     var immediateCallbackForTestFaking: (() -> Bool)?
 
-    /// All pending callbacks while a check is being performed.
-    private var pendingCallbacks: [(Bool) -> Void]?
+    private let condition: AuthCondition
 
     /// Initializes the instance.
     /// - Parameter application: The application.
@@ -69,56 +68,53 @@
       self.application = application
       self.appCredentialManager = appCredentialManager
       timeout = kProbingTimeout
+      condition = AuthCondition()
     }
 
-    /// Checks whether or not remote notifications are being forwarded to this class.
-    /// - Parameter callback: The block to be called either immediately or in future once a result
-    /// is available.
-    func checkNotificationForwardingInternal(withCallback callback: @escaping (Bool) -> Void) {
-      if pendingCallbacks != nil {
-        pendingCallbacks?.append(callback)
-        return
+    private actor PendingCount {
+      private var count = 0
+      func increment() -> Int {
+        count = count + 1
+        return count
       }
+    }
+
+    private let pendingCount = PendingCount()
+
+    /// Checks whether or not remote notifications are being forwarded to this class.
+    func checkNotificationForwarding() async -> Bool {
       if let getValueFunc = immediateCallbackForTestFaking {
-        callback(getValueFunc())
-        return
+        return getValueFunc()
       }
       if hasCheckedNotificationForwarding {
-        callback(isNotificationBeingForwarded)
-        return
+        return isNotificationBeingForwarded
       }
-      hasCheckedNotificationForwarding = true
-      pendingCallbacks = [callback]
-
-      DispatchQueue.main.async {
-        let proberNotification = [self.kNotificationDataKey: [self.kNotificationProberKey:
-            "This fake notification should be forwarded to Firebase Auth."]]
-        if let delegate = self.application.delegate,
-           delegate
-           .responds(to: #selector(UIApplicationDelegate
-               .application(_:didReceiveRemoteNotification:fetchCompletionHandler:))) {
-          delegate.application?(self.application,
-                                didReceiveRemoteNotification: proberNotification) { _ in
+      if await pendingCount.increment() == 1 {
+        DispatchQueue.main.async {
+          let proberNotification = [self.kNotificationDataKey: [self.kNotificationProberKey:
+              "This fake notification should be forwarded to Firebase Auth."]]
+          if let delegate = self.application.delegate,
+             delegate
+             .responds(to: #selector(UIApplicationDelegate
+                 .application(_:didReceiveRemoteNotification:fetchCompletionHandler:))) {
+            delegate.application?(self.application,
+                                  didReceiveRemoteNotification: proberNotification) { _ in
+            }
+          } else {
+            AuthLog.logWarning(
+              code: "I-AUT000015",
+              message: "The UIApplicationDelegate must handle " +
+                "remote notification for phone number authentication to work."
+            )
+          }
+          kAuthGlobalWorkQueue.asyncAfter(deadline: .now() + .seconds(Int(self.timeout))) {
+            self.condition.signal()
           }
-        } else {
-          AuthLog.logWarning(
-            code: "I-AUT000015",
-            message: "The UIApplicationDelegate must handle " +
-              "remote notification for phone number authentication to work."
-          )
-        }
-        kAuthGlobalWorkQueue.asyncAfter(deadline: .now() + .seconds(Int(self.timeout))) {
-          self.callback()
-        }
-      }
-    }
-
-    func checkNotificationForwarding() async -> Bool {
-      return await withUnsafeContinuation { continuation in
-        checkNotificationForwardingInternal { value in
-          continuation.resume(returning: value)
         }
       }
+      await condition.wait()
+      hasCheckedNotificationForwarding = true
+      return isNotificationBeingForwarded
     }
 
     /// Attempts to handle the remote notification.
@@ -140,12 +136,12 @@
         return false
       }
       if dictionary[kNotificationProberKey] != nil {
-        if pendingCallbacks == nil {
+        if hasCheckedNotificationForwarding {
           // The prober notification probably comes from another instance, so pass it along.
           return false
         }
         isNotificationBeingForwarded = true
-        callback()
+        condition.signal()
         return true
       }
       guard let receipt = dictionary[kNotificationReceiptKey] as? String,
@@ -154,17 +150,5 @@
       }
       return appCredentialManager.canFinishVerification(withReceipt: receipt, secret: secret)
     }
-
-    // MARK: Internal methods
-
-    private func callback() {
-      guard let pendingCallbacks else {
-        return
-      }
-      self.pendingCallbacks = nil
-      for callback in pendingCallbacks {
-        callback(isNotificationBeingForwarded)
-      }
-    }
   }
 #endif

FirebaseAuth/Sources/Swift/Utilities/AuthCondition.swift

@@ -0,0 +1,40 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+/// Utility struct to make the execution of one task dependent upon a signal from another task.
+@available(iOS 13, tvOS 13, macOS 10.15, macCatalyst 13, watchOS 7, *)
+struct AuthCondition {
+  private let waiter: () async -> Void
+  private let stream: AsyncStream<Void>.Continuation
+
+  init() {
+    let (stream, continuation) = AsyncStream<Void>.makeStream()
+    waiter = {
+      for await _ in stream {}
+    }
+    self.stream = continuation
+  }
+
+  // Signal to unblock the waiter.
+  func signal() {
+    stream.finish()
+  }
+
+  /// Wait for the condition.
+  func wait() async {
+    await waiter()
+  }
+}

FirebaseAuth/Tests/Unit/AuthNotificationManagerTests.swift

@@ -33,7 +33,7 @@
     /** @property notificationManager
         @brief The notification manager to forward.
      */
-    private var notificationManager: AuthNotificationManager?
+    private var notificationManager: AuthNotificationManager!
 
     /** @var modernDelegate
         @brief The modern fake UIApplicationDelegate for testing.
@@ -75,7 +75,8 @@
     private func verify(forwarding: Bool, delegate: FakeForwardingDelegate) throws {
       delegate.forwardsNotification = forwarding
       let expectation = self.expectation(description: "callback")
-      notificationManager?.checkNotificationForwardingInternal { forwarded in
+      Task {
+        let forwarded = await notificationManager.checkNotificationForwarding()
         XCTAssertEqual(forwarded, forwarding)
         expectation.fulfill()
       }
@@ -93,12 +94,13 @@
       let delegate = try XCTUnwrap(modernDelegate)
       try verify(forwarding: false, delegate: delegate)
       modernDelegate?.notificationReceived = false
-      var calledBack = false
-      notificationManager?.checkNotificationForwardingInternal { isNotificationBeingForwarded in
+      let expectation = self.expectation(description: "callback")
+      Task {
+        let isNotificationBeingForwarded = await notificationManager.checkNotificationForwarding()
         XCTAssertFalse(isNotificationBeingForwarded)
-        calledBack = true
+        expectation.fulfill()
       }
-      XCTAssertTrue(calledBack)
+      waitForExpectations(timeout: 5)
       XCTAssertFalse(delegate.notificationReceived)
     }
 
@@ -136,7 +138,7 @@
         .canHandle(notification: ["com.google.firebase.auth": ["secret": kSecret]]))
       // Probing notification does not belong to this instance.
       XCTAssertFalse(manager
-        .canHandle(notification: ["com.google.firebase.auth": ["warning": "asdf"]]))
+        .canHandle(notification: ["com.google.firebase.auth": ["error": "asdf"]]))
     }
 
     private class FakeApplication: UIApplication {}