dataconnect(chore): Create AuthUid value class to add some type safety compared to raw strings (#8198)

Author: dconeybe

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

@@ -56,7 +56,13 @@ internal class DataConnectAuth(
       GetAuthTokenResult(it.token, it.getAuthUid())
     }
 
-  data class GetAuthTokenResult(override val token: String?, val authUid: String?) : GetTokenResult
+  @JvmInline
+  value class AuthUid(val string: String) {
+    override fun toString() = "AuthUid($string)"
+  }
+
+  data class GetAuthTokenResult(override val token: String?, val authUid: AuthUid?) :
+    GetTokenResult
 
   private class IdTokenListenerImpl(private val logger: Logger) : IdTokenListener {
     override fun onIdTokenChanged(tokenResult: InternalTokenResult) {
@@ -68,6 +74,9 @@ internal class DataConnectAuth(
 
     // 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
+    fun com.google.firebase.auth.GetTokenResult.getAuthUid(): AuthUid? {
+      val sub = claims["sub"] as? String
+      return if (sub === null) null else AuthUid(sub)
+    }
   }
 }

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

@@ -16,6 +16,7 @@
 
 package com.google.firebase.dataconnect.core
 
+import com.google.firebase.dataconnect.core.DataConnectAuth.AuthUid
 import com.google.firebase.dataconnect.core.LoggerGlobals.debug
 import com.google.firebase.dataconnect.util.CoroutineUtils.completedFlow
 import com.google.firebase.dataconnect.util.GrpcBidiFlow
@@ -72,7 +73,7 @@ import kotlinx.coroutines.launch
  * @param coroutineScope The [CoroutineScope] to whose lifetime this object belongs.
  */
 internal class DataConnectBidiConnectStream(
-  flow: Flow<GrpcBidiFlow.Event<StreamRequestProto, StreamResponseProto, String?>>,
+  flow: Flow<GrpcBidiFlow.Event<StreamRequestProto, StreamResponseProto, AuthUid?>>,
   private val coroutineScope: CoroutineScope,
   private val logger: Logger,
 ) {
@@ -278,11 +279,11 @@ internal class DataConnectBidiConnectStream(
 
     class Connected(
       val connectionId: String,
-      val authUid: String?,
+      val authUid: AuthUid?,
       val outgoingRequests: SendChannel<StreamRequestProto>,
     ) : Connection {
       constructor(
-        event: GrpcBidiFlow.Event.ConnectionInfo<StreamRequestProto, String?>
+        event: GrpcBidiFlow.Event.ConnectionInfo<StreamRequestProto, AuthUid?>
       ) : this(event.connectionId, event.connectionCookie, event.outgoingRequests)
 
       override fun toString() = "Connected(connectionId=$connectionId)"
@@ -342,7 +343,7 @@ internal class DataConnectBidiConnectStream(
   ): Flow<SubscriptionEvent.Message> {
 
     fun SendChannel<StreamRequestProto>.trySendOrThrow(
-      authUid: String?,
+      authUid: AuthUid?,
       request: StreamRequestProto
     ) {
       val sendResult = trySend(request)
@@ -359,7 +360,7 @@ internal class DataConnectBidiConnectStream(
     }
 
     suspend fun SendChannel<StreamRequestProto>.subscribeOrResumeLoop(
-      authUid: String?,
+      authUid: AuthUid?,
       subscribeOrResumeSignal: ConflatedSignal,
     ) {
       var subscribed = false

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

@@ -20,6 +20,7 @@ import android.os.Build
 import com.google.firebase.FirebaseApp
 import com.google.firebase.dataconnect.BuildConfig
 import com.google.firebase.dataconnect.FirebaseDataConnect
+import com.google.firebase.dataconnect.core.DataConnectAuth.AuthUid
 import com.google.firebase.dataconnect.core.Globals.toScrubbedAccessToken
 import com.google.firebase.dataconnect.core.LoggerGlobals.Logger
 import com.google.firebase.dataconnect.core.LoggerGlobals.debug
@@ -100,7 +101,7 @@ internal class DataConnectGrpcMetadata(
   companion object {
     // TODO: Move this to ProtoUtil.kt where it would live alongside other related methods.
     // NOTE: Keep the implementation of this method in parity with StructProtoBuilder.putHeaders().
-    fun Metadata.toStructProto(authUid: String?): Struct = buildStructProto {
+    fun Metadata.toStructProto(authUid: AuthUid?): Struct = buildStructProto {
       val keys: List<Metadata.Key<String>> = run {
         val keySet: MutableSet<String> = keys().toMutableSet()
         // Always explicitly include the auth header in the returned string, even if it is absent.
@@ -116,7 +117,8 @@ internal class DataConnectGrpcMetadata(
           else {
             values.map {
               when (key.name()) {
-                firebaseAuthTokenHeader.name() -> it.toScrubbedAccessToken() + " (authUid=$authUid)"
+                firebaseAuthTokenHeader.name() ->
+                  it.toScrubbedAccessToken() + " (authUid=${authUid?.string})"
                 firebaseAppCheckTokenHeader.name() -> it.toScrubbedAccessToken()
                 else -> it
               }
@@ -131,7 +133,11 @@ internal class DataConnectGrpcMetadata(
 
     // TODO: Move this to ProtoUtil.kt where it would live alongside other related methods.
     // NOTE: Keep the implementation of this method in parity with Metadata.toStructProto().
-    fun StructProtoBuilder.putHeaders(authUid: String?, key: String, headers: Map<String, String>) {
+    fun StructProtoBuilder.putHeaders(
+      authUid: AuthUid?,
+      key: String,
+      headers: Map<String, String>
+    ) {
       putStruct(key) {
         val keys: List<String> =
           buildSet {
@@ -148,7 +154,8 @@ internal class DataConnectGrpcMetadata(
           val scrubbedValue =
             value?.let {
               when (key) {
-                firebaseAuthTokenHeader.name() -> it.toScrubbedAccessToken() + " (authUid=$authUid)"
+                firebaseAuthTokenHeader.name() ->
+                  it.toScrubbedAccessToken() + " (authUid=${authUid?.string})"
                 firebaseAppCheckTokenHeader.name() -> it.toScrubbedAccessToken()
                 else -> it
               }

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

@@ -24,6 +24,7 @@ import com.google.firebase.dataconnect.DataConnectPath
 import com.google.firebase.dataconnect.DataConnectPathSegment
 import com.google.firebase.dataconnect.FirebaseDataConnect
 import com.google.firebase.dataconnect.QueryRef.FetchPolicy
+import com.google.firebase.dataconnect.core.DataConnectAuth.AuthUid
 import com.google.firebase.dataconnect.core.DataConnectGrpcMetadata.Companion.toStructProto
 import com.google.firebase.dataconnect.core.LoggerGlobals.Logger
 import com.google.firebase.dataconnect.core.LoggerGlobals.debug
@@ -240,7 +241,7 @@ internal class DataConnectGrpcRPCs(
 
   private class QueryCacheInfo(
     val cacheDb: DataConnectCacheDatabase,
-    val authUid: String?,
+    val authUid: AuthUid?,
     val queryId: ImmutableByteArray,
     val maxAge: DurationProto,
   )
@@ -427,7 +428,7 @@ internal class DataConnectGrpcRPCs(
 
   private inner class ConnectGrpcBidiFlowListener(
     private val streamId: String,
-    private val authUid: String?,
+    private val authUid: AuthUid?,
     private val initRequest: StreamRequest,
     private val kotlinMethodName: String,
   ) : GrpcBidiFlow.Listener<StreamRequest, StreamResponse> {
@@ -511,7 +512,7 @@ internal class DataConnectGrpcRPCs(
   }
 
   @Suppress("unused")
-  private class ConnectGrpcBidiFlowListenerFormatter(private val authUid: String?) :
+  private class ConnectGrpcBidiFlowListenerFormatter(private val authUid: AuthUid?) :
     GrpcBidiFlowListenerMessageFormatter.Formatter<StreamRequest, StreamResponse>() {
     override fun connectionStartingHeaders(headers: Metadata?): String =
       headers?.toStructProto(authUid)?.toCompactString().toString()
@@ -682,7 +683,7 @@ internal class DataConnectGrpcRPCs(
       metadata: Metadata?,
       request: () -> Struct,
       requestTypeName: String,
-      authUid: String?,
+      authUid: AuthUid?,
     ) = debug {
       val requestStruct = request()
       val struct = buildStructProto {

firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/sqlite/DataConnectCacheDatabase.kt

@@ -18,6 +18,7 @@ package com.google.firebase.dataconnect.sqlite
 
 import android.annotation.SuppressLint
 import android.database.sqlite.SQLiteDatabase
+import com.google.firebase.dataconnect.core.DataConnectAuth.AuthUid
 import com.google.firebase.dataconnect.core.Logger
 import com.google.firebase.dataconnect.core.LoggerGlobals.warn
 import com.google.firebase.dataconnect.sqlite.DataConnectCacheDatabase.GetQueryResultResult
@@ -198,17 +199,17 @@ internal class DataConnectCacheDatabase(
 
   @JvmInline private value class SqliteEntityId(val sqliteRowId: Long)
 
-  private fun SQLiteDatabase.getOrInsertAuthUid(authUid: String?): SqliteUserId {
-    execSQL(logger, "INSERT OR IGNORE INTO users (auth_uid) VALUES (?)", arrayOf(authUid))
+  private fun SQLiteDatabase.getOrInsertAuthUid(authUid: AuthUid?): SqliteUserId {
+    execSQL(logger, "INSERT OR IGNORE INTO users (auth_uid) VALUES (?)", arrayOf(authUid?.string))
     return rawQuery(
       logger,
       "SELECT id FROM users WHERE auth_uid ${if (authUid === null) "IS NULL" else "= ?"}",
-      if (authUid === null) emptyArray() else arrayOf(authUid),
+      if (authUid === null) emptyArray() else arrayOf(authUid.string),
     ) { cursor ->
       if (cursor.moveToNext()) {
         SqliteUserId(cursor.getLong(0))
       } else {
-        throw AuthUidNotFoundException("authUid=$authUid (internal error m5m52ahrxz)")
+        throw AuthUidNotFoundException("authUid=${authUid?.string} (internal error m5m52ahrxz)")
       }
     }
   }
@@ -451,7 +452,7 @@ internal class DataConnectCacheDatabase(
   }
 
   suspend fun getQueryResult(
-    authUid: String?,
+    authUid: AuthUid?,
     queryId: ImmutableByteArray,
     currentTimeMillis: Long,
     staleResult: KClass<out GetQueryResultResult>,
@@ -524,7 +525,7 @@ internal class DataConnectCacheDatabase(
   }
 
   suspend fun insertQueryResult(
-    authUid: String?,
+    authUid: AuthUid?,
     queryId: ImmutableByteArray,
     queryData: Struct,
     maxAge: DurationProto,

firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/util/ProtoUtil.kt

@@ -18,6 +18,7 @@ package com.google.firebase.dataconnect.util
 
 import com.google.firebase.dataconnect.DataConnectPath
 import com.google.firebase.dataconnect.DataConnectPathSegment
+import com.google.firebase.dataconnect.core.DataConnectAuth.AuthUid
 import com.google.firebase.dataconnect.core.DataConnectGrpcMetadata.Companion.putHeaders
 import com.google.firebase.dataconnect.core.DataConnectSerialization.Companion.toErrorInfoImpl
 import com.google.firebase.dataconnect.toPathString
@@ -297,11 +298,11 @@ internal object ProtoUtil {
     }
 
   fun StreamRequest.toCompactString(
-    authUid: String?,
+    authUid: AuthUid?,
     keySortSelector: ((String) -> String)? = ::streamRequestKeySortSelector,
   ): String = toStructProto(authUid).toCompactString(keySortSelector)
 
-  fun StreamRequest.toStructProto(authUid: String?): Struct = buildStructProto {
+  fun StreamRequest.toStructProto(authUid: AuthUid?): Struct = buildStructProto {
     name.takeIf { it.isNotEmpty() }?.let { put("name", it) }
     requestId.takeIf { it.isNotEmpty() }?.let { put("requestId", it) }
     dataEtag.takeIf { it.isNotEmpty() }?.let { put("data_etag", it) }

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

@@ -23,6 +23,7 @@ import com.google.firebase.auth.GetTokenResult
 import com.google.firebase.auth.internal.IdTokenListener
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.dataconnect.DataConnectException
+import com.google.firebase.dataconnect.core.DataConnectAuth.AuthUid
 import com.google.firebase.dataconnect.core.Globals.toScrubbedAccessToken
 import com.google.firebase.dataconnect.testutil.DataConnectLogLevelRule
 import com.google.firebase.dataconnect.testutil.DelayedDeferred
@@ -326,7 +327,7 @@ class DataConnectAuthUnitTest {
 
     val result = dataConnectAuth.getToken(requestId)
 
-    result.shouldNotBeNull().authUid shouldBe uid
+    result.shouldNotBeNull().authUid shouldBe AuthUid(uid)
   }
 
   @Test