dataconnect(chore): add token StateFlow to observe token changes (#8273)

Author: dconeybe

firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectAppCheck.kt

@@ -22,8 +22,8 @@ import com.google.firebase.appcheck.interop.AppCheckTokenListener
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.dataconnect.core.DataConnectAppCheck.GetAppCheckTokenResult
 import com.google.firebase.dataconnect.core.Globals.toScrubbedAccessToken
-import com.google.firebase.dataconnect.core.LoggerGlobals.debug
 import com.google.firebase.dataconnect.util.IdStringGenerator
+import java.lang.ref.WeakReference
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.tasks.await
@@ -42,7 +42,10 @@ internal class DataConnectAppCheck(
     blockingDispatcher = blockingDispatcher,
     logger = logger,
   ) {
-  private val appCheckTokenListener = AppCheckTokenListenerImpl(logger)
+
+  @Suppress("LeakingThis") private val weakThis = WeakReference(this)
+
+  private val appCheckTokenListener = AppCheckTokenListenerImpl(weakThis)
 
   @DeferredApi
   override fun addTokenListener(provider: InteropAppCheckTokenProvider) =
@@ -54,11 +57,19 @@ internal class DataConnectAppCheck(
   override suspend fun getToken(provider: InteropAppCheckTokenProvider, forceRefresh: Boolean) =
     provider.getToken(forceRefresh).await().let { GetAppCheckTokenResult(it.token) }
 
-  data class GetAppCheckTokenResult(override val token: String?) : GetTokenResult
+  override fun onClose() {
+    weakThis.clear()
+  }
+
+  data class GetAppCheckTokenResult(override val token: String?) : GetTokenResult {
+    override fun toString() = "GetAppCheckTokenResult(token=${token?.toScrubbedAccessToken()})"
+  }
 
-  private class AppCheckTokenListenerImpl(private val logger: Logger) : AppCheckTokenListener {
+  private class AppCheckTokenListenerImpl(
+    private val dataConnectAppCheckRef: WeakReference<DataConnectAppCheck>,
+  ) : AppCheckTokenListener {
     override fun onAppCheckTokenChanged(tokenResult: AppCheckTokenResult) {
-      logger.debug { "onAppCheckTokenChanged(token=${tokenResult.token.toScrubbedAccessToken()})" }
+      dataConnectAppCheckRef.get()?.onTokenChanged(tokenResult.token)
     }
   }
 }

firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectAuth.kt

@@ -21,9 +21,9 @@ import com.google.firebase.auth.internal.IdTokenListener
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.dataconnect.core.DataConnectAuth.GetAuthTokenResult
 import com.google.firebase.dataconnect.core.Globals.toScrubbedAccessToken
-import com.google.firebase.dataconnect.core.LoggerGlobals.debug
 import com.google.firebase.dataconnect.util.IdStringGenerator
 import com.google.firebase.internal.InternalTokenResult
+import java.lang.ref.WeakReference
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.tasks.await
@@ -42,7 +42,10 @@ internal class DataConnectAuth(
     blockingDispatcher = blockingDispatcher,
     logger = logger,
   ) {
-  private val idTokenListener = IdTokenListenerImpl(logger)
+
+  @Suppress("LeakingThis") private val weakThis = WeakReference(this)
+
+  private val idTokenListener = IdTokenListenerImpl(weakThis)
 
   @DeferredApi
   override fun addTokenListener(provider: InternalAuthProvider) =
@@ -56,17 +59,26 @@ internal class DataConnectAuth(
       GetAuthTokenResult(it.token, it.getAuthUid())
     }
 
+  override fun onClose() {
+    weakThis.clear()
+  }
+
   @JvmInline
   value class AuthUid(val string: String) {
     override fun toString() = "AuthUid($string)"
   }
 
   data class GetAuthTokenResult(override val token: String?, val authUid: AuthUid?) :
-    GetTokenResult
+    GetTokenResult {
+    override fun toString() =
+      "GetAuthTokenResult(authUid=$authUid, token=${token?.toScrubbedAccessToken()})"
+  }
 
-  private class IdTokenListenerImpl(private val logger: Logger) : IdTokenListener {
+  private class IdTokenListenerImpl(
+    private val dataConnectAuthRef: WeakReference<DataConnectAuth>
+  ) : IdTokenListener {
     override fun onIdTokenChanged(tokenResult: InternalTokenResult) {
-      logger.debug { "onIdTokenChanged(token=${tokenResult.token?.toScrubbedAccessToken()})" }
+      dataConnectAuthRef.get()?.onTokenChanged(tokenResult.token)
     }
   }
 

firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt

@@ -32,7 +32,6 @@ import com.google.firebase.inject.Deferred.DeferredHandler
 import com.google.firebase.inject.Provider
 import com.google.firebase.internal.api.FirebaseNoSignedInUserException
 import java.lang.ref.WeakReference
-import kotlin.coroutines.coroutineContext
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineName
@@ -41,8 +40,11 @@ import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.async
 import kotlinx.coroutines.cancel
+import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.getAndUpdate
@@ -110,6 +112,17 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any, R : GetTokenRe
   /** The current state of this object. */
   private val state = MutableStateFlow<State<T, R>>(State.New)
 
+  private val _token = MutableStateFlow<R?>(null)
+
+  /**
+   * The last token returned from [getToken]
+   *
+   * Returns null if [getToken] has never been called or never completed successfully.
+   *
+   * After [close] the value of this flow is _not_ cleared, and will remain unchanged indefinitely.
+   */
+  val token: StateFlow<R?> = _token.asStateFlow()
+
   /**
    * Adds the token listener to the given provider.
    *
@@ -130,6 +143,9 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any, R : GetTokenRe
    */
   protected abstract suspend fun getToken(provider: T, forceRefresh: Boolean): R
 
+  /** Invoked synchronously by [close]. */
+  protected abstract fun onClose()
+
   /**
    * Initializes this object.
    *
@@ -173,6 +189,7 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any, R : GetTokenRe
 
     weakThis.clear()
     coroutineScope.cancel()
+    onClose()
 
     val oldState = state.getAndUpdate { State.Closed }
     when (oldState) {
@@ -339,7 +356,7 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any, R : GetTokenRe
 
       // Ensure that any exception checking below is due to an exception that happened in the
       // coroutine that called getToken(), not from the calling coroutine being cancelled.
-      coroutineContext.ensureActive()
+      currentCoroutineContext().ensureActive()
 
       val sequencedResult = jobResult.getOrNull()
       if (sequencedResult !== null && sequencedResult.sequenceNumber < attemptSequenceNumber) {
@@ -357,6 +374,7 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any, R : GetTokenRe
           logger.debug {
             "$invocationId getToken() returns null (FirebaseAuth reports no signed-in user)"
           }
+          _token.value = null
           return null
         } else if (exception is CancellationException) {
           logger.warn(exception) {
@@ -371,17 +389,16 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any, R : GetTokenRe
       }
 
       val tokenResult: R = sequencedResult!!.ref.getOrThrow()
-      logger.debug {
-        "$invocationId getToken() returns retrieved token: " +
-          tokenResult.token?.toScrubbedAccessToken()
-      }
+      logger.debug { "$invocationId getToken() returns $tokenResult" }
+      _token.value = tokenResult
       return tokenResult
     }
   }
 
   private sealed class GetTokenRetry(message: String) : Exception(message)
   private class ForceRefresh(message: String) : GetTokenRetry(message)
   private class NewProvider(message: String) : GetTokenRetry(message)
+  private class NewToken(message: String) : GetTokenRetry(message)
 
   @DeferredApi
   private fun onProviderAvailable(newProvider: T) {
@@ -422,6 +439,37 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any, R : GetTokenRe
     }
   }
 
+  protected fun onTokenChanged(newToken: String?) {
+    if (token.value?.token == newToken) {
+      return
+    }
+
+    val invocationId = idStringGenerator.next("otc")
+    logger.debug { "$invocationId onTokenChanged(newToken=${newToken?.toScrubbedAccessToken()})" }
+
+    while (true) {
+      val currentState = state.value
+
+      val activeState =
+        when (currentState) {
+          State.New -> return
+          is State.Initialized -> break
+          is State.Idle -> break
+          is State.Active -> currentState
+          State.Closed -> return
+        }
+
+      val newState = State.Idle(activeState.provider, forceTokenRefresh = false)
+      if (state.compareAndSet(currentState, newState)) {
+        val message = "$invocationId a new token is available (j567n2577q)"
+        activeState.job.cancel(message, NewToken(message))
+        break
+      }
+    }
+
+    coroutineScope.launch(CoroutineName(invocationId)) { getToken(invocationId) }
+  }
+
   /**
    * An implementation of [DeferredHandler] to be registered with the [Deferred] given to the
    * constructor.