[FAL] Add limited-use token support (#15099)

Author: daymxn

FirebaseAI/CHANGELOG.md

@@ -1,6 +1,13 @@
 # 12.2.0
 - [feature] Added support for returning thought summaries, which are synthesized
   versions of a model's internal reasoning process. (#15096)
+- [feature] Added a new configuration option to use 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. (#15099)
 
 # 12.0.0
 - [feature] Added support for Grounding with Google Search. (#15014)

FirebaseAI/Sources/FirebaseAI.swift

@@ -32,13 +32,33 @@ 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).
+  ///   - 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.
   /// - Returns: A `FirebaseAI` instance, configured with the custom `FirebaseApp`.
   public static func firebaseAI(app: FirebaseApp? = nil,
-                                backend: Backend = .googleAI()) -> FirebaseAI {
+                                backend: Backend = .googleAI(),
+                                useLimitedUseAppCheckTokens: Bool = false) -> FirebaseAI {
     let instance = createInstance(
       app: app,
       location: backend.location,
-      apiConfig: backend.apiConfig
+      apiConfig: backend.apiConfig,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     // 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 +110,8 @@ public final class FirebaseAI: Sendable {
       tools: tools,
       toolConfig: toolConfig,
       systemInstruction: systemInstruction,
-      requestOptions: requestOptions
+      requestOptions: requestOptions,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
   }
 
@@ -126,7 +147,8 @@ public final class FirebaseAI: Sendable {
       apiConfig: apiConfig,
       generationConfig: generationConfig,
       safetySettings: safetySettings,
-      requestOptions: requestOptions
+      requestOptions: requestOptions,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
   }
 
@@ -141,6 +163,8 @@ public final class FirebaseAI: Sendable {
 
   let apiConfig: APIConfig
 
+  let useLimitedUseAppCheckTokens: Bool
+
   /// 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,8 @@ public final class FirebaseAI: Sendable {
   )
 
   static func createInstance(app: FirebaseApp?, location: String?,
-                             apiConfig: APIConfig) -> FirebaseAI {
+                             apiConfig: APIConfig,
+                             useLimitedUseAppCheckTokens: Bool) -> FirebaseAI {
     guard let app = app ?? FirebaseApp.app() else {
       fatalError("No instance of the default Firebase app was found.")
     }
@@ -166,16 +191,27 @@ 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,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+    )
     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,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+    )
     instances[instanceKey] = newInstance
     return newInstance
   }
 
-  init(app: FirebaseApp, location: String?, apiConfig: APIConfig) {
+  init(app: FirebaseApp, location: String?, apiConfig: APIConfig,
+       useLimitedUseAppCheckTokens: Bool) {
     guard let projectID = app.options.projectID else {
       fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.")
     }
@@ -195,6 +231,7 @@ public final class FirebaseAI: Sendable {
     )
     self.apiConfig = apiConfig
     self.location = location
+    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
   }
 
   func modelResourceName(modelName: String) -> String {
@@ -249,5 +286,6 @@ public final class FirebaseAI: Sendable {
     let appName: String
     let location: String?
     let apiConfig: APIConfig
+    let useLimitedUseAppCheckTokens: Bool
   }
 }

FirebaseAI/Sources/GenerativeAIService.swift

@@ -30,9 +30,12 @@ struct GenerativeAIService {
 
   private let urlSession: URLSession
 
-  init(firebaseInfo: FirebaseInfo, urlSession: URLSession) {
+  private let useLimitedUseAppCheckTokens: Bool
+
+  init(firebaseInfo: FirebaseInfo, urlSession: URLSession, useLimitedUseAppCheckTokens: Bool) {
     self.firebaseInfo = firebaseInfo
     self.urlSession = urlSession
+    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
   }
 
   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,53 @@ struct GenerativeAIService {
     return urlRequest
   }
 
+  private func fetchAppCheckToken(appCheck: AppCheckInterop) async throws
+    -> FIRAppCheckTokenResultInterop {
+    if useLimitedUseAppCheckTokens {
+      if let token = await getLimitedUseAppCheckToken(appCheck: appCheck) {
+        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 getLimitedUseAppCheckToken(appCheck: AppCheckInterop) async
+    -> FIRAppCheckTokenResultInterop? {
+    // At the moment, `await` doesn’t get along with Objective-C’s optional protocol methods.
+    await withCheckedContinuation { (continuation: CheckedContinuation<
+      FIRAppCheckTokenResultInterop?,
+      Never
+    >) in
+      guard
+        useLimitedUseAppCheckTokens,
+        // `getLimitedUseToken(completion:)` is an optional protocol method. Optional binding
+        // is performed to make sure `continuation` is called even if the method’s not implemented.
+        let limitedUseTokenClosure = appCheck.getLimitedUseToken
+      else {
+        return continuation.resume(returning: nil)
+      }
+
+      limitedUseTokenClosure { tokenResult in
+        // The placeholder token should be used in the case of App Check error.
+        continuation.resume(returning: tokenResult)
+      }
+    }
+  }
+
   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,8 @@ 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`.
+  ///   - useLimitedUseAppCheckTokens: Use App Check's limited-use tokens instead of the standard
+  /// cached tokens.
   init(modelName: String,
        modelResourceName: String,
        firebaseInfo: FirebaseInfo,
@@ -86,13 +88,15 @@ public final class GenerativeModel: Sendable {
        toolConfig: ToolConfig? = nil,
        systemInstruction: ModelContent? = nil,
        requestOptions: RequestOptions,
-       urlSession: URLSession = GenAIURLSession.default) {
+       urlSession: URLSession = GenAIURLSession.default,
+       useLimitedUseAppCheckTokens: Bool = false) {
     self.modelName = modelName
     self.modelResourceName = modelResourceName
     self.apiConfig = apiConfig
     generativeAIService = GenerativeAIService(
       firebaseInfo: firebaseInfo,
-      urlSession: urlSession
+      urlSession: urlSession,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings

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

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

FirebaseAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift

@@ -121,7 +121,8 @@ extension FirebaseAI {
       return FirebaseAI.createInstance(
         app: instanceConfig.app,
         location: location,
-        apiConfig: instanceConfig.apiConfig
+        apiConfig: instanceConfig.apiConfig,
+        useLimitedUseAppCheckTokens: false
       )
     case .googleAI:
       assert(
@@ -131,7 +132,8 @@ extension FirebaseAI {
       return FirebaseAI.createInstance(
         app: instanceConfig.app,
         location: nil,
-        apiConfig: instanceConfig.apiConfig
+        apiConfig: instanceConfig.apiConfig,
+        useLimitedUseAppCheckTokens: false
       )
     }
   }

FirebaseAI/Tests/Unit/Fakes/AppCheckInteropFake.swift

@@ -40,6 +40,10 @@ class AppCheckInteropFake: NSObject, AppCheckInterop {
     return AppCheckTokenResultInteropFake(token: token, error: error)
   }
 
+  func getLimitedUseToken() async -> any FIRAppCheckTokenResultInterop {
+    return AppCheckTokenResultInteropFake(token: "limited_use_\(token)", error: error)
+  }
+
   func tokenDidChangeNotificationName() -> String {
     fatalError("\(#function) not implemented.")
   }
@@ -52,9 +56,10 @@ class AppCheckInteropFake: NSObject, AppCheckInterop {
     fatalError("\(#function) not implemented.")
   }
 
-  private class AppCheckTokenResultInteropFake: NSObject, FIRAppCheckTokenResultInterop {
-    var token: String
-    var error: Error?
+  private class AppCheckTokenResultInteropFake: NSObject, FIRAppCheckTokenResultInterop,
+    @unchecked Sendable {
+    let token: String
+    let error: Error?
 
     init(token: String, error: Error?) {
       self.token = token

FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift

@@ -501,6 +501,31 @@ final class GenerativeModelVertexAITests: XCTestCase {
     _ = try await model.generateContent(testPrompt)
   }
 
+  func testGenerateContent_appCheck_validToken_limitedUse() async throws {
+    let appCheckToken = "test-valid-token"
+    model = GenerativeModel(
+      modelName: testModelName,
+      modelResourceName: testModelResourceName,
+      firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(
+        appCheck: AppCheckInteropFake(token: appCheckToken)
+      ),
+      apiConfig: apiConfig,
+      tools: nil,
+      requestOptions: RequestOptions(),
+      urlSession: urlSession,
+      useLimitedUseAppCheckTokens: true
+    )
+    MockURLProtocol
+      .requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+        forResource: "unary-success-basic-reply-short",
+        withExtension: "json",
+        subdirectory: vertexSubdirectory,
+        appCheckToken: "limited_use_\(appCheckToken)"
+      )
+
+    _ = try await model.generateContent(testPrompt)
+  }
+
   func testGenerateContent_dataCollectionOff() async throws {
     let appCheckToken = "test-valid-token"
     model = GenerativeModel(

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.