fix: migrate storage to instanceName-apiKey to isolate by instance (#114

)

* fix: migrate storage to instanceName-apiKey to isolate by instance

* fix: only migration instance storage for sandboxed apps

* chore: add STORAGE_VERSION to userdefaults to keep track of migration status

* chore: add Configuration tests to check new storage paths

* chore: added tests to check migrations in sandboxed and non-sandboxed environments

Sources/Amplitude/Amplitude.swift

@@ -41,10 +41,10 @@ public class Amplitude {
 
         migrateApiKeyStorages()
         migrateDefaultInstanceStorages()
-
-        if configuration.migrateLegacyData {
+        if configuration.migrateLegacyData && getStorageVersion() < .API_KEY_AND_INSTANCE_NAME {
             RemnantDataMigration(self).execute()
         }
+        migrateInstanceOnlyStorages()
 
         if let deviceId: String? = configuration.storageProvider.read(key: .DEVICE_ID) {
             state.deviceId = deviceId
@@ -336,7 +336,17 @@ public class Amplitude {
         }
     }
 
+    private func getStorageVersion() -> PersistentStorageVersion {
+        let storageVersionInt: Int? = configuration.storageProvider.read(key: .STORAGE_VERSION)
+        let storageVersion: PersistentStorageVersion = (storageVersionInt == nil) ? PersistentStorageVersion.NO_VERSION : PersistentStorageVersion(rawValue: storageVersionInt!)!
+        return storageVersion
+    }
+
     private func migrateApiKeyStorages() {
+        if getStorageVersion() >= PersistentStorageVersion.API_KEY {
+            return
+        }
+        configuration.loggerProvider.debug(message: "Running migrateApiKeyStorages")
         if let persistentStorage = configuration.storageProvider as? PersistentStorage {
             let apiKeyStorage = PersistentStorage(storagePrefix: "\(PersistentStorage.DEFAULT_STORAGE_PREFIX)-\(configuration.apiKey)")
             StoragePrefixMigration(source: apiKeyStorage, destination: persistentStorage, logger: logger).execute()
@@ -349,10 +359,11 @@ public class Amplitude {
     }
 
     private func migrateDefaultInstanceStorages() {
-        if configuration.instanceName != Constants.Configuration.DEFAULT_INSTANCE {
+        if getStorageVersion() >= PersistentStorageVersion.INSTANCE_NAME ||
+            configuration.instanceName != Constants.Configuration.DEFAULT_INSTANCE {
             return
         }
-
+        configuration.loggerProvider.debug(message: "Running migrateDefaultInstanceStorages")
         let legacyDefaultInstanceName = "default_instance"
         if let persistentStorage = configuration.storageProvider as? PersistentStorage {
             let legacyStorage = PersistentStorage(storagePrefix: "storage-\(legacyDefaultInstanceName)")
@@ -364,4 +375,54 @@ public class Amplitude {
             StoragePrefixMigration(source: legacyIdentifyStorage, destination: persistentIdentifyStorage, logger: logger).execute()
         }
     }
+
+    internal func migrateInstanceOnlyStorages() {
+        if getStorageVersion() >= .API_KEY_AND_INSTANCE_NAME {
+            configuration.loggerProvider.debug(message: "Skipping migrateInstanceOnlyStorages based on STORAGE_VERSION")
+            return
+        }
+        configuration.loggerProvider.debug(message: "Running migrateInstanceOnlyStorages")
+
+        let skipEventMigration = !isSandboxEnabled()
+        // Only migrate sandboxed apps to avoid potential data pollution
+        if skipEventMigration {
+            configuration.loggerProvider.debug(message: "Skipping event migration in non-sandboxed app. Transfering UserDefaults only.")
+        }
+
+        let instanceName = configuration.getNormalizeInstanceName()
+        if let persistentStorage = configuration.storageProvider as? PersistentStorage {
+            let instanceOnlyEventPrefix = "\(PersistentStorage.DEFAULT_STORAGE_PREFIX)-storage-\(instanceName)"
+            let instanceNameOnlyStorage = PersistentStorage(storagePrefix: instanceOnlyEventPrefix)
+            StoragePrefixMigration(
+                source: instanceNameOnlyStorage,
+                destination: persistentStorage,
+                logger: logger
+            ).execute(skipEventFiles: skipEventMigration)
+        }
+
+        if let persistentIdentifyStorage = configuration.identifyStorageProvider as? PersistentStorage {
+            let instanceOnlyIdentifyPrefix = "\(PersistentStorage.DEFAULT_STORAGE_PREFIX)-identify-\(instanceName)"
+            let instanceNameOnlyIdentifyStorage = PersistentStorage(storagePrefix: instanceOnlyIdentifyPrefix)
+            StoragePrefixMigration(
+                source: instanceNameOnlyIdentifyStorage,
+                destination: persistentIdentifyStorage,
+                logger: logger
+            ).execute(skipEventFiles: skipEventMigration)
+        }
+
+        do {
+            // Store the current storage version
+            try configuration.storageProvider.write(
+                key: .STORAGE_VERSION,
+                value: PersistentStorageVersion.API_KEY_AND_INSTANCE_NAME.rawValue as Int
+            )
+            configuration.loggerProvider.debug(message: "Updated STORAGE_VERSION to .API_KEY_AND_INSTANCE_NAME")
+        } catch {
+            configuration.loggerProvider.error(message: "Unable to set STORAGE_VERSION in storageProvider during migration")
+        }
+    }
+
+    internal func isSandboxEnabled() -> Bool {
+        return SandboxHelper().isSandboxEnabled()
+    }
 }

Sources/Amplitude/Configuration.swift

@@ -64,17 +64,17 @@ public class Configuration {
         migrateLegacyData: Bool = true,
         offline: Bool? = false
     ) {
-        let normalizedInstanceName = instanceName == "" ? Constants.Configuration.DEFAULT_INSTANCE : instanceName
+        let normalizedInstanceName = Configuration.getNormalizeInstanceName(instanceName)
 
         self.apiKey = apiKey
         self.flushQueueSize = flushQueueSize
         self.flushIntervalMillis = flushIntervalMillis
         self.instanceName = normalizedInstanceName
         self.optOut = optOut
         self.storageProvider = storageProvider
-            ?? PersistentStorage(storagePrefix: "storage-\(normalizedInstanceName)")
+            ?? PersistentStorage(storagePrefix: PersistentStorage.getEventStoragePrefix(apiKey, normalizedInstanceName))
         self.identifyStorageProvider = identifyStorageProvider
-            ?? PersistentStorage(storagePrefix: "identify-\(normalizedInstanceName)")
+            ?? PersistentStorage(storagePrefix: PersistentStorage.getIdentifyStoragePrefix(apiKey, normalizedInstanceName))
         self.logLevel = logLevel
         self.loggerProvider = loggerProvider
         self.minIdLength = minIdLength
@@ -103,4 +103,12 @@ public class Configuration {
             && minTimeBetweenSessionsMillis > 0
             && (minIdLength == nil || minIdLength! > 0)
     }
+
+    private class func getNormalizeInstanceName(_ instanceName: String) -> String {
+        return instanceName == "" ? Constants.Configuration.DEFAULT_INSTANCE : instanceName
+    }
+
+    internal func getNormalizeInstanceName() -> String {
+        return Configuration.getNormalizeInstanceName(self.instanceName)
+    }
 }

Sources/Amplitude/Migration/StoragePrefixMigration.swift

@@ -11,12 +11,15 @@ class StoragePrefixMigration {
         self.logger = logger
     }
 
-    func execute() {
+    func execute(skipEventFiles: Bool = false) {
         if source.storagePrefix == destination.storagePrefix {
             return
         }
 
-        moveSourceEventFilesToDestination()
+        if !skipEventFiles {
+            moveSourceEventFilesToDestination()
+        }
+
         moveUserDefaults()
     }
 

Sources/Amplitude/Storages/PersistentStorage.swift

@@ -10,6 +10,14 @@ import Foundation
 class PersistentStorage: Storage {
     typealias EventBlock = URL
 
+    static internal func getEventStoragePrefix(_ apiKey: String, _ instanceName: String) -> String {
+        return "storage-\(apiKey)-\(instanceName)"
+    }
+
+    static internal func getIdentifyStoragePrefix(_ apiKey: String, _ instanceName: String) -> String {
+        return "identify-\(apiKey)-\(instanceName)"
+    }
+
     let storagePrefix: String
     let userDefaults: UserDefaults?
     let fileManager: FileManager

Sources/Amplitude/Types.swift

@@ -74,6 +74,23 @@ public enum StorageKey: String, CaseIterable {
     case DEVICE_ID = "device_id"
     case APP_BUILD = "app_build"
     case APP_VERSION = "app_version"
+    // The version of PersistentStorage, used for data migrations
+    // Value should be a PersistentStorageVersion value
+    // Note the first version is 2, which corresponds to apiKey-instanceName based storage
+    case STORAGE_VERSION = "storage_version"
+}
+
+public enum PersistentStorageVersion: Int, Comparable {
+    public static func < (lhs: PersistentStorageVersion, rhs: PersistentStorageVersion) -> Bool {
+        return lhs.rawValue < rhs.rawValue
+    }
+
+    case NO_VERSION = -1
+    // Note that versioning was added after these storage changes (0, 1)
+    case API_KEY = 0
+    case INSTANCE_NAME = 1
+    // This is the first version (2) we set a value in storageProvider.read(.StorageVersion)
+    case API_KEY_AND_INSTANCE_NAME = 2
 }
 
 public protocol Logger {

Tests/AmplitudeTests/AmplitudeTests.swift

@@ -324,6 +324,165 @@ final class AmplitudeTests: XCTestCase {
         ])
     }
 
+    func testMigrationToApiKeyAndInstanceNameStorage() throws {
+        let legacyUserId = "legacy-user-id"
+        let config = Configuration(
+            apiKey: "amp-migration-api-key",
+            // don't transfer any events
+            flushQueueSize: 1000,
+            flushIntervalMillis: 99999,
+            logLevel: LogLevelEnum.DEBUG,
+            defaultTracking: DefaultTrackingOptions.NONE
+        )
+
+        // Create storages using instance name only
+        let legacyEventStorage = PersistentStorage(storagePrefix: "storage-\(config.getNormalizeInstanceName())")
+        let legacyIdentityStorage = PersistentStorage(storagePrefix: "identify-\(config.getNormalizeInstanceName())")
+
+        // Init Amplitude using legacy storage
+        let legacyStorageAmplitude = FakeAmplitudeWithNoInstNameOnlyMigration(configuration: Configuration(
+            apiKey: config.apiKey,
+            flushQueueSize: config.flushQueueSize,
+            flushIntervalMillis: config.flushIntervalMillis,
+            storageProvider: legacyEventStorage,
+            identifyStorageProvider: legacyIdentityStorage,
+            logLevel: config.logLevel,
+            defaultTracking: config.defaultTracking
+        ))
+
+        let legacyDeviceId = legacyStorageAmplitude.getDeviceId()
+
+        // set userId
+        legacyStorageAmplitude.setUserId(userId: legacyUserId)
+        XCTAssertEqual(legacyUserId, legacyStorageAmplitude.getUserId())
+
+        // track events to legacy storage
+        legacyStorageAmplitude.identify(identify: Identify().set(property: "user-prop", value: true))
+        legacyStorageAmplitude.track(event: BaseEvent(eventType: "Legacy Storage Event"))
+
+        guard let legacyEventFiles: [URL]? = legacyEventStorage.read(key: StorageKey.EVENTS) else { return }
+
+        var legacyEventsString = ""
+        legacyEventFiles?.forEach { file in
+            legacyEventsString = legacyEventStorage.getEventsString(eventBlock: file) ?? ""
+        }
+
+        XCTAssertEqual(legacyEventFiles?.count ?? 0, 1)
+
+        let amplitude = Amplitude(configuration: config)
+        let deviceId = amplitude.getDeviceId()
+        let userId = amplitude.getUserId()
+
+        guard let eventFiles: [URL]? = amplitude.storage.read(key: StorageKey.EVENTS) else { return }
+
+        var eventsString = ""
+        eventFiles?.forEach { file in
+            eventsString = legacyEventStorage.getEventsString(eventBlock: file) ?? ""
+        }
+
+        XCTAssertEqual(legacyDeviceId != nil, true)
+        XCTAssertEqual(deviceId != nil, true)
+        XCTAssertEqual(legacyDeviceId, deviceId)
+
+        XCTAssertEqual(legacyUserId, userId)
+
+        XCTAssertNotNil(legacyEventsString)
+
+        #if os(macOS)
+        // We don't want to transfer event data in non-sanboxed apps
+        XCTAssertFalse(amplitude.isSandboxEnabled())
+        XCTAssertEqual(eventFiles?.count ?? 0, 0)
+        #else
+        XCTAssertTrue(eventsString != "")
+        XCTAssertEqual(legacyEventsString, eventsString)
+        XCTAssertEqual(eventFiles?.count ?? 0, 1)
+        #endif
+
+        // clear storage
+        amplitude.storage.reset()
+        amplitude.identifyStorage.reset()
+        legacyStorageAmplitude.storage.reset()
+        legacyStorageAmplitude.identifyStorage.reset()
+    }
+
+    #if os(macOS)
+    func testMigrationToApiKeyAndInstanceNameStorageMacSandboxEnabled() throws {
+        let legacyUserId = "legacy-user-id"
+        let config = Configuration(
+            apiKey: "amp-mac-migration-api-key",
+            // don't transfer any events
+            flushQueueSize: 1000,
+            flushIntervalMillis: 99999,
+            logLevel: LogLevelEnum.DEBUG,
+            defaultTracking: DefaultTrackingOptions.NONE
+        )
+
+        // Create storages using instance name only
+        let legacyEventStorage = FakePersistentStorageAppSandboxEnabled(storagePrefix: "storage-\(config.getNormalizeInstanceName())")
+        let legacyIdentityStorage = FakePersistentStorageAppSandboxEnabled(storagePrefix: "identify-\(config.getNormalizeInstanceName())")
+
+        // Init Amplitude using legacy storage
+        let legacyStorageAmplitude = FakeAmplitudeWithNoInstNameOnlyMigration(configuration: Configuration(
+            apiKey: config.apiKey,
+            flushQueueSize: config.flushQueueSize,
+            flushIntervalMillis: config.flushIntervalMillis,
+            storageProvider: legacyEventStorage,
+            identifyStorageProvider: legacyIdentityStorage,
+            logLevel: config.logLevel,
+            defaultTracking: config.defaultTracking
+        ))
+
+        let legacyDeviceId = legacyStorageAmplitude.getDeviceId()
+
+        // set userId
+        legacyStorageAmplitude.setUserId(userId: legacyUserId)
+        XCTAssertEqual(legacyUserId, legacyStorageAmplitude.getUserId())
+
+        // track events to legacy storage
+        legacyStorageAmplitude.identify(identify: Identify().set(property: "user-prop", value: true))
+        legacyStorageAmplitude.track(event: BaseEvent(eventType: "Legacy Storage Event"))
+
+        guard let legacyEventFiles: [URL]? = legacyEventStorage.read(key: StorageKey.EVENTS) else { return }
+
+        var legacyEventsString = ""
+        legacyEventFiles?.forEach { file in
+            legacyEventsString = legacyEventStorage.getEventsString(eventBlock: file) ?? ""
+        }
+
+        XCTAssertEqual(legacyEventFiles?.count ?? 0, 1)
+
+        let amplitude = FakeAmplitudeWithSandboxEnabled(configuration: config)
+        let deviceId = amplitude.getDeviceId()
+        let userId = amplitude.getUserId()
+
+        guard let eventFiles: [URL]? = amplitude.storage.read(key: StorageKey.EVENTS) else { return }
+
+        var eventsString = ""
+        eventFiles?.forEach { file in
+            eventsString = legacyEventStorage.getEventsString(eventBlock: file) ?? ""
+        }
+
+        XCTAssertEqual(legacyDeviceId != nil, true)
+        XCTAssertEqual(deviceId != nil, true)
+        XCTAssertEqual(legacyDeviceId, deviceId)
+
+        XCTAssertEqual(legacyUserId, userId)
+
+        XCTAssertNotNil(legacyEventsString)
+
+        // Transfer event data in sandboxed apps
+        XCTAssertTrue(eventsString == "")
+        XCTAssertNotEqual(legacyEventsString, eventsString)
+        XCTAssertEqual(eventFiles?.count ?? 0, 0)
+
+        // clear storage
+        amplitude.storage.reset()
+        amplitude.identifyStorage.reset()
+        legacyStorageAmplitude.storage.reset()
+        legacyStorageAmplitude.identifyStorage.reset()
+    }
+    #endif
+
     func testInit_Offline() {
         XCTAssertEqual(Amplitude(configuration: configuration).configuration.offline, false)
     }