From 1737e8ad6ce97146903fbd6a9228e76070f77c86 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Tue, 3 Nov 2026 01:35:53 +0700 Subject: [PATCH 1/2] fix(plugin-snowflake): keep Snowflake sign-in working after a downgrade --- .claude/rules/data-sync-storage.md | 129 ++++++++++++++++++ .../SnowflakeTokenCache.swift | 104 ++++++++++++++ .../SnowflakeTokenCacheTests.swift | 111 +++++++++++++++ .../SnowflakeConnection.swift | 12 +- .../SnowflakeIdTokenStore.swift | 86 +----------- .../SnowflakeMFATokenStore.swift | 100 ++------------ 6 files changed, 370 insertions(+), 172 deletions(-) create mode 100644 apple/Packages/TableProCore/Sources/TableProSnowflakeCore/SnowflakeTokenCache.swift create mode 100644 apple/Packages/TableProCore/Tests/TableProSnowflakeCoreTests/SnowflakeTokenCacheTests.swift diff --git a/.claude/rules/data-sync-storage.md b/.claude/rules/data-sync-storage.md index 4ff769e5c5..b508dd018e 100644 --- a/.claude/rules/data-sync-storage.md +++ b/.claude/rules/data-sync-storage.md @@ -7,6 +7,134 @@ paths: - "apple/TablePro/Core/Services/Query/**/*" - "apple/Packages/TableProCore/Sources/TableProSyncTransport/**/*" - "apple/CloudKit/**/*" + - "apple/TablePro/Core/Storage/**/*" + - "apple/TablePro/Models/Settings/**/*" + - "apple/TablePro/Models/Connection/*Mode.swift" + - "apple/TablePro/Models/Connection/*Configuration.swift" + - "apple/TablePro/Models/Connection/*Profile.swift" + - "apple/TablePro/Models/Query/QueryHistory*.swift" + - "apple/TablePro/Models/Query/SQLFavorite*.swift" + - "apple/TablePro/Core/Services/Execution/ExecutionAudit*.swift" + - "apple/TablePro/Core/Services/Infrastructure/TabPersistenceCoordinator*.swift" + - "apple/**/*Keychain*.swift" + - "apple/**/*TokenStore.swift" + - "apple/Plugins/TableProPluginKit/Plugin*Storage.swift" + - "apple/Plugins/*/*ExportModels.swift" + - "apple/Plugins/*/*portOptions.swift" + - "apple/TableProMobile/TableProMobile/Sync/**/*" + - "apple/Packages/TableProCore/Sources/TableProSync/**/*" + - "apple/TablePro/Models/Connection/DatabaseConnection.swift" + - "apple/TablePro/Models/Connection/DatabaseType.swift" + - "apple/TablePro/Models/Connection/PasswordSource.swift" + - "apple/TablePro/Models/Connection/VaultDatabaseRole.swift" + - "apple/TablePro/Models/Connection/SafeModeLevel.swift" + - "apple/TablePro/Models/Connection/ConnectionGroup.swift" + - "apple/TablePro/Models/Connection/ConnectionTag.swift" + - "apple/TablePro/Models/Connection/ConnectionOrigin.swift" + - "apple/TablePro/Models/AI/AIConversation.swift" + - "apple/TablePro/Models/AI/AIModels.swift" + - "apple/TablePro/Models/AI/AgentSession.swift" + - "apple/TablePro/Models/AI/AgentSessionStatus.swift" + - "apple/TablePro/Models/AI/CustomSlashCommand.swift" + - "apple/TablePro/Models/Query/QueryTabState.swift" + - "apple/TablePro/Models/Query/QueryParameter.swift" + - "apple/TablePro/Models/Query/DatabaseObjectRef.swift" + - "apple/TablePro/Models/Query/ObjectAttribute.swift" + - "apple/TablePro/Models/Query/LinkedSQLFolder.swift" + - "apple/TablePro/Models/Query/QueryInsights.swift" + - "apple/TablePro/Models/UI/FilterPreset.swift" + - "apple/TablePro/Models/UI/FilterState.swift" + - "apple/TablePro/Models/UI/KeyboardShortcutModels.swift" + - "apple/TablePro/Models/UI/BoundKey.swift" + - "apple/TablePro/Models/UI/WindowSidebarState.swift" + - "apple/TablePro/Models/UI/SharedSidebarState.swift" + - "apple/TablePro/Models/UI/TrailingPaneState.swift" + - "apple/TablePro/Models/UI/RowInspectorState.swift" + - "apple/TablePro/Models/UI/SidebarRowSizePreference.swift" + - "apple/TablePro/Models/UI/HistoryDateRange.swift" + - "apple/TablePro/Models/AuditEntry.swift" + - "apple/TablePro/Models/Database/TableFilter.swift" + - "apple/TablePro/Models/Favorites/FavoriteDatabaseEntry.swift" + - "apple/TablePro/Models/Favorites/FavoriteDatabaseEnvironment.swift" + - "apple/TablePro/Models/Highlight/HighlightRule.swift" + - "apple/TablePro/Models/ServerDashboard/DashboardTile.swift" + - "apple/TablePro/Models/Sidebar/TableFolder.swift" + - "apple/TablePro/Models/TeamLibrary/TeamLibraryModels.swift" + - "apple/TablePro/Models/Workspace/WorkspaceID.swift" + - "apple/TablePro/Core/AI/Chat/ChatTurn.swift" + - "apple/TablePro/Core/AI/ChatGPTCodex/ChatGPTCodexModels.swift" + - "apple/TablePro/Core/ConnectionLibrary/RecentConnectionsStore.swift" + - "apple/TablePro/Core/ConnectionLibrary/ConnectionListPreferences.swift" + - "apple/TablePro/Core/Diagnostics/TableLoadHistoryStore.swift" + - "apple/TablePro/Core/Diagnostics/TableLoadPerformanceRecord.swift" + - "apple/TablePro/Core/DataWrite/Rewind/RewindCipher.swift" + - "apple/TablePro/Core/DataWrite/Rewind/RewindRecord.swift" + - "apple/TablePro/Core/MCP/Outside/MCPServerStore.swift" + - "apple/TablePro/Core/Modules/AppModule.swift" + - "apple/TablePro/Core/Modules/AppModuleSettings.swift" + - "apple/TablePro/Core/Plugins/PluginManager.swift" + - "apple/TablePro/Core/Plugins/Registry/RegistryManifestVerifier.swift" + - "apple/TablePro/Core/RoutineDebugging/RoutineDebugOptionStore.swift" + - "apple/TablePro/Core/RowGeneration/RowGenerationModel.swift" + - "apple/TablePro/Core/SSH/HostKeyStore.swift" + - "apple/TablePro/Core/Services/Export/ImportColumnMatcher.swift" + - "apple/TablePro/Core/Services/Formatting/ValueDisplayFormat.swift" + - "apple/TablePro/Core/Services/TeamLibrary/TeamLibraryStore.swift" + - "apple/TablePro/Core/VersionHistory/VersionHistoryModels.swift" + - "apple/TablePro/Core/UsersRoles/PrivilegeExpansionStore.swift" + - "apple/TablePro/Core/Cloudflare/CloudflareTunnelManager.swift" + - "apple/TablePro/Core/TunnelCommand/TunnelCommandManager.swift" + - "apple/TablePro/Core/CloudSQL/CloudSQLProxyManager.swift" + - "apple/TablePro/Core/Services/Infrastructure/MainWindowToolbar.swift" + - "apple/TablePro/Core/Services/Infrastructure/MacAnalyticsProvider.swift" + - "apple/TablePro/Core/Services/Infrastructure/UpdateChannel.swift" + - "apple/TablePro/Core/Services/Infrastructure/SparkleUpdatePreferences.swift" + - "apple/TablePro/Core/Services/Infrastructure/SampleDatabaseService.swift" + - "apple/TablePro/Core/Services/Policy/ManagedPolicy.swift" + - "apple/TablePro/Core/Services/Support/SupportAskCoordinator.swift" + - "apple/TablePro/Core/Services/Support/UsageDayCounter.swift" + - "apple/TablePro/Theme/ThemeCatalog.swift" + - "apple/TablePro/ViewModels/SidebarViewModel.swift" + - "apple/TablePro/ViewModels/FavoritesExpansionState.swift" + - "apple/TablePro/ViewModels/WelcomeGroupExpansionStore.swift" + - "apple/TablePro/ViewModels/WelcomeViewModel+Sample.swift" + - "apple/TablePro/ViewModels/QuickSwitcherViewModel.swift" + - "apple/TablePro/Views/Sidebar/SidebarPersistenceKey.swift" + - "apple/TablePro/Views/Structure/StructureGridStateStore.swift" + - "apple/TablePro/Views/Structure/StructureDefinitionEditor.swift" + - "apple/TablePro/Views/Settings/SettingsView.swift" + - "apple/TablePro/Views/Settings/SettingsWindowController.swift" + - "apple/TablePro/Views/WhatsNew/WhatsNewNotice.swift" + - "apple/TablePro/Views/Compare/CompareSyncWindowController.swift" + - "apple/TablePro/Views/Components/TransferResultAlert.swift" + - "apple/TablePro/Views/Shared/ObjectSourceView.swift" + - "apple/TablePro/Views/Compare/CompareApplySheetView.swift" + - "apple/TablePro/Views/Compare/CompareScriptPane.swift" + - "apple/TablePro/Views/Structure/TriggerEditorView.swift" + - "apple/TablePro/Views/Structure/StructureDDLPane.swift" + - "apple/TablePro/Views/Structure/CreateTableView.swift" + - "apple/TablePro/Views/Export/ExportDialog.swift" + - "apple/Packages/TableProMacCore/Sources/TableProSSHConfig/Models/SSHTypes.swift" + - "apple/Packages/TableProMacCore/Sources/TableProUtilities/File/FileTextEncoding.swift" + - "apple/Packages/TableProMacCore/Sources/TableProUtilities/UI/QuickSwitcherFrecencyStore.swift" + - "apple/Packages/TableProMacCore/Sources/TableProMCP/Outside/MCPServerConfiguration.swift" + - "apple/Packages/TableProCore/Sources/TableProSSHTransport/PublishedSSHRoute.swift" + - "apple/Packages/TableProCore/Sources/TableProSSHTransport/TOTPConfiguration.swift" + - "apple/Packages/TableProCore/Sources/TableProConnectionLibrary/RecentConnectionsLedger.swift" + - "apple/Packages/TableProCore/Sources/TableProSnowflakeCore/SnowflakeAccount.swift" + - "apple/Packages/TableProCore/Sources/TableProSnowflakeCore/SnowflakeTokenCache.swift" + - "apple/Packages/TableProCore/Sources/TableProAnalytics/AnalyticsHeartbeatService.swift" + - "apple/Packages/TableProCore/Sources/TableProTheme/ThemeFileFormat.swift" + - "apple/Packages/TableProCore/Sources/TableProTheme/ThemeDefinition.swift" + - "apple/Packages/TableProCore/Sources/TableProTheme/ThemeSlot.swift" + - "apple/Packages/TableProCore/Sources/TableProModels/SSHConfiguration.swift" + - "apple/Packages/TableProCore/Sources/TableProModels/SSLConfiguration.swift" + - "apple/Packages/TableProCore/Sources/TableProModels/SafeModeLevel.swift" + - "apple/Packages/TableProCore/Sources/TableProModels/ConnectionColor.swift" + - "apple/Plugins/TableProPluginKit/SettablePlugin.swift" + - "apple/Plugins/TableProPluginKit/SSLConfiguration.swift" + - "apple/Plugins/TableProPluginKit/PluginExportRowScope.swift" + - "apple/Plugins/TableProPluginKit/Entra/EntraCredentialResolver.swift" --- # Storage, sync and the database core @@ -46,6 +174,7 @@ This path holds user data. A silent wrong answer here looks exactly like an empt | Trusted external links (loopback only) | UserDefaults | `ExternalConnectionTrustStore` | | Highest plugin manifest `serial` accepted, so an older signed manifest is refused | UserDefaults, device-local, never synced | `RegistryManifestVerifier` through `RegistryClient` | | Approved environment variable references, per connection, names and destination digest | UserDefaults, device-local, never synced | `EnvVarReferenceApprovalStore` | +| Snowflake SSO and MFA tokens, per host and user. A bare account's `NAME.USER` item stays for 0.83.1, which reads only that key | Keychain, device-local | `SnowflakeTokenCache` | | Where a connection came from (`origin`): this Mac, imported, iCloud | `connections.json`, never synced or exported | `ConnectionStorage` | | Table load timings (no identifying data, never uploaded) | JSON lines | `TableLoadHistoryStore` | | Reopen Last Session list (`LastSession.json`, plus `LastOpenConnections.json` holding only the ids that reconnect, which older builds read) | JSON | `LastOpenConnectionsStorage` | diff --git a/apple/Packages/TableProCore/Sources/TableProSnowflakeCore/SnowflakeTokenCache.swift b/apple/Packages/TableProCore/Sources/TableProSnowflakeCore/SnowflakeTokenCache.swift new file mode 100644 index 0000000000..064e1ee143 --- /dev/null +++ b/apple/Packages/TableProCore/Sources/TableProSnowflakeCore/SnowflakeTokenCache.swift @@ -0,0 +1,104 @@ +import Foundation +import os +import Security + +public protocol SnowflakeTokenStorage: Sendable { + func read(key: String) -> String? + @discardableResult + func write(_ value: String, key: String) -> Bool + func delete(key: String) +} + +public final class SnowflakeTokenCache: Sendable { + private let storage: any SnowflakeTokenStorage + private let memory = OSAllocatedUnfairLock(initialState: [String: String]()) + + public init(storage: any SnowflakeTokenStorage) { + self.storage = storage + } + + public func token(for account: SnowflakeAccount, user: String) -> String? { + let key = account.tokenKey(user: user) + if let cached = memory.withLock({ $0[key] }) { + return cached + } + if let stored = storage.read(key: key) { + memory.withLock { $0[key] = stored } + return stored + } + guard let legacyKey = account.legacyTokenKey(user: user), let legacy = storage.read(key: legacyKey) else { + return nil + } + memory.withLock { $0[key] = legacy } + // The old item stays for 0.83.1, which reads only it; once copied, this build never reads it again. + storage.write(legacy, key: key) + return legacy + } + + public func store(_ token: String, for account: SnowflakeAccount, user: String) { + let key = account.tokenKey(user: user) + memory.withLock { $0[key] = token } + storage.write(token, key: key) + } + + public func clear(for account: SnowflakeAccount, user: String) { + let key = account.tokenKey(user: user) + memory.withLock { $0[key] = nil } + storage.delete(key: key) + if let legacyKey = account.legacyTokenKey(user: user) { + storage.delete(key: legacyKey) + } + } +} + +public struct SnowflakeKeychainTokenStorage: SnowflakeTokenStorage { + private static let logger = Logger(subsystem: "com.TablePro", category: "SnowflakeTokenCache") + + public let service: String + + public init(service: String) { + self.service = service + } + + public func read(key: String) -> String? { + var query = baseQuery(key: key) + query[kSecReturnData as String] = true + query[kSecMatchLimit as String] = kSecMatchLimitOne + var result: AnyObject? + let status = SecItemCopyMatching(query as CFDictionary, &result) + guard status == errSecSuccess, let data = result as? Data else { return nil } + return String(data: data, encoding: .utf8) + } + + @discardableResult + public func write(_ value: String, key: String) -> Bool { + var addQuery = baseQuery(key: key) + addQuery[kSecValueData as String] = Data(value.utf8) + addQuery[kSecAttrAccessible as String] = kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly + var status = SecItemAdd(addQuery as CFDictionary, nil) + if status == errSecDuplicateItem { + status = SecItemUpdate( + baseQuery(key: key) as CFDictionary, + [kSecValueData as String: Data(value.utf8)] as CFDictionary + ) + } + if status != errSecSuccess { + Self.logger.warning( + "Could not persist a token in \(service, privacy: .public) (OSStatus \(status)); it is cached in memory only" + ) + } + return status == errSecSuccess + } + + public func delete(key: String) { + SecItemDelete(baseQuery(key: key) as CFDictionary) + } + + private func baseQuery(key: String) -> [String: Any] { + [ + kSecClass as String: kSecClassGenericPassword, + kSecAttrService as String: service, + kSecAttrAccount as String: key + ] + } +} diff --git a/apple/Packages/TableProCore/Tests/TableProSnowflakeCoreTests/SnowflakeTokenCacheTests.swift b/apple/Packages/TableProCore/Tests/TableProSnowflakeCoreTests/SnowflakeTokenCacheTests.swift new file mode 100644 index 0000000000..9b62c4ffd9 --- /dev/null +++ b/apple/Packages/TableProCore/Tests/TableProSnowflakeCoreTests/SnowflakeTokenCacheTests.swift @@ -0,0 +1,111 @@ +import Foundation +import os +import TableProSnowflakeCore +import Testing + +private final class InMemoryTokenStorage: SnowflakeTokenStorage { + private let items: OSAllocatedUnfairLock<[String: String]> + + init(_ initial: [String: String] = [:]) { + items = OSAllocatedUnfairLock(initialState: initial) + } + + var snapshot: [String: String] { items.withLock { $0 } } + + func read(key: String) -> String? { items.withLock { $0[key] } } + + func write(_ value: String, key: String) -> Bool { + items.withLock { $0[key] = value } + return true + } + + func delete(key: String) { items.withLock { $0[key] = nil } } + + func set(_ value: String?, key: String) { items.withLock { $0[key] = value } } +} + +struct SnowflakeTokenCacheTests { + private let hostKey = "acme.snowflakecomputing.com|BOB" + private let legacyKey = "ACME.BOB" + + @Test("A token under the old key is copied to the host key and the old item stays for the previous release") + func legacyItemSurvivesTheCopy() throws { + let storage = InMemoryTokenStorage([legacyKey: "old"]) + let cache = SnowflakeTokenCache(storage: storage) + + #expect(cache.token(for: try #require(SnowflakeAccount("acme")), user: "bob") == "old") + #expect(storage.snapshot == [legacyKey: "old", hostKey: "old"]) + } + + @Test("Once copied, a token the previous release later writes under the old key is never read") + func hostKeyWinsAfterTheCopy() throws { + let account = try #require(SnowflakeAccount("acme")) + let storage = InMemoryTokenStorage([legacyKey: "old"]) + _ = SnowflakeTokenCache(storage: storage).token(for: account, user: "bob") + + storage.set("written-for-acme.us-east-1", key: legacyKey) + + #expect(SnowflakeTokenCache(storage: storage).token(for: account, user: "bob") == "old") + } + + @Test("An account other than a bare one on the default domain never reads the old key", arguments: [ + "acme.us-east-1", + "acme.us-east-2.aws", + "acme.snowflakecomputing.cn", + ]) + func otherHostsNeverReadTheOldKey(text: String) throws { + let storage = InMemoryTokenStorage([legacyKey: "old"]) + let cache = SnowflakeTokenCache(storage: storage) + + #expect(cache.token(for: try #require(SnowflakeAccount(text)), user: "bob") == nil) + #expect(storage.snapshot == [legacyKey: "old"]) + } + + @Test("A token stored for one host is never returned for another host of the same name") + func tokensStayOnTheirHost() throws { + let cache = SnowflakeTokenCache(storage: InMemoryTokenStorage()) + cache.store("regional", for: try #require(SnowflakeAccount("acme.us-east-1")), user: "bob") + + #expect(cache.token(for: try #require(SnowflakeAccount("acme")), user: "bob") == nil) + } + + @Test("Storing writes the host key only and leaves the old item as the previous release left it") + func storeLeavesTheOldItem() throws { + let storage = InMemoryTokenStorage([legacyKey: "old"]) + let cache = SnowflakeTokenCache(storage: storage) + let account = try #require(SnowflakeAccount("acme")) + + cache.store("new", for: account, user: "bob") + + #expect(storage.snapshot == [legacyKey: "old", hostKey: "new"]) + #expect(cache.token(for: account, user: "bob") == "new") + } + + @Test("Clearing a refused token removes the copy and the old item, so it never comes back") + func clearRemovesBothItems() throws { + let storage = InMemoryTokenStorage([legacyKey: "old"]) + let cache = SnowflakeTokenCache(storage: storage) + let account = try #require(SnowflakeAccount("acme")) + _ = cache.token(for: account, user: "bob") + + cache.clear(for: account, user: "bob") + + #expect(storage.snapshot.isEmpty) + #expect(cache.token(for: account, user: "bob") == nil) + #expect(SnowflakeTokenCache(storage: storage).token(for: account, user: "bob") == nil) + } + + @Test("The Keychain storage round-trips when a Keychain is available") + func keychainRoundTrip() { + let storage = SnowflakeKeychainTokenStorage(service: "com.TablePro.tests.snowflake.\(UUID().uuidString)") + let key = "acme.snowflakecomputing.com|BOB" + defer { storage.delete(key: key) } + + guard storage.write("first", key: key) else { return } + #expect(storage.read(key: key) == "first") + #expect(storage.write("second", key: key)) + #expect(storage.read(key: key) == "second") + storage.delete(key: key) + #expect(storage.read(key: key) == nil) + } +} diff --git a/apple/Plugins/SnowflakeDriverPlugin/SnowflakeConnection.swift b/apple/Plugins/SnowflakeDriverPlugin/SnowflakeConnection.swift index 26c2e6febb..2bbb828161 100644 --- a/apple/Plugins/SnowflakeDriverPlugin/SnowflakeConnection.swift +++ b/apple/Plugins/SnowflakeDriverPlugin/SnowflakeConnection.swift @@ -199,8 +199,10 @@ final class SnowflakeConnection: @unchecked Sendable { ) return } catch { - SnowflakeMFATokenStore.clear(for: account, user: params.user) - Self.logger.info("Cached MFA token rejected; retrying with credentials") + if case .loginFailed? = error as? SnowflakeError { + SnowflakeMFATokenStore.clear(for: account, user: params.user) + } + Self.logger.info("Login with the cached MFA token failed; retrying with credentials") } } @@ -244,8 +246,10 @@ final class SnowflakeConnection: @unchecked Sendable { try await login(authenticator: "ID_TOKEN", extra: ["TOKEN": idToken], deadline: deadline) return } catch { - SnowflakeIdTokenStore.clear(for: account, user: params.user) - Self.logger.info("Cached SSO id token rejected; falling back to browser authentication") + if case .loginFailed? = error as? SnowflakeError { + SnowflakeIdTokenStore.clear(for: account, user: params.user) + } + Self.logger.info("Login with the cached SSO id token failed; falling back to browser authentication") } } diff --git a/apple/Plugins/SnowflakeDriverPlugin/SnowflakeIdTokenStore.swift b/apple/Plugins/SnowflakeDriverPlugin/SnowflakeIdTokenStore.swift index 2ee62dcdac..55b0cfb866 100644 --- a/apple/Plugins/SnowflakeDriverPlugin/SnowflakeIdTokenStore.swift +++ b/apple/Plugins/SnowflakeDriverPlugin/SnowflakeIdTokenStore.swift @@ -1,94 +1,20 @@ -// -// SnowflakeIdTokenStore.swift -// SnowflakeDriverPlugin -// -// Caches the idToken Snowflake issues after a successful external-browser SSO -// login (CLIENT_STORE_TEMPORARY_CREDENTIAL), so subsequent connects use the -// ID_TOKEN authenticator instead of reopening the browser. Mirrors the -// official connectors' temporary credential cache, stored in the Keychain. -// - import Foundation -import os -import Security import TableProSnowflakeCore enum SnowflakeIdTokenStore { - private static let cache = OSAllocatedUnfairLock(initialState: [String: String]()) - private static let service = "com.TablePro.SnowflakeDriverPlugin.idToken" - private static let logger = Logger(subsystem: "com.TablePro", category: "SnowflakeIdTokenStore") + private static let cache = SnowflakeTokenCache( + storage: SnowflakeKeychainTokenStorage(service: "com.TablePro.SnowflakeDriverPlugin.idToken") + ) static func token(for account: SnowflakeAccount, user: String) -> String? { - let key = account.tokenKey(user: user) - if let cached = cache.withLock({ $0[key] }) { - return cached - } - if let stored = readKeychain(key: key) { - cache.withLock { $0[key] = stored } - return stored - } - guard let legacyKey = account.legacyTokenKey(user: user), let legacy = readKeychain(key: legacyKey) else { - return nil - } - cache.withLock { $0[key] = legacy } - if writeKeychain(legacy, key: key) { - deleteKeychain(key: legacyKey) - } - return legacy + cache.token(for: account, user: user) } static func store(_ token: String, for account: SnowflakeAccount, user: String) { - let key = account.tokenKey(user: user) - cache.withLock { $0[key] = token } - writeKeychain(token, key: key) + cache.store(token, for: account, user: user) } static func clear(for account: SnowflakeAccount, user: String) { - let key = account.tokenKey(user: user) - cache.withLock { $0[key] = nil } - deleteKeychain(key: key) - if let legacyKey = account.legacyTokenKey(user: user) { - deleteKeychain(key: legacyKey) - } - } - - private static func baseQuery(key: String) -> [String: Any] { - [ - kSecClass as String: kSecClassGenericPassword, - kSecAttrService as String: service, - kSecAttrAccount as String: key - ] - } - - private static func readKeychain(key: String) -> String? { - var query = baseQuery(key: key) - query[kSecReturnData as String] = true - query[kSecMatchLimit as String] = kSecMatchLimitOne - var result: AnyObject? - let status = SecItemCopyMatching(query as CFDictionary, &result) - guard status == errSecSuccess, let data = result as? Data else { return nil } - return String(data: data, encoding: .utf8) - } - - @discardableResult - private static func writeKeychain(_ value: String, key: String) -> Bool { - var addQuery = baseQuery(key: key) - addQuery[kSecValueData as String] = Data(value.utf8) - addQuery[kSecAttrAccessible as String] = kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly - var status = SecItemAdd(addQuery as CFDictionary, nil) - if status == errSecDuplicateItem { - status = SecItemUpdate( - baseQuery(key: key) as CFDictionary, - [kSecValueData as String: Data(value.utf8)] as CFDictionary - ) - } - if status != errSecSuccess { - logger.warning("Could not persist SSO id token (OSStatus \(status)); it will be cached in memory only") - } - return status == errSecSuccess - } - - private static func deleteKeychain(key: String) { - SecItemDelete(baseQuery(key: key) as CFDictionary) + cache.clear(for: account, user: user) } } diff --git a/apple/Plugins/SnowflakeDriverPlugin/SnowflakeMFATokenStore.swift b/apple/Plugins/SnowflakeDriverPlugin/SnowflakeMFATokenStore.swift index 2dac26af5b..b2850483ed 100644 --- a/apple/Plugins/SnowflakeDriverPlugin/SnowflakeMFATokenStore.swift +++ b/apple/Plugins/SnowflakeDriverPlugin/SnowflakeMFATokenStore.swift @@ -1,59 +1,23 @@ -// -// SnowflakeMFATokenStore.swift -// SnowflakeDriverPlugin -// -// Caches the MFA token Snowflake issues after a successful TOTP login -// (CLIENT_REQUEST_MFA_TOKEN), so subsequent connects don't need a fresh -// passcode until the token expires. Mirrors the official connectors. -// - import Foundation import os -import Security import TableProSnowflakeCore enum SnowflakeMFATokenStore { - private struct TokenState { - var cache: [String: String] = [:] - var rejectedPasscodes: [String: Date] = [:] - } - - private static let state = OSAllocatedUnfairLock(initialState: TokenState()) - private static let service = "com.TablePro.SnowflakeDriverPlugin.mfaToken" - private static let logger = Logger(subsystem: "com.TablePro", category: "SnowflakeMFATokenStore") + private static let cache = SnowflakeTokenCache( + storage: SnowflakeKeychainTokenStorage(service: "com.TablePro.SnowflakeDriverPlugin.mfaToken") + ) + private static let rejectedPasscodes = OSAllocatedUnfairLock(initialState: [String: Date]()) static func token(for account: SnowflakeAccount, user: String) -> String? { - let key = account.tokenKey(user: user) - if let cached = state.withLock({ $0.cache[key] }) { - return cached - } - if let stored = readKeychain(key: key) { - state.withLock { $0.cache[key] = stored } - return stored - } - guard let legacyKey = account.legacyTokenKey(user: user), let legacy = readKeychain(key: legacyKey) else { - return nil - } - state.withLock { $0.cache[key] = legacy } - if writeKeychain(legacy, key: key) { - deleteKeychain(key: legacyKey) - } - return legacy + cache.token(for: account, user: user) } static func store(_ token: String, for account: SnowflakeAccount, user: String) { - let key = account.tokenKey(user: user) - state.withLock { $0.cache[key] = token } - writeKeychain(token, key: key) + cache.store(token, for: account, user: user) } static func clear(for account: SnowflakeAccount, user: String) { - let key = account.tokenKey(user: user) - state.withLock { $0.cache[key] = nil } - deleteKeychain(key: key) - if let legacyKey = account.legacyTokenKey(user: user) { - deleteKeychain(key: legacyKey) - } + cache.clear(for: account, user: user) } static let rejectedPasscodeLifetime: TimeInterval = 300 @@ -61,60 +25,20 @@ enum SnowflakeMFATokenStore { static func markPasscodeRejected(_ passcode: String, for account: SnowflakeAccount, user: String) { guard !passcode.isEmpty else { return } let key = "\(account.tokenKey(user: user)):\(passcode)" - state.withLock { state in - state.rejectedPasscodes = state.rejectedPasscodes.filter { + rejectedPasscodes.withLock { rejected in + rejected = rejected.filter { Date().timeIntervalSince($0.value) < rejectedPasscodeLifetime } - state.rejectedPasscodes[key] = Date() + rejected[key] = Date() } } static func isPasscodeRejected(_ passcode: String, for account: SnowflakeAccount, user: String) -> Bool { guard !passcode.isEmpty else { return false } let key = "\(account.tokenKey(user: user)):\(passcode)" - return state.withLock { - guard let rejectedAt = $0.rejectedPasscodes[key] else { return false } + return rejectedPasscodes.withLock { + guard let rejectedAt = $0[key] else { return false } return Date().timeIntervalSince(rejectedAt) < rejectedPasscodeLifetime } } - - private static func baseQuery(key: String) -> [String: Any] { - [ - kSecClass as String: kSecClassGenericPassword, - kSecAttrService as String: service, - kSecAttrAccount as String: key - ] - } - - private static func readKeychain(key: String) -> String? { - var query = baseQuery(key: key) - query[kSecReturnData as String] = true - query[kSecMatchLimit as String] = kSecMatchLimitOne - var result: AnyObject? - let status = SecItemCopyMatching(query as CFDictionary, &result) - guard status == errSecSuccess, let data = result as? Data else { return nil } - return String(data: data, encoding: .utf8) - } - - @discardableResult - private static func writeKeychain(_ value: String, key: String) -> Bool { - var addQuery = baseQuery(key: key) - addQuery[kSecValueData as String] = Data(value.utf8) - addQuery[kSecAttrAccessible as String] = kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly - var status = SecItemAdd(addQuery as CFDictionary, nil) - if status == errSecDuplicateItem { - status = SecItemUpdate( - baseQuery(key: key) as CFDictionary, - [kSecValueData as String: Data(value.utf8)] as CFDictionary - ) - } - if status != errSecSuccess { - logger.warning("Could not persist MFA token (OSStatus \(status)); it will be cached in memory only") - } - return status == errSecSuccess - } - - private static func deleteKeychain(key: String) { - SecItemDelete(baseQuery(key: key) as CFDictionary) - } } From 2bbac51664a920d60ba0ca66cc2a50880ee42995 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Tue, 3 Nov 2026 01:36:14 +0700 Subject: [PATCH 2/2] docs(changelog): add the Snowflake token fragment --- apple/changelog.d/3996.fixed.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 apple/changelog.d/3996.fixed.md diff --git a/apple/changelog.d/3996.fixed.md b/apple/changelog.d/3996.fixed.md new file mode 100644 index 0000000000..51311a67c9 --- /dev/null +++ b/apple/changelog.d/3996.fixed.md @@ -0,0 +1 @@ +A Snowflake connect that fails on the network no longer forgets the cached SSO or MFA token.