dataconnect: auth token internal refactor to track authUid, part 2 (#7485)

Author: dconeybe

firebase-dataconnect/CHANGELOG.md

@@ -1,7 +1,8 @@
 # Unreleased
 
 - [changed] Internal refactor for managing Auth and App Check tokens
-  ([#7184](https://github.com/firebase/firebase-android-sdk/pull/7184))
+  ([#7484](https://github.com/firebase/firebase-android-sdk/pull/7484),
+  [#7485](https://github.com/firebase/firebase-android-sdk/pull/7485))
 
 # 17.1.0
 

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

@@ -49,13 +49,22 @@ internal class DataConnectAuth(
     provider.removeIdTokenListener(idTokenListener)
 
   override suspend fun getToken(provider: InternalAuthProvider, forceRefresh: Boolean) =
-    provider.getAccessToken(forceRefresh).await().let { GetAuthTokenResult(it.token) }
+    provider.getAccessToken(forceRefresh).await().let {
+      GetAuthTokenResult(it.token, it.getAuthUid())
+    }
 
-  data class GetAuthTokenResult(override val token: String?) : GetTokenResult
+  data class GetAuthTokenResult(override val token: String?, val authUid: String?) : GetTokenResult
 
   private class IdTokenListenerImpl(private val logger: Logger) : IdTokenListener {
     override fun onIdTokenChanged(tokenResult: InternalTokenResult) {
       logger.debug { "onIdTokenChanged(token=${tokenResult.token?.toScrubbedAccessToken()})" }
     }
   }
+
+  private companion object {
+
+    // The "sub" claim is documented to be "a non-empty string and must be the uid of the user or
+    // device". See http://goo.gle/4oGjEQt for the relevant Firebase documentation.
+    fun com.google.firebase.auth.GetTokenResult.getAuthUid(): String? = claims["sub"] as? String
+  }
 }

firebase-dataconnect/src/test/kotlin/com/google/firebase/dataconnect/core/DataConnectAuthUnitTest.kt

@@ -54,7 +54,9 @@ import io.kotest.matchers.shouldBe
 import io.kotest.matchers.types.shouldBeSameInstanceAs
 import io.kotest.property.Arb
 import io.kotest.property.RandomSource
+import io.kotest.property.arbitrary.map
 import io.kotest.property.arbitrary.next
+import io.kotest.property.arbs.products.brand
 import io.mockk.coEvery
 import io.mockk.confirmVerified
 import io.mockk.every
@@ -311,6 +313,46 @@ class DataConnectAuthUnitTest {
     mockLogger.shouldNotHaveLoggedAnyMessagesContaining(accessToken)
   }
 
+  @Test
+  fun `getToken() should populate authUid from sub claim`() = runTest {
+    val dataConnectAuth = newDataConnectAuth()
+    dataConnectAuth.initialize()
+    advanceUntilIdle()
+    val uid = Arb.brand().map { it.value }.next(rs)
+    coEvery { mockInternalAuthProvider.getAccessToken(any()) } returns
+      taskForToken(accessToken, mapOf("sub" to uid))
+
+    val result = dataConnectAuth.getToken(requestId)
+
+    result.shouldNotBeNull().authUid shouldBe uid
+  }
+
+  @Test
+  fun `getToken() should populate null authUid if sub claim is missing`() = runTest {
+    val dataConnectAuth = newDataConnectAuth()
+    dataConnectAuth.initialize()
+    advanceUntilIdle()
+    coEvery { mockInternalAuthProvider.getAccessToken(any()) } returns
+      taskForToken(accessToken, emptyMap())
+
+    val result = dataConnectAuth.getToken(requestId)
+
+    result.shouldNotBeNull().authUid.shouldBeNull()
+  }
+
+  @Test
+  fun `getToken() should populate null authUid if sub claim is not a String`() = runTest {
+    val dataConnectAuth = newDataConnectAuth()
+    dataConnectAuth.initialize()
+    advanceUntilIdle()
+    coEvery { mockInternalAuthProvider.getAccessToken(any()) } returns
+      taskForToken(accessToken, mapOf("sub" to 42))
+
+    val result = dataConnectAuth.getToken(requestId)
+
+    result.shouldNotBeNull().authUid.shouldBeNull()
+  }
+
   @Test
   fun `getToken() should return re-throw the exception from the task returned from FirebaseAuth`() =
     runTest {
@@ -613,7 +655,7 @@ class DataConnectAuthUnitTest {
       interval = 100.milliseconds
     }
 
-    fun taskForToken(token: String?): Task<GetTokenResult> =
-      Tasks.forResult(mockk(relaxed = true) { every { getToken() } returns token })
+    fun taskForToken(token: String?, claims: Map<String, Any> = emptyMap()): Task<GetTokenResult> =
+      Tasks.forResult(GetTokenResult(token, claims))
   }
 }

firebase-dataconnect/src/test/kotlin/com/google/firebase/dataconnect/testutil/property/arbitrary/arbs.kt

@@ -334,7 +334,9 @@ internal inline fun <Data, reified Variables> DataConnectArb.operationRefConstru
 
 internal fun DataConnectArb.authTokenResult(
   accessToken: Arb<String?> = accessToken().orNull(nullProbability = 0.33),
-): Arb<GetAuthTokenResult> = accessToken.map { GetAuthTokenResult(it) }
+  authUid: Arb<String?> =
+    Arb.string(0..10, Codepoint.alphanumeric()).orNull(nullProbability = 0.33),
+): Arb<GetAuthTokenResult> = Arb.bind(accessToken, authUid, ::GetAuthTokenResult)
 
 internal fun DataConnectArb.appCheckTokenResult(
   accessToken: Arb<String?> = accessToken().orNull(nullProbability = 0.33),