[FAL] Add limited-use token support

FirebaseAI/CHANGELOG.md

@@ -1,3 +1,12 @@
+# Unreleased
+- [feature] Added a new configuration option to require 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 `FirebaseAI.Config` struct 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.
+
 # 12.0.0
 - [added] Added support for Grounding with Google Search. (#15014)
 - [removed] Removed `CountTokensResponse.totalBillableCharacters` which was

FirebaseAI/Sources/AppCheckOptions.swift

@@ -0,0 +1,48 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// Configurable options for how App Check is used within a ``FirebaseAI`` instance.
+///
+/// Can be set when creating a ``FirebaseAI.Config``.
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public struct AppCheckOptions: Sendable, Hashable, Equatable {
+  /// Use `limitedUseTokens`, instead of the standard cached tokens, when sending requests
+  /// to the backend.
+  let requireLimitedUseTokens: Bool
+
+  /// Creates a new ``AppCheckOptions`` value.
+  ///
+  /// - Parameters:
+  ///   - requireLimitedUseTokens: 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.
+  public init(requireLimitedUseTokens: Bool = false) {
+    self.requireLimitedUseTokens = requireLimitedUseTokens
+  }
+}

FirebaseAI/Sources/FirebaseAI.swift

@@ -32,13 +32,18 @@ public final class FirebaseAI: Sendable {
   ///     ``FirebaseApp``.
   ///   - backend: The backend API for the Firebase AI SDK; if not specified, uses the default
   ///     ``Backend/googleAI()`` (Gemini Developer API).
+  ///   - config: Configuration options for the Firebase AI SDK that propogate across all models
+  ///     created. Uses default options when not specified, see the ``FirebaseAI.Config``
+  ///     documentation for more information.
   /// - Returns: A `FirebaseAI` instance, configured with the custom `FirebaseApp`.
   public static func firebaseAI(app: FirebaseApp? = nil,
-                                backend: Backend = .googleAI()) -> FirebaseAI {
+                                backend: Backend = .googleAI(),
+                                config: FirebaseAI.Config = .config()) -> FirebaseAI {
     let instance = createInstance(
       app: app,
       location: backend.location,
-      apiConfig: backend.apiConfig
+      apiConfig: backend.apiConfig,
+      aiConfig: config
     )
     // Verify that the `FirebaseAI` instance is always configured with the production endpoint since
     // this is the public API surface for creating an instance.
@@ -90,7 +95,8 @@ public final class FirebaseAI: Sendable {
       tools: tools,
       toolConfig: toolConfig,
       systemInstruction: systemInstruction,
-      requestOptions: requestOptions
+      requestOptions: requestOptions,
+      aiConfig: aiConfig
     )
   }
 
@@ -126,21 +132,39 @@ public final class FirebaseAI: Sendable {
       apiConfig: apiConfig,
       generationConfig: generationConfig,
       safetySettings: safetySettings,
-      requestOptions: requestOptions
+      requestOptions: requestOptions,
+      aiConfig: aiConfig
     )
   }
 
   /// Class to enable FirebaseAI to register via the Objective-C based Firebase component system
   /// to include FirebaseAI in the userAgent.
   @objc(FIRVertexAIComponent) class FirebaseVertexAIComponent: NSObject {}
 
+  /// Configuration options for ``FirebaseAI``, which persists across all models.
+  @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+  public struct Config: Sendable, Hashable, Equatable {
+    /// Options for App Check specific behavior within a ``FirebaseAI`` instance.
+    let appCheck: AppCheckOptions
+
+    /// Creates a new ``FirebaseAI.Config`` value.
+    ///
+    /// - Parameters:
+    ///   - appCheck: Optionally configure certain behavior with how App Check is used.
+    public static func config(appCheck: AppCheckOptions = AppCheckOptions()) -> Config {
+      Config(appCheck: appCheck)
+    }
+  }
+
   // MARK: - Private
 
   /// Firebase data relevant to Firebase AI.
   let firebaseInfo: FirebaseInfo
 
   let apiConfig: APIConfig
 
+  let aiConfig: FirebaseAI.Config
+
   /// A map of active `FirebaseAI` instances keyed by the `FirebaseApp` name and the `location`,
   /// in the format `appName:location`.
   private nonisolated(unsafe) static var instances: [InstanceKey: FirebaseAI] = [:]
@@ -156,7 +180,7 @@ public final class FirebaseAI: Sendable {
   )
 
   static func createInstance(app: FirebaseApp?, location: String?,
-                             apiConfig: APIConfig) -> FirebaseAI {
+                             apiConfig: APIConfig, aiConfig: FirebaseAI.Config) -> FirebaseAI {
     guard let app = app ?? FirebaseApp.app() else {
       fatalError("No instance of the default Firebase app was found.")
     }
@@ -166,16 +190,26 @@ public final class FirebaseAI: Sendable {
     // Unlock before the function returns.
     defer { os_unfair_lock_unlock(&instancesLock) }
 
-    let instanceKey = InstanceKey(appName: app.name, location: location, apiConfig: apiConfig)
+    let instanceKey = InstanceKey(
+      appName: app.name,
+      location: location,
+      apiConfig: apiConfig,
+      aiConfig: aiConfig
+    )
     if let instance = instances[instanceKey] {
       return instance
     }
-    let newInstance = FirebaseAI(app: app, location: location, apiConfig: apiConfig)
+    let newInstance = FirebaseAI(
+      app: app,
+      location: location,
+      apiConfig: apiConfig,
+      aiConfig: aiConfig
+    )
     instances[instanceKey] = newInstance
     return newInstance
   }
 
-  init(app: FirebaseApp, location: String?, apiConfig: APIConfig) {
+  init(app: FirebaseApp, location: String?, apiConfig: APIConfig, aiConfig: FirebaseAI.Config) {
     guard let projectID = app.options.projectID else {
       fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.")
     }
@@ -195,6 +229,7 @@ public final class FirebaseAI: Sendable {
     )
     self.apiConfig = apiConfig
     self.location = location
+    self.aiConfig = aiConfig
   }
 
   func modelResourceName(modelName: String) -> String {
@@ -249,5 +284,6 @@ public final class FirebaseAI: Sendable {
     let appName: String
     let location: String?
     let apiConfig: APIConfig
+    let aiConfig: FirebaseAI.Config
   }
 }

FirebaseAI/Sources/GenerativeAIService.swift

@@ -30,9 +30,12 @@ struct GenerativeAIService {
 
   private let urlSession: URLSession
 
-  init(firebaseInfo: FirebaseInfo, urlSession: URLSession) {
+  private let aiConfig: FirebaseAI.Config
+
+  init(firebaseInfo: FirebaseInfo, urlSession: URLSession, aiConfig: FirebaseAI.Config) {
     self.firebaseInfo = firebaseInfo
     self.urlSession = urlSession
+    self.aiConfig = aiConfig
   }
 
   func loadRequest<T: GenerativeAIRequest>(request: T) async throws -> T.Response {
@@ -177,7 +180,7 @@ struct GenerativeAIService {
     urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type")
 
     if let appCheck = firebaseInfo.appCheck {
-      let tokenResult = await appCheck.getToken(forcingRefresh: false)
+      let tokenResult = try await fetchAppCheckToken(appCheck: appCheck)
       urlRequest.setValue(tokenResult.token, forHTTPHeaderField: "X-Firebase-AppCheck")
       if let error = tokenResult.error {
         AILog.error(
@@ -207,6 +210,30 @@ struct GenerativeAIService {
     return urlRequest
   }
 
+  private func fetchAppCheckToken(appCheck: AppCheckInterop) async throws
+    -> FIRAppCheckTokenResultInterop {
+    if aiConfig.appCheck.requireLimitedUseTokens {
+      if let token = await appCheck.getLimitedUseToken?() {
+        return token
+      }
+
+      let errorMessage =
+        "The provided App Check token provider doesn't implement getLimitedUseToken(), but requireLimitedUseTokens was enabled."
+
+      #if Debug
+        fatalError(errorMessage)
+      #else
+        throw NSError(
+          domain: "\(Constants.baseErrorDomain).\(Self.self)",
+          code: AILog.MessageCode.appCheckTokenFetchFailed.rawValue,
+          userInfo: [NSLocalizedDescriptionKey: errorMessage]
+        )
+      #endif
+    }
+
+    return await appCheck.getToken(forcingRefresh: false)
+  }
+
   private func httpResponse(urlResponse: URLResponse) throws -> HTTPURLResponse {
     // The following condition should always be true: "Whenever you make HTTP URL load requests, any
     // response objects you get back from the URLSession, NSURLConnection, or NSURLDownload class

FirebaseAI/Sources/GenerativeModel.swift

@@ -76,6 +76,7 @@ public final class GenerativeModel: Sendable {
   ///     only text content is supported.
   ///   - requestOptions: Configuration parameters for sending requests to the backend.
   ///   - urlSession: The `URLSession` to use for requests; defaults to `URLSession.shared`.
+  ///   - aiConfig: Configuration for various behavior shared across models.
   init(modelName: String,
        modelResourceName: String,
        firebaseInfo: FirebaseInfo,
@@ -86,13 +87,15 @@ public final class GenerativeModel: Sendable {
        toolConfig: ToolConfig? = nil,
        systemInstruction: ModelContent? = nil,
        requestOptions: RequestOptions,
-       urlSession: URLSession = GenAIURLSession.default) {
+       urlSession: URLSession = GenAIURLSession.default,
+       aiConfig: FirebaseAI.Config = .config()) {
     self.modelName = modelName
     self.modelResourceName = modelResourceName
     self.apiConfig = apiConfig
     generativeAIService = GenerativeAIService(
       firebaseInfo: firebaseInfo,
-      urlSession: urlSession
+      urlSession: urlSession,
+      aiConfig: aiConfig
     )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings

FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift

@@ -53,12 +53,13 @@ public final class ImagenModel {
        generationConfig: ImagenGenerationConfig?,
        safetySettings: ImagenSafetySettings?,
        requestOptions: RequestOptions,
-       urlSession: URLSession = GenAIURLSession.default) {
+       urlSession: URLSession = GenAIURLSession.default, aiConfig: FirebaseAI.Config) {
     self.modelResourceName = modelResourceName
     self.apiConfig = apiConfig
     generativeAIService = GenerativeAIService(
       firebaseInfo: firebaseInfo,
-      urlSession: urlSession
+      urlSession: urlSession,
+      aiConfig: aiConfig
     )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings

FirebaseAI/Tests/Unit/README.md

@@ -1,3 +1,3 @@
 See the Firebase AI SDK
-[README](https://github.com/firebase/firebase-ios-sdk/tree/main/FirebaseVertexAI#unit-tests)
+[README](https://github.com/firebase/firebase-ios-sdk/tree/main/FirebaseAI#unit-tests)
 for required setup instructions.

FirebaseAI/Tests/Unit/VertexComponentTests.swift

@@ -155,12 +155,14 @@ class VertexComponentTests: XCTestCase {
     let vertex1 = FirebaseAI.createInstance(
       app: VertexComponentTests.app,
       location: location,
-      apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta)
+      apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta),
+      aiConfig: .config()
     )
     let vertex2 = FirebaseAI.createInstance(
       app: VertexComponentTests.app,
       location: location,
-      apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1)
+      apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1),
+      aiConfig: .config()
     )
 
     // Ensure they are different instances.
@@ -181,7 +183,8 @@ class VertexComponentTests: XCTestCase {
       let vertex = FirebaseAI(
         app: app1,
         location: "transitory location",
-        apiConfig: FirebaseAI.defaultVertexAIAPIConfig
+        apiConfig: FirebaseAI.defaultVertexAIAPIConfig,
+        aiConfig: .config()
       )
       weakVertex = vertex
       XCTAssertNotNil(weakVertex)
@@ -208,7 +211,12 @@ class VertexComponentTests: XCTestCase {
   func testModelResourceName_developerAPI_generativeLanguage() throws {
     let app = try XCTUnwrap(VertexComponentTests.app)
     let apiConfig = APIConfig(service: .googleAI(endpoint: .googleAIBypassProxy), version: .v1beta)
-    let vertex = FirebaseAI.createInstance(app: app, location: nil, apiConfig: apiConfig)
+    let vertex = FirebaseAI.createInstance(
+      app: app,
+      location: nil,
+      apiConfig: apiConfig,
+      aiConfig: .config()
+    )
     let model = "test-model-name"
 
     let modelResourceName = vertex.modelResourceName(modelName: model)
@@ -222,7 +230,12 @@ class VertexComponentTests: XCTestCase {
       service: .googleAI(endpoint: .firebaseProxyStaging),
       version: .v1beta
     )
-    let vertex = FirebaseAI.createInstance(app: app, location: nil, apiConfig: apiConfig)
+    let vertex = FirebaseAI.createInstance(
+      app: app,
+      location: nil,
+      apiConfig: apiConfig,
+      aiConfig: .config()
+    )
     let model = "test-model-name"
     let projectID = vertex.firebaseInfo.projectID
 
@@ -253,7 +266,12 @@ class VertexComponentTests: XCTestCase {
       service: .googleAI(endpoint: .firebaseProxyStaging),
       version: .v1beta
     )
-    let vertex = FirebaseAI.createInstance(app: app, location: nil, apiConfig: apiConfig)
+    let vertex = FirebaseAI.createInstance(
+      app: app,
+      location: nil,
+      apiConfig: apiConfig,
+      aiConfig: .config()
+    )
     let modelResourceName = vertex.modelResourceName(modelName: modelName)
     let expectedSystemInstruction = ModelContent(role: nil, parts: systemInstruction.parts)