[FAL] Add support for limited-use tokens (#7285)

Per [b/440338679](https://b.corp.google.com/issues/440338679),

This adds support for the `useLimitedUseAppCheckTokens` parameter to the
`getInstance` constructor and the `.ai` helper method.

To help assist this process, I've also update the
`FirebaseAIMultiResourceComponent` to use a hashed data class containing
all the unique data of each `getInstance`. Previously, it only tracked
instances by the location. Now, it tracks instances by the location,
backend, and the value of `useLimitedUseAppCheckTokens`. In all
practicality, this is unlikely to get much usage, but it helps avoid any
potential edge-case issues where customers expect different instances.

Documentation is present for the added changes, as well as a changelog
entry. These both match what we did on the iCore side, but feel free to
adjust as needed. Furthermore, the `api.txt` file has been updated with
the corresponding changes.

Author: daymxn

firebase-ai/CHANGELOG.md

@@ -7,6 +7,11 @@
   configured.
 * [changed] Added a `dilation` parameter to `ImagenMaskReference.generateMaskAndPadForOutpainting`
   (#7260)
+* [feature] Added a new configuration option to enable limited-use App Check tokens for attesting
+  Firebase AI Logic requests. This enhances security against replay attacks. To use this feature,
+  configure it explicitly via the new `useLimitedUseAppCheckTokens` parameter when initializing
+  `FirebaseAI`. We recommend migrating to limited-use tokens now, so your app will be ready to take
+  advantage of replay protection when it becomes available for Firebase AI Logic. 
 
 # 17.1.0
 =======

firebase-ai/api.txt

@@ -23,8 +23,10 @@ package com.google.firebase.ai {
     method public com.google.firebase.ai.GenerativeModel generativeModel(String modelName, com.google.firebase.ai.type.GenerationConfig? generationConfig = null, java.util.List<com.google.firebase.ai.type.SafetySetting>? safetySettings = null, java.util.List<com.google.firebase.ai.type.Tool>? tools = null, com.google.firebase.ai.type.ToolConfig? toolConfig = null, com.google.firebase.ai.type.Content? systemInstruction = null, com.google.firebase.ai.type.RequestOptions requestOptions = com.google.firebase.ai.type.RequestOptions());
     method public static com.google.firebase.ai.FirebaseAI getInstance();
     method public static com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.ai.type.GenerativeBackend backend);
+    method public static com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.ai.type.GenerativeBackend backend, boolean useLimitedUseAppCheckTokens);
     method public static com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.FirebaseApp app);
     method public static com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.FirebaseApp app = Firebase.app, com.google.firebase.ai.type.GenerativeBackend backend);
+    method public static com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.FirebaseApp app = Firebase.app, com.google.firebase.ai.type.GenerativeBackend backend, boolean useLimitedUseAppCheckTokens);
     method @com.google.firebase.ai.type.PublicPreviewAPI public com.google.firebase.ai.ImagenModel imagenModel(String modelName);
     method @com.google.firebase.ai.type.PublicPreviewAPI public com.google.firebase.ai.ImagenModel imagenModel(String modelName, com.google.firebase.ai.type.ImagenGenerationConfig? generationConfig = null);
     method @com.google.firebase.ai.type.PublicPreviewAPI public com.google.firebase.ai.ImagenModel imagenModel(String modelName, com.google.firebase.ai.type.ImagenGenerationConfig? generationConfig = null, com.google.firebase.ai.type.ImagenSafetySettings? safetySettings = null);
@@ -41,13 +43,16 @@ package com.google.firebase.ai {
   public static final class FirebaseAI.Companion {
     method public com.google.firebase.ai.FirebaseAI getInstance();
     method public com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.ai.type.GenerativeBackend backend);
+    method public com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.ai.type.GenerativeBackend backend, boolean useLimitedUseAppCheckTokens);
     method public com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.FirebaseApp app);
     method public com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.FirebaseApp app = Firebase.app, com.google.firebase.ai.type.GenerativeBackend backend);
+    method public com.google.firebase.ai.FirebaseAI getInstance(com.google.firebase.FirebaseApp app = Firebase.app, com.google.firebase.ai.type.GenerativeBackend backend, boolean useLimitedUseAppCheckTokens);
     property public final com.google.firebase.ai.FirebaseAI instance;
   }
 
   public final class FirebaseAIKt {
     method public static com.google.firebase.ai.FirebaseAI ai(com.google.firebase.Firebase, com.google.firebase.FirebaseApp app = Firebase.app, com.google.firebase.ai.type.GenerativeBackend backend = GenerativeBackend.googleAI());
+    method public static com.google.firebase.ai.FirebaseAI ai(com.google.firebase.Firebase, com.google.firebase.FirebaseApp app = Firebase.app, com.google.firebase.ai.type.GenerativeBackend backend = GenerativeBackend.googleAI(), boolean useLimitedUseAppCheckTokens);
     method public static com.google.firebase.ai.FirebaseAI getAi(com.google.firebase.Firebase);
   }
 

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,
@@ -218,9 +222,40 @@ internal constructor(
     public fun getInstance(
       app: FirebaseApp = Firebase.app,
       backend: GenerativeBackend
+    ): FirebaseAI {
+      return getInstance(app, backend, false)
+    }
+
+    /**
+     * 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.
+     */
+    // TODO(b/440356335): Update docs above when web page goes live in M170
+    @JvmStatic
+    @JvmOverloads
+    public fun getInstance(
+      app: FirebaseApp = Firebase.app,
+      backend: GenerativeBackend,
+      useLimitedUseAppCheckTokens: Boolean,
     ): FirebaseAI {
       val multiResourceComponent = app[FirebaseAIMultiResourceComponent::class.java]
-      return multiResourceComponent.get(backend)
+      return multiResourceComponent.get(InstanceKey(backend, useLimitedUseAppCheckTokens))
     }
 
     /** The [FirebaseAI] instance for the provided [FirebaseApp] using the Google AI Backend. */
@@ -249,3 +284,17 @@ public fun Firebase.ai(
   app: FirebaseApp = Firebase.app,
   backend: GenerativeBackend = GenerativeBackend.googleAI()
 ): FirebaseAI = FirebaseAI.getInstance(app, backend)
+
+/**
+ * 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].
+ */
+// TODO(b/440356335): Update docs above when web page goes live in M170
+public fun Firebase.ai(
+  app: FirebaseApp = Firebase.app,
+  backend: GenerativeBackend = GenerativeBackend.googleAI(),
+  useLimitedUseAppCheckTokens: Boolean
+): FirebaseAI = FirebaseAI.getInstance(app, backend, useLimitedUseAppCheckTokens)

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

@@ -37,18 +37,24 @@ 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 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)