diff --git a/FirebaseAI/CHANGELOG.md b/FirebaseAI/CHANGELOG.md index ea7fd44941e..b5f3199a797 100644 --- a/FirebaseAI/CHANGELOG.md +++ b/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) diff --git a/FirebaseAI/Sources/FirebaseAI.swift b/FirebaseAI/Sources/FirebaseAI.swift index 48f7183d4e6..d4c2e9d545d 100644 --- a/FirebaseAI/Sources/FirebaseAI.swift +++ b/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 } } diff --git a/FirebaseAI/Sources/GenerativeAIService.swift b/FirebaseAI/Sources/GenerativeAIService.swift index e1538af997f..5b676162c94 100644 --- a/FirebaseAI/Sources/GenerativeAIService.swift +++ b/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(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 diff --git a/FirebaseAI/Sources/GenerativeModel.swift b/FirebaseAI/Sources/GenerativeModel.swift index 8d3f5e043a7..2807022de59 100644 --- a/FirebaseAI/Sources/GenerativeModel.swift +++ b/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 diff --git a/FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift b/FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift index 729fad9f28d..542f842362c 100644 --- a/FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift +++ b/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 diff --git a/FirebaseAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift b/FirebaseAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift index 21554d28250..2c6dceb24fd 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift +++ b/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 ) } } diff --git a/FirebaseAI/Tests/Unit/Fakes/AppCheckInteropFake.swift b/FirebaseAI/Tests/Unit/Fakes/AppCheckInteropFake.swift index 62fc753ae68..81b6ac1bb7c 100644 --- a/FirebaseAI/Tests/Unit/Fakes/AppCheckInteropFake.swift +++ b/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 diff --git a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift index 7c23726f152..80a383eb461 100644 --- a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift +++ b/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( diff --git a/FirebaseAI/Tests/Unit/README.md b/FirebaseAI/Tests/Unit/README.md index 9463d595294..88019041f9f 100644 --- a/FirebaseAI/Tests/Unit/README.md +++ b/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. diff --git a/FirebaseAI/Tests/Unit/VertexComponentTests.swift b/FirebaseAI/Tests/Unit/VertexComponentTests.swift index 7202e01f4d6..702c6e50871 100644 --- a/FirebaseAI/Tests/Unit/VertexComponentTests.swift +++ b/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), + useLimitedUseAppCheckTokens: false ) 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), + useLimitedUseAppCheckTokens: false ) // 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, + useLimitedUseAppCheckTokens: false ) 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, + useLimitedUseAppCheckTokens: false + ) 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, + useLimitedUseAppCheckTokens: false + ) 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, + useLimitedUseAppCheckTokens: false + ) let modelResourceName = vertex.modelResourceName(modelName: modelName) let expectedSystemInstruction = ModelContent(role: nil, parts: systemInstruction.parts) diff --git a/FirebaseAppCheck/Interop/Public/FirebaseAppCheckInterop/FIRAppCheckTokenResultInterop.h b/FirebaseAppCheck/Interop/Public/FirebaseAppCheckInterop/FIRAppCheckTokenResultInterop.h index cb86a1bffe8..f87820b9790 100644 --- a/FirebaseAppCheck/Interop/Public/FirebaseAppCheckInterop/FIRAppCheckTokenResultInterop.h +++ b/FirebaseAppCheck/Interop/Public/FirebaseAppCheckInterop/FIRAppCheckTokenResultInterop.h @@ -18,6 +18,7 @@ NS_ASSUME_NONNULL_BEGIN +NS_SWIFT_SENDABLE @protocol FIRAppCheckTokenResultInterop /// App Check token in the case of success or a dummy token in the case of a failure. diff --git a/FirebaseAppCheck/Sources/Core/FIRAppCheckTokenResult.h b/FirebaseAppCheck/Sources/Core/FIRAppCheckTokenResult.h index 37b57b2adef..9cf6ddadec1 100644 --- a/FirebaseAppCheck/Sources/Core/FIRAppCheckTokenResult.h +++ b/FirebaseAppCheck/Sources/Core/FIRAppCheckTokenResult.h @@ -20,6 +20,7 @@ NS_ASSUME_NONNULL_BEGIN +NS_SWIFT_SENDABLE @interface FIRAppCheckTokenResult : NSObject - (instancetype)initWithToken:(NSString *)token error:(nullable NSError *)error;