Add support for limited use tokens

Author: daymxn

firebase-ai/src/main/kotlin/com/google/firebase/ai/FirebaseAI.kt

@@ -46,6 +46,7 @@ internal constructor(
   @Blocking private val blockingDispatcher: CoroutineContext,
   private val appCheckProvider: Provider<InteropAppCheckTokenProvider>,
   private val internalAuthProvider: Provider<InternalAuthProvider>,
+  private val useLimitedUseAppCheckTokens: Boolean
 ) {
 
   /**
@@ -92,6 +93,7 @@ internal constructor(
       modelUri,
       firebaseApp.options.apiKey,
       firebaseApp,
+      useLimitedUseAppCheckTokens,
       generationConfig,
       safetySettings,
       tools,
@@ -152,7 +154,8 @@ internal constructor(
       requestOptions,
       appCheckProvider.get(),
       internalAuthProvider.get(),
-      backend
+      backend,
+      useLimitedUseAppCheckTokens,
     )
   }
 
@@ -194,6 +197,7 @@ internal constructor(
       modelUri,
       firebaseApp.options.apiKey,
       firebaseApp,
+      useLimitedUseAppCheckTokens,
       generationConfig,
       safetySettings,
       requestOptions,
@@ -212,15 +216,33 @@ internal constructor(
      * Returns the [FirebaseAI] instance for the provided [FirebaseApp] and [backend].
      *
      * @param backend the backend reference to make generative AI requests to.
+     * @param useLimitedUseAppCheckTokens when sending tokens to the backend, this option enables
+     * the usage of App Check's limited-use tokens instead of the standard cached tokens.
+     *
+     * A new limited-use tokens will be generated for each request; providing a smaller attack
+     * surface for malicious parties to hijack tokens. When used alongside replay protection,
+     * limited-use tokens are also _consumed_ after each request, ensuring they can't be used again.
+     *
+     * _This flag is set to `false` by default._
+     *
+     * > **Important:** Replay protection is not currently supported for the FirebaseAI backend. >
+     * While this feature is being developed, you can still migrate to using > limited-use tokens.
+     * Because limited-use tokens are backwards compatible, you can still > use them without replay
+     * protection. Due to their shorter TTL over standard App Check > tokens, they still provide a
+     * security benefit. > Migrating to limited-use tokens sooner minimizes disruption when support
+     * for replay > protection is added.
      */
     @JvmStatic
     @JvmOverloads
     public fun getInstance(
       app: FirebaseApp = Firebase.app,
-      backend: GenerativeBackend
+      backend: GenerativeBackend,
+      useLimitedUseAppCheckTokens: Boolean = false,
     ): FirebaseAI {
       val multiResourceComponent = app[FirebaseAIMultiResourceComponent::class.java]
-      return multiResourceComponent.get(backend)
+      return multiResourceComponent.get(
+        InstanceKey(backend.location, backend, useLimitedUseAppCheckTokens)
+      )
     }
 
     /** The [FirebaseAI] instance for the provided [FirebaseApp] using the Google AI Backend. */
@@ -244,8 +266,11 @@ public val Firebase.ai: FirebaseAI
  * Returns the [FirebaseAI] instance for the provided [FirebaseApp] and [backend].
  *
  * @param backend the backend reference to make generative AI requests to.
+ * @param useLimitedUseAppCheckTokens use App Check's limited-use tokens when sending requests to
+ * the backend. To learn more about what this means, see the full docs on [FirebaseAI.getInstance].
  */
 public fun Firebase.ai(
   app: FirebaseApp = Firebase.app,
-  backend: GenerativeBackend = GenerativeBackend.googleAI()
-): FirebaseAI = FirebaseAI.getInstance(app, backend)
+  backend: GenerativeBackend = GenerativeBackend.googleAI(),
+  useLimitedUseAppCheckTokens: Boolean = false
+): FirebaseAI = FirebaseAI.getInstance(app, backend, useLimitedUseAppCheckTokens)

firebase-ai/src/main/kotlin/com/google/firebase/ai/FirebaseAIMultiResourceComponent.kt

@@ -37,18 +37,25 @@ internal class FirebaseAIMultiResourceComponent(
   private val internalAuthProvider: Provider<InternalAuthProvider>,
 ) {
 
-  @GuardedBy("this") private val instances: MutableMap<String, FirebaseAI> = mutableMapOf()
+  @GuardedBy("this") private val instances: MutableMap<InstanceKey, FirebaseAI> = mutableMapOf()
 
-  fun get(backend: GenerativeBackend): FirebaseAI =
+  fun get(key: InstanceKey): FirebaseAI =
     synchronized(this) {
-      instances.getOrPut(backend.location) {
+      instances.getOrPut(key) {
         FirebaseAI(
           app,
-          backend,
+          key.backend,
           blockingDispatcher,
           appCheckProvider,
           internalAuthProvider,
+          key.useLimitedUseAppCheckTokens
         )
       }
     }
 }
+
+internal data class InstanceKey(
+  val location: String,
+  val backend: GenerativeBackend,
+  val useLimitedUseAppCheckTokens: Boolean
+)

firebase-ai/src/main/kotlin/com/google/firebase/ai/GenerativeModel.kt

@@ -65,6 +65,7 @@ internal constructor(
     modelName: String,
     apiKey: String,
     firebaseApp: FirebaseApp,
+    useLimitedUseAppCheckTokens: Boolean,
     generationConfig: GenerationConfig? = null,
     safetySettings: List<SafetySetting>? = null,
     tools: List<Tool>? = null,
@@ -73,7 +74,7 @@ internal constructor(
     requestOptions: RequestOptions = RequestOptions(),
     generativeBackend: GenerativeBackend,
     appCheckTokenProvider: InteropAppCheckTokenProvider? = null,
-    internalAuthProvider: InternalAuthProvider? = null,
+    internalAuthProvider: InternalAuthProvider? = null
   ) : this(
     modelName,
     generationConfig,
@@ -88,7 +89,12 @@ internal constructor(
       requestOptions,
       "gl-kotlin/${KotlinVersion.CURRENT}-ai fire/${BuildConfig.VERSION_NAME}",
       firebaseApp,
-      AppCheckHeaderProvider(TAG, appCheckTokenProvider, internalAuthProvider),
+      AppCheckHeaderProvider(
+        TAG,
+        useLimitedUseAppCheckTokens,
+        appCheckTokenProvider,
+        internalAuthProvider
+      ),
     ),
   )
 

firebase-ai/src/main/kotlin/com/google/firebase/ai/ImagenModel.kt

@@ -58,6 +58,7 @@ internal constructor(
     modelName: String,
     apiKey: String,
     firebaseApp: FirebaseApp,
+    useLimitedUseAppCheckTokens: Boolean,
     generationConfig: ImagenGenerationConfig? = null,
     safetySettings: ImagenSafetySettings? = null,
     requestOptions: RequestOptions = RequestOptions(),
@@ -73,7 +74,12 @@ internal constructor(
       requestOptions,
       "gl-kotlin/${KotlinVersion.CURRENT}-ai fire/${BuildConfig.VERSION_NAME}",
       firebaseApp,
-      AppCheckHeaderProvider(TAG, appCheckTokenProvider, internalAuthProvider),
+      AppCheckHeaderProvider(
+        TAG,
+        useLimitedUseAppCheckTokens,
+        appCheckTokenProvider,
+        internalAuthProvider
+      ),
     ),
   )
 

firebase-ai/src/main/kotlin/com/google/firebase/ai/LiveGenerativeModel.kt

@@ -70,6 +70,7 @@ internal constructor(
     appCheckTokenProvider: InteropAppCheckTokenProvider? = null,
     internalAuthProvider: InternalAuthProvider? = null,
     generativeBackend: GenerativeBackend,
+    useLimitedUseAppCheckTokens: Boolean,
   ) : this(
     modelName,
     blockingDispatcher,
@@ -83,7 +84,12 @@ internal constructor(
       requestOptions,
       "gl-kotlin/${KotlinVersion.CURRENT}-ai fire/${BuildConfig.VERSION_NAME}",
       firebaseApp,
-      AppCheckHeaderProvider(TAG, appCheckTokenProvider, internalAuthProvider),
+      AppCheckHeaderProvider(
+        TAG,
+        useLimitedUseAppCheckTokens,
+        appCheckTokenProvider,
+        internalAuthProvider
+      ),
       generativeBackend
     ),
   )

firebase-ai/src/main/kotlin/com/google/firebase/ai/common/AppCheckHeaderProvider.kt

@@ -25,6 +25,7 @@ import kotlinx.coroutines.tasks.await
 
 internal class AppCheckHeaderProvider(
   private val logTag: String,
+  private val useLimitedUseAppCheckTokens: Boolean,
   private val appCheckTokenProvider: InteropAppCheckTokenProvider? = null,
   private val internalAuthProvider: InternalAuthProvider? = null,
 ) : HeaderProvider {
@@ -36,7 +37,14 @@ internal class AppCheckHeaderProvider(
     if (appCheckTokenProvider == null) {
       Log.w(logTag, "AppCheck not registered, skipping")
     } else {
-      val token = appCheckTokenProvider.getToken(false).await()
+      val result =
+        if (useLimitedUseAppCheckTokens) {
+          appCheckTokenProvider.limitedUseToken
+        } else {
+          appCheckTokenProvider.getToken(false)
+        }
+
+      val token = result.await()
 
       if (token.error != null) {
         Log.w(logTag, "Error obtaining AppCheck token", token.error)