[Auth] Use async flush method to get header value (#13853)

firebase · Oct 9, 2024 · c3f8bc9 · c3f8bc9
1 parent 1e63a55
commit c3f8bc9
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 15 deletions.
diff --git a/FirebaseAuth/Sources/Swift/Backend/AuthBackend.swift b/FirebaseAuth/Sources/Swift/Backend/AuthBackend.swift
@@ -98,6 +98,11 @@ class AuthBackend {
   class func request(withURL url: URL,
                      contentType: String,
                      requestConfiguration: AuthRequestConfiguration) async -> URLRequest {
+    // Kick off tasks for the async header values.
+    async let heartbeatsHeaderValue = requestConfiguration.heartbeatLogger?.asyncHeaderValue()
+    async let appCheckTokenHeaderValue = requestConfiguration.appCheck?
+      .getToken(forcingRefresh: true)
+
     var request = URLRequest(url: url)
     request.setValue(contentType, forHTTPHeaderField: "Content-Type")
     let additionalFrameworkMarker = requestConfiguration
@@ -106,13 +111,6 @@ class AuthBackend {
     request.setValue(clientVersion, forHTTPHeaderField: "X-Client-Version")
     request.setValue(Bundle.main.bundleIdentifier, forHTTPHeaderField: "X-Ios-Bundle-Identifier")
     request.setValue(requestConfiguration.appID, forHTTPHeaderField: "X-Firebase-GMPID")
-    if let heartbeatLogger = requestConfiguration.heartbeatLogger {
-      // The below call synchronously dispatches to a queue. To avoid blocking
-      // the shared concurrency queue, `async let` will spawn the process on
-      // a separate thread.
-      async let heartbeatsHeaderValue = heartbeatLogger.headerValue()
-      await request.setValue(heartbeatsHeaderValue, forHTTPHeaderField: "X-Firebase-Client")
-    }
     request.httpMethod = requestConfiguration.httpMethod
     let preferredLocalizations = Bundle.main.preferredLocalizations
     if preferredLocalizations.count > 0 {
@@ -122,8 +120,9 @@ class AuthBackend {
        languageCode.count > 0 {
       request.setValue(languageCode, forHTTPHeaderField: "X-Firebase-Locale")
     }
-    if let appCheck = requestConfiguration.appCheck {
-      let tokenResult = await appCheck.getToken(forcingRefresh: false)
+    // Wait for the async header values.
+    await request.setValue(heartbeatsHeaderValue, forHTTPHeaderField: "X-Firebase-Client")
+    if let tokenResult = await appCheckTokenHeaderValue {
       if let error = tokenResult.error {
         AuthLog.logWarning(code: "I-AUT000018",
                            message: "Error getting App Check token; using placeholder " +

diff --git a/FirebaseAuth/Tests/Unit/AuthBackendRPCImplementationTests.swift b/FirebaseAuth/Tests/Unit/AuthBackendRPCImplementationTests.swift
@@ -534,7 +534,7 @@ class AuthBackendRPCImplementationTests: RPCBaseTests {
 
   #if COCOAPODS || SWIFT_PACKAGE
     private class FakeHeartbeatLogger: NSObject, FIRHeartbeatLoggerProtocol {
-      func headerValue() -> String? {
+      func asyncHeaderValue() async -> String? {
         let payload = flushHeartbeatsIntoPayload()
         guard !payload.isEmpty else {
           return nil

diff --git a/FirebaseCore/Extension/FIRHeartbeatLogger.h b/FirebaseCore/Extension/FIRHeartbeatLogger.h
@@ -35,14 +35,19 @@ typedef NS_ENUM(NSInteger, FIRDailyHeartbeatCode) {
 /// Asynchronously logs a heartbeat.
 - (void)log;
 
-#ifndef FIREBASE_BUILD_CMAKE
-/// Return the headerValue for the HeartbeatLogger.
-- (NSString *_Nullable)headerValue;
-#endif  // FIREBASE_BUILD_CMAKE
-
 /// Gets the heartbeat code for today.
 - (FIRDailyHeartbeatCode)heartbeatCodeForToday;
 
+#ifndef FIREBASE_BUILD_CMAKE
+/// Returns the header value for the heartbeat logger via the given completion handler..
+- (void)asyncHeaderValueWithCompletionHandler:(void (^)(NSString *_Nullable))completionHandler
+    API_AVAILABLE(ios(13.0), macosx(10.15), macCatalyst(13.0), tvos(13.0), watchos(6.0));
+
+/// Return the header value for the heartbeat logger.
+- (NSString *_Nullable)
+    headerValue NS_SWIFT_UNAVAILABLE("Use `asyncHeaderValue() async -> String?` instead.");
+#endif  // FIREBASE_BUILD_CMAKE
+
 @end
 
 #ifndef FIREBASE_BUILD_CMAKE

diff --git a/FirebaseCore/Sources/FIRHeartbeatLogger.m b/FirebaseCore/Sources/FIRHeartbeatLogger.m
@@ -74,6 +74,13 @@ - (NSString *_Nullable)headerValue {
   return FIRHeaderValueFromHeartbeatsPayload([self flushHeartbeatsIntoPayload]);
 }
 
+- (void)asyncHeaderValueWithCompletionHandler:(void (^)(NSString *_Nullable))completionHandler
+    API_AVAILABLE(ios(13.0), macosx(10.15), macCatalyst(13.0), tvos(13.0), watchos(6.0)) {
+  [self flushHeartbeatsIntoPayloadWithCompletionHandler:^(FIRHeartbeatsPayload *payload) {
+    completionHandler(FIRHeaderValueFromHeartbeatsPayload(payload));
+  }];
+}
+
 - (FIRHeartbeatsPayload *)flushHeartbeatsIntoPayload {
   FIRHeartbeatsPayload *payload = [_heartbeatController flush];
   return payload;

diff --git a/FirebaseInstallations/Source/Tests/Unit/FIRInstallationsAPIServiceTests.m b/FirebaseInstallations/Source/Tests/Unit/FIRInstallationsAPIServiceTests.m
@@ -69,6 +69,11 @@ - (NSString *_Nullable)headerValue {
   return FIRHeaderValueFromHeartbeatsPayload([self flushHeartbeatsIntoPayload]);
 }
 
+- (void)asyncHeaderValueWithCompletionHandler:
+    (nonnull void (^)(NSString *_Nullable))completionHandler {
+  [self doesNotRecognizeSelector:_cmd];
+}
+
 @end
 
 #pragma mark - FIRInstallationsAPIService + Internal

diff --git a/FirebaseMessaging/Tests/UnitTests/FIRMessagingTokenOperationsTest.m b/FirebaseMessaging/Tests/UnitTests/FIRMessagingTokenOperationsTest.m
@@ -102,6 +102,11 @@ - (NSString *_Nullable)headerValue {
   return @"unimplemented";
 }
 
+- (void)asyncHeaderValueWithCompletionHandler:
+    (nonnull void (^)(NSString *_Nullable))completionHandler {
+  [self doesNotRecognizeSelector:_cmd];
+}
+
 @end
 
 #pragma mark - FIRMessagingTokenOperationsTest