328 lines
12 KiB
Swift
328 lines
12 KiB
Swift
import Foundation
|
|
|
|
public struct VertexAIOAuthCredentials: Sendable {
|
|
public let accessToken: String
|
|
public let refreshToken: String
|
|
public let clientId: String
|
|
public let clientSecret: String
|
|
public let projectId: String?
|
|
public let email: String?
|
|
public let expiryDate: Date?
|
|
|
|
public init(
|
|
accessToken: String,
|
|
refreshToken: String,
|
|
clientId: String,
|
|
clientSecret: String,
|
|
projectId: String?,
|
|
email: String?,
|
|
expiryDate: Date?)
|
|
{
|
|
self.accessToken = accessToken
|
|
self.refreshToken = refreshToken
|
|
self.clientId = clientId
|
|
self.clientSecret = clientSecret
|
|
self.projectId = projectId
|
|
self.email = email
|
|
self.expiryDate = expiryDate
|
|
}
|
|
|
|
public var needsRefresh: Bool {
|
|
guard let expiryDate else { return true }
|
|
// Refresh 5 minutes before expiry
|
|
return Date().addingTimeInterval(300) > expiryDate
|
|
}
|
|
}
|
|
|
|
public enum VertexAIOAuthCredentialsError: LocalizedError, Sendable {
|
|
case notFound
|
|
case decodeFailed(String)
|
|
case missingTokens
|
|
case missingClientCredentials
|
|
|
|
public var errorDescription: String? {
|
|
switch self {
|
|
case .notFound:
|
|
"gcloud credentials not found. Run `gcloud auth application-default login` to authenticate."
|
|
case let .decodeFailed(message):
|
|
"Failed to decode gcloud credentials: \(message)"
|
|
case .missingTokens:
|
|
"gcloud credentials exist but contain no tokens."
|
|
case .missingClientCredentials:
|
|
"gcloud credentials missing client ID or secret."
|
|
}
|
|
}
|
|
}
|
|
|
|
public enum VertexAIOAuthCredentialsStore {
|
|
#if DEBUG
|
|
@TaskLocal static var gcloudAccessTokenOverrideForTesting: (@Sendable ([String: String]) async throws -> String)?
|
|
#endif
|
|
|
|
private struct ServiceAccountMetadata {
|
|
let email: String
|
|
let projectId: String?
|
|
}
|
|
|
|
private static func credentialsFilePath(
|
|
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL
|
|
{
|
|
if let path = environment["GOOGLE_APPLICATION_CREDENTIALS"]?.trimmingCharacters(
|
|
in: .whitespacesAndNewlines),
|
|
!path.isEmpty
|
|
{
|
|
return URL(fileURLWithPath: path)
|
|
}
|
|
|
|
let home = FileManager.default.homeDirectoryForCurrentUser
|
|
// gcloud application default credentials location
|
|
if let configDir = environment["CLOUDSDK_CONFIG"]?.trimmingCharacters(
|
|
in: .whitespacesAndNewlines),
|
|
!configDir.isEmpty
|
|
{
|
|
return URL(fileURLWithPath: configDir)
|
|
.appendingPathComponent("application_default_credentials.json")
|
|
}
|
|
return home
|
|
.appendingPathComponent(".config")
|
|
.appendingPathComponent("gcloud")
|
|
.appendingPathComponent("application_default_credentials.json")
|
|
}
|
|
|
|
private static func projectFilePath(
|
|
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL
|
|
{
|
|
let home = FileManager.default.homeDirectoryForCurrentUser
|
|
if let configDir = environment["CLOUDSDK_CONFIG"]?.trimmingCharacters(
|
|
in: .whitespacesAndNewlines),
|
|
!configDir.isEmpty
|
|
{
|
|
return URL(fileURLWithPath: configDir)
|
|
.appendingPathComponent("configurations")
|
|
.appendingPathComponent("config_default")
|
|
}
|
|
return home
|
|
.appendingPathComponent(".config")
|
|
.appendingPathComponent("gcloud")
|
|
.appendingPathComponent("configurations")
|
|
.appendingPathComponent("config_default")
|
|
}
|
|
|
|
public static func hasCredentials(
|
|
environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool
|
|
{
|
|
let url = self.credentialsFilePath(environment: environment)
|
|
guard FileManager.default.fileExists(atPath: url.path),
|
|
let data = try? Data(contentsOf: url),
|
|
let json = try? self.parseJSONObject(data: data)
|
|
else {
|
|
return false
|
|
}
|
|
|
|
if self.parseServiceAccountMetadata(json: json) != nil {
|
|
return true
|
|
}
|
|
|
|
return (try? self.parseUserCredentials(json: json, environment: environment)) != nil
|
|
}
|
|
|
|
public static func load(
|
|
environment: [String: String] = ProcessInfo.processInfo.environment) throws -> VertexAIOAuthCredentials
|
|
{
|
|
let url = self.credentialsFilePath(environment: environment)
|
|
guard FileManager.default.fileExists(atPath: url.path) else {
|
|
throw VertexAIOAuthCredentialsError.notFound
|
|
}
|
|
|
|
let data = try Data(contentsOf: url)
|
|
return try self.parse(data: data, environment: environment)
|
|
}
|
|
|
|
public static func loadForFetch(
|
|
environment: [String: String] = ProcessInfo.processInfo.environment) async throws -> VertexAIOAuthCredentials
|
|
{
|
|
let url = self.credentialsFilePath(environment: environment)
|
|
guard FileManager.default.fileExists(atPath: url.path) else {
|
|
throw VertexAIOAuthCredentialsError.notFound
|
|
}
|
|
|
|
let data = try Data(contentsOf: url)
|
|
let json = try self.parseJSONObject(data: data)
|
|
if let serviceAccount = self.parseServiceAccountMetadata(json: json) {
|
|
let token = try await self.printAccessToken(environment: environment)
|
|
return VertexAIOAuthCredentials(
|
|
accessToken: token,
|
|
refreshToken: "",
|
|
clientId: "",
|
|
clientSecret: "",
|
|
projectId: serviceAccount.projectId ?? self.loadProjectId(environment: environment),
|
|
email: serviceAccount.email,
|
|
expiryDate: Date().addingTimeInterval(50 * 60))
|
|
}
|
|
|
|
return try self.parseUserCredentials(json: json, environment: environment)
|
|
}
|
|
|
|
public static func parse(data: Data) throws -> VertexAIOAuthCredentials {
|
|
try self.parse(data: data, environment: ProcessInfo.processInfo.environment)
|
|
}
|
|
|
|
public static func parse(
|
|
data: Data,
|
|
environment: [String: String]) throws -> VertexAIOAuthCredentials
|
|
{
|
|
let json = try self.parseJSONObject(data: data)
|
|
return try self.parseUserCredentials(json: json, environment: environment)
|
|
}
|
|
|
|
private static func parseJSONObject(data: Data) throws -> [String: Any] {
|
|
guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
|
|
throw VertexAIOAuthCredentialsError.decodeFailed("Invalid JSON")
|
|
}
|
|
return json
|
|
}
|
|
|
|
private static func parseUserCredentials(
|
|
json: [String: Any],
|
|
environment: [String: String]) throws -> VertexAIOAuthCredentials
|
|
{
|
|
// Check for service account credentials
|
|
if self.parseServiceAccountMetadata(json: json) != nil {
|
|
throw VertexAIOAuthCredentialsError.decodeFailed(
|
|
"Service account credentials require `gcloud auth application-default print-access-token`.")
|
|
}
|
|
|
|
// User credentials from gcloud auth application-default login
|
|
guard let clientId = json["client_id"] as? String,
|
|
let clientSecret = json["client_secret"] as? String
|
|
else {
|
|
throw VertexAIOAuthCredentialsError.missingClientCredentials
|
|
}
|
|
|
|
guard let refreshToken = json["refresh_token"] as? String, !refreshToken.isEmpty else {
|
|
throw VertexAIOAuthCredentialsError.missingTokens
|
|
}
|
|
|
|
// Access token may not be present in the file; we'll need to refresh
|
|
let accessToken = json["access_token"] as? String ?? ""
|
|
|
|
// Try to get project ID from gcloud config
|
|
let projectId = Self.loadProjectId(environment: environment)
|
|
|
|
// Try to extract email from ID token if present
|
|
let email = Self.extractEmailFromIdToken(json["id_token"] as? String)
|
|
|
|
// Parse expiry if present
|
|
var expiryDate: Date?
|
|
if let expiryStr = json["token_expiry"] as? String {
|
|
let formatter = ISO8601DateFormatter()
|
|
expiryDate = formatter.date(from: expiryStr)
|
|
}
|
|
|
|
return VertexAIOAuthCredentials(
|
|
accessToken: accessToken,
|
|
refreshToken: refreshToken,
|
|
clientId: clientId,
|
|
clientSecret: clientSecret,
|
|
projectId: projectId,
|
|
email: email,
|
|
expiryDate: expiryDate)
|
|
}
|
|
|
|
public static func save(_ credentials: VertexAIOAuthCredentials) throws {
|
|
// We don't modify gcloud's credentials file; just cache the access token in memory
|
|
// The refresh happens on each app launch if needed
|
|
}
|
|
|
|
private static func parseServiceAccountMetadata(json: [String: Any]) -> ServiceAccountMetadata? {
|
|
guard let email = (json["client_email"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines),
|
|
!email.isEmpty,
|
|
let privateKey = (json["private_key"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines),
|
|
!privateKey.isEmpty
|
|
else {
|
|
return nil
|
|
}
|
|
|
|
let projectId = (json["project_id"] as? String)?
|
|
.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
return ServiceAccountMetadata(
|
|
email: email,
|
|
projectId: projectId?.isEmpty == false ? projectId : nil)
|
|
}
|
|
|
|
private static func printAccessToken(environment: [String: String]) async throws -> String {
|
|
#if DEBUG
|
|
if let override = self.gcloudAccessTokenOverrideForTesting {
|
|
let token = try await override(environment)
|
|
return try self.cleanAccessToken(token)
|
|
}
|
|
#endif
|
|
|
|
let env = TTYCommandRunner.enrichedEnvironment(baseEnv: environment)
|
|
let result = try await SubprocessRunner.run(
|
|
binary: "/usr/bin/env",
|
|
arguments: ["gcloud", "auth", "application-default", "print-access-token"],
|
|
environment: env,
|
|
timeout: 20,
|
|
label: "vertexai-gcloud-adc-token")
|
|
return try self.cleanAccessToken(result.stdout)
|
|
}
|
|
|
|
private static func cleanAccessToken(_ token: String) throws -> String {
|
|
let trimmed = token.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
guard !trimmed.isEmpty else {
|
|
throw VertexAIOAuthCredentialsError.missingTokens
|
|
}
|
|
return trimmed
|
|
}
|
|
|
|
private static func loadProjectId(environment: [String: String]) -> String? {
|
|
let configPath = self.projectFilePath(environment: environment)
|
|
guard let content = try? String(contentsOf: configPath, encoding: .utf8) else {
|
|
return environment["GOOGLE_CLOUD_PROJECT"]
|
|
?? environment["GCLOUD_PROJECT"]
|
|
?? environment["CLOUDSDK_CORE_PROJECT"]
|
|
}
|
|
|
|
// Parse INI-style config for project
|
|
for line in content.components(separatedBy: .newlines) {
|
|
let trimmed = line.trimmingCharacters(in: .whitespaces)
|
|
if trimmed.hasPrefix("project") {
|
|
let parts = trimmed.components(separatedBy: "=")
|
|
if parts.count >= 2 {
|
|
return parts[1].trimmingCharacters(in: .whitespaces)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Try environment variable
|
|
return environment["GOOGLE_CLOUD_PROJECT"]
|
|
?? environment["GCLOUD_PROJECT"]
|
|
?? environment["CLOUDSDK_CORE_PROJECT"]
|
|
}
|
|
|
|
private static func extractEmailFromIdToken(_ token: String?) -> String? {
|
|
guard let token, !token.isEmpty else { return nil }
|
|
|
|
let parts = token.components(separatedBy: ".")
|
|
guard parts.count >= 2 else { return nil }
|
|
|
|
var payload = parts[1]
|
|
.replacingOccurrences(of: "-", with: "+")
|
|
.replacingOccurrences(of: "_", with: "/")
|
|
|
|
let remainder = payload.count % 4
|
|
if remainder > 0 {
|
|
payload += String(repeating: "=", count: 4 - remainder)
|
|
}
|
|
|
|
guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters),
|
|
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
|
|
else {
|
|
return nil
|
|
}
|
|
|
|
return json["email"] as? String
|
|
}
|
|
}
|