diff --git a/FirebaseAI/Sources/AppCheckOptions.swift b/FirebaseAI/Sources/AppCheckOptions.swift new file mode 100644 index 00000000000..6411eb39e12 --- /dev/null +++ b/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, Encodable { + /// 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 `limitedUseTokens` instead of the standard cached tokens. + /// + /// A new `limitedUseToken` will be generated for each request; providing a lower attack + /// surface for malicious parties to hijack tokens. When used alongside [replay protection](https://firebase.google.com/docs/app-check/custom-resource-backend#replay-protection), + /// `limitedUseTokens` are also _consumed_ after each request, ensuring they can't be used + /// again. + /// + /// _To prevent breakage, 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 + /// > `limitedUseTokens`. Because `limitedUseTokens` are backwards compatable, 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 `limitedUseTokens` ahead of time will also allow you to enable replay + /// > protection down the road (when support is added), without breaking your users. + public init(requireLimitedUseTokens: Bool = false) { + self.requireLimitedUseTokens = requireLimitedUseTokens + } +} diff --git a/FirebaseAI/Sources/FirebaseAI.swift b/FirebaseAI/Sources/FirebaseAI.swift index 48f7183d4e6..b9938ced25b 100644 --- a/FirebaseAI/Sources/FirebaseAI.swift +++ b/FirebaseAI/Sources/FirebaseAI.swift @@ -32,13 +32,19 @@ 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(appCheck: AppCheckOptions())) -> 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 +96,8 @@ public final class FirebaseAI: Sendable { tools: tools, toolConfig: toolConfig, systemInstruction: systemInstruction, - requestOptions: requestOptions + requestOptions: requestOptions, + aiConfig: aiConfig ) } @@ -126,7 +133,8 @@ public final class FirebaseAI: Sendable { apiConfig: apiConfig, generationConfig: generationConfig, safetySettings: safetySettings, - requestOptions: requestOptions + requestOptions: requestOptions, + aiConfig: aiConfig ) } @@ -134,6 +142,21 @@ public final class FirebaseAI: Sendable { /// 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, Encodable { + /// 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. @@ -141,6 +164,8 @@ public final class FirebaseAI: Sendable { 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 +181,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 +191,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 +230,7 @@ public final class FirebaseAI: Sendable { ) self.apiConfig = apiConfig self.location = location + self.aiConfig = aiConfig } func modelResourceName(modelName: String) -> String { @@ -249,5 +285,6 @@ public final class FirebaseAI: Sendable { let appName: String let location: String? let apiConfig: APIConfig + let aiConfig: FirebaseAI.Config } } diff --git a/FirebaseAI/Sources/GenerativeAIService.swift b/FirebaseAI/Sources/GenerativeAIService.swift index e1538af997f..beb0fcda3c3 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 aiConfig: FirebaseAI.Config + + init(firebaseInfo: FirebaseInfo, urlSession: URLSession, aiConfig: FirebaseAI.Config) { self.firebaseInfo = firebaseInfo self.urlSession = urlSession + self.aiConfig = aiConfig } 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 = await fetchAppCheckToken(appCheck: appCheck) urlRequest.setValue(tokenResult.token, forHTTPHeaderField: "X-Firebase-AppCheck") if let error = tokenResult.error { AILog.error( @@ -207,6 +210,23 @@ struct GenerativeAIService { return urlRequest } + private func fetchAppCheckToken(appCheck: AppCheckInterop) async + -> FIRAppCheckTokenResultInterop { + if aiConfig.appCheck.requireLimitedUseTokens { + if let token = await appCheck.getLimitedUseToken?() { + return token + } + + AILog.error( + code: .appCheckTokenFetchFailed, + "Missing getLimitedUseToken() function, but requireLimitedUseTokens was enabled." + ) + // falls back to standard token + } + + 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 diff --git a/FirebaseAI/Sources/GenerativeModel.swift b/FirebaseAI/Sources/GenerativeModel.swift index 8d3f5e043a7..0a229d3489c 100644 --- a/FirebaseAI/Sources/GenerativeModel.swift +++ b/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,14 @@ public final class GenerativeModel: Sendable { toolConfig: ToolConfig? = nil, systemInstruction: ModelContent? = nil, requestOptions: RequestOptions, - urlSession: URLSession = GenAIURLSession.default) { + urlSession: URLSession = GenAIURLSession.default, aiConfig: FirebaseAI.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 diff --git a/FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift b/FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift index e6f96df511a..254e1fe21ea 100644 --- a/FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift +++ b/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