DataConnectAuth.kt: add authUids to GetAuthTokenResult

This is a partial cherry-pick of #7485

Author: dconeybe

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

@@ -49,13 +49,31 @@ 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.getAuthUids())
+    }
 
-  data class GetAuthTokenResult(override val token: String?) : GetTokenResult
+  data class GetAuthTokenResult(override val token: String?, val authUids: Set<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 {
+
+    val authUidClaimNames = listOf("user_id", "sub")
+
+    fun com.google.firebase.auth.GetTokenResult.getAuthUids(): Set<String> = buildSet {
+      authUidClaimNames.forEach { claimName ->
+        claims[claimName]?.let { claimValue ->
+          if (claimValue is String) {
+            add(claimValue)
+          }
+        }
+      }
+    }
+  }
 }

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

@@ -32,6 +32,7 @@ import com.google.firebase.dataconnect.testutil.UnavailableDeferred
 import com.google.firebase.dataconnect.testutil.newBackgroundScopeThatAdvancesLikeForeground
 import com.google.firebase.dataconnect.testutil.newMockLogger
 import com.google.firebase.dataconnect.testutil.property.arbitrary.dataConnect
+import com.google.firebase.dataconnect.testutil.property.arbitrary.distinctPair
 import com.google.firebase.dataconnect.testutil.shouldContainWithNonAbuttingText
 import com.google.firebase.dataconnect.testutil.shouldContainWithNonAbuttingTextIgnoringCase
 import com.google.firebase.dataconnect.testutil.shouldHaveLoggedAtLeastOneMessageContaining
@@ -46,15 +47,19 @@ import io.kotest.assertions.nondeterministic.eventually
 import io.kotest.assertions.nondeterministic.eventuallyConfig
 import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.assertions.withClue
+import io.kotest.matchers.collections.shouldBeEmpty
 import io.kotest.matchers.collections.shouldContain
 import io.kotest.matchers.collections.shouldContainExactly
+import io.kotest.matchers.collections.shouldContainExactlyInAnyOrder
 import io.kotest.matchers.nulls.shouldBeNull
 import io.kotest.matchers.nulls.shouldNotBeNull
 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 +316,74 @@ class DataConnectAuthUnitTest {
     mockLogger.shouldNotHaveLoggedAnyMessagesContaining(accessToken)
   }
 
+  @Test
+  fun `getToken() should populate authUids from user_id 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("user_id" to uid))
+
+    val result = dataConnectAuth.getToken(requestId)
+
+    result.shouldNotBeNull().authUids.shouldContainExactly(uid)
+  }
+
+  @Test
+  fun `getToken() should populate authUids 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().authUids.shouldContainExactly(uid)
+  }
+
+  @Test
+  fun `getToken() should populate authUids from user_id and sub claims`() = runTest {
+    val dataConnectAuth = newDataConnectAuth()
+    dataConnectAuth.initialize()
+    advanceUntilIdle()
+    val (uid1, uid2) = Arb.brand().map { it.value }.distinctPair().next(rs)
+    coEvery { mockInternalAuthProvider.getAccessToken(any()) } returns
+      taskForToken(accessToken, mapOf("user_id" to uid1, "sub" to uid2))
+
+    val result = dataConnectAuth.getToken(requestId)
+
+    result.shouldNotBeNull().authUids.shouldContainExactlyInAnyOrder(uid1, uid2)
+  }
+
+  @Test
+  fun `getToken() should populate empty authUids if claims are missing`() = runTest {
+    val dataConnectAuth = newDataConnectAuth()
+    dataConnectAuth.initialize()
+    advanceUntilIdle()
+    coEvery { mockInternalAuthProvider.getAccessToken(any()) } returns
+      taskForToken(accessToken, emptyMap())
+
+    val result = dataConnectAuth.getToken(requestId)
+
+    result.shouldNotBeNull().authUids.shouldBeEmpty()
+  }
+
+  @Test
+  fun `getToken() should ignore non-string uid claims`() = runTest {
+    val dataConnectAuth = newDataConnectAuth()
+    dataConnectAuth.initialize()
+    advanceUntilIdle()
+    coEvery { mockInternalAuthProvider.getAccessToken(any()) } returns
+      taskForToken(accessToken, mapOf("user_id" to 123, "sub" to true))
+
+    val result = dataConnectAuth.getToken(requestId)
+
+    result.shouldNotBeNull().authUids shouldBe emptySet()
+  }
+
   @Test
   fun `getToken() should return re-throw the exception from the task returned from FirebaseAuth`() =
     runTest {
@@ -613,7 +686,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

@@ -51,6 +51,7 @@ import io.kotest.property.arbitrary.int
 import io.kotest.property.arbitrary.list
 import io.kotest.property.arbitrary.map
 import io.kotest.property.arbitrary.orNull
+import io.kotest.property.arbitrary.set
 import io.kotest.property.arbitrary.string
 import io.mockk.coEvery
 import io.mockk.mockk
@@ -333,9 +334,10 @@ internal inline fun <Data, reified Variables> DataConnectArb.operationRefConstru
 }
 
 internal fun DataConnectArb.authTokenResult(
-  accessToken: Arb<String> = accessToken()
-): Arb<GetAuthTokenResult> = accessToken.map { GetAuthTokenResult(it) }
+  accessToken: Arb<String?> = accessToken().orNull(nullProbability = 0.33),
+  authUids: Arb<Set<String>> = Arb.set(string(0..10, Codepoint.alphanumeric()), 0..10),
+): Arb<GetAuthTokenResult> = Arb.bind(accessToken, authUids, ::GetAuthTokenResult)
 
 internal fun DataConnectArb.appCheckTokenResult(
   accessToken: Arb<String> = accessToken()
-): Arb<GetAppCheckTokenResult> = accessToken.map { GetAppCheckTokenResult(it) }
+): Arb<GetAppCheckTokenResult> = accessToken.map(::GetAppCheckTokenResult)