DataConnectAuth.kt: just get the authUid from the sub claim, not the sub AND user_id claims

Author: dconeybe

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

@@ -50,11 +50,10 @@ internal class DataConnectAuth(
 
   override suspend fun getToken(provider: InternalAuthProvider, forceRefresh: Boolean) =
     provider.getAccessToken(forceRefresh).await().let {
-      GetAuthTokenResult(it.token, it.getAuthUids())
+      GetAuthTokenResult(it.token, it.getAuthUid())
     }
 
-  data class GetAuthTokenResult(override val token: String?, val authUids: Set<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) {
@@ -64,16 +63,6 @@ internal class DataConnectAuth(
 
   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)
-          }
-        }
-      }
-    }
+    fun com.google.firebase.auth.GetTokenResult.getAuthUid(): String? = claims["sub"] as? String
   }
 }

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

@@ -32,7 +32,6 @@ 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
@@ -47,10 +46,8 @@ 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
@@ -317,21 +314,7 @@ class DataConnectAuthUnitTest {
   }
 
   @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 {
+  fun `getToken() should populate authUid from sub claim`() = runTest {
     val dataConnectAuth = newDataConnectAuth()
     dataConnectAuth.initialize()
     advanceUntilIdle()
@@ -341,25 +324,11 @@ class DataConnectAuthUnitTest {
 
     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)
+    result.shouldNotBeNull().authUid shouldBe uid
   }
 
   @Test
-  fun `getToken() should populate empty authUids if claims are missing`() = runTest {
+  fun `getToken() should populate null authUid if sub claim is missing`() = runTest {
     val dataConnectAuth = newDataConnectAuth()
     dataConnectAuth.initialize()
     advanceUntilIdle()
@@ -368,20 +337,20 @@ class DataConnectAuthUnitTest {
 
     val result = dataConnectAuth.getToken(requestId)
 
-    result.shouldNotBeNull().authUids.shouldBeEmpty()
+    result.shouldNotBeNull().authUid.shouldBeNull()
   }
 
   @Test
-  fun `getToken() should ignore non-string uid claims`() = runTest {
+  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("user_id" to 123, "sub" to true))
+      taskForToken(accessToken, mapOf("sub" to 42))
 
     val result = dataConnectAuth.getToken(requestId)
 
-    result.shouldNotBeNull().authUids shouldBe emptySet()
+    result.shouldNotBeNull().authUid.shouldBeNull()
   }
 
   @Test

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

@@ -51,7 +51,6 @@ 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
@@ -335,8 +334,9 @@ internal inline fun <Data, reified Variables> DataConnectArb.operationRefConstru
 
 internal fun DataConnectArb.authTokenResult(
   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)
+  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()