[new login flow] Migrate oAuth data from V1 to V2

When new login flow is adopted, we try to make the
experience seamless here: we restore the stored oAuth
data from V1 and use it in V2. Also in order to avoid
future incorrect/unexpected data migrations, we delete
the old oAuth data after migration.

Bug: 325079863
Test: added
Change-Id: I6fc3d176faa1f579043620cfd9533383ecbe4fd5

3 files changed, 321 insertions, 2 deletions

diff --git a/google-login-plugin/src/com/google/gct/login2/GoogleLoginService.kt b/google-login-plugin/src/com/google/gct/login2/GoogleLoginService.kt
index a62c815..e5f655e 100644
--- a/google-login-plugin/src/com/google/gct/login2/GoogleLoginService.kt
+++ b/google-login-plugin/src/com/google/gct/login2/GoogleLoginService.kt
@@ -22,6 +22,7 @@ import com.android.tools.idea.flags.StudioFlags
 import com.android.tools.idea.io.grpc.ClientInterceptor
 import com.google.api.client.auth.oauth2.TokenResponseException
 import com.google.gct.login.GoogleLogin
+import com.google.gct.login.GoogleLoginPrefs
 import com.google.gct.login2.common.OAuthDataStore
 import com.google.gct.login2.common.UiFacade
 import com.google.gct.login2.settings.GoogleLoginApplicationSettings
@@ -141,6 +142,7 @@ internal class GoogleLoginServiceImpl(
   private val oAuthServer: OAuthServer = GoogleOAuthServer,
   private val loginApplicationSettings: GoogleLoginApplicationSettings =
     service<GoogleLoginApplicationSettings>(),
+  private val oldLoginApplicationSettings: GoogleLoginPrefs = GoogleLoginPrefs.INSTANCE,
 ) : GoogleLoginService, Disposable {
 
   private val logger: Logger = thisLogger()
@@ -414,9 +416,13 @@ internal class GoogleLoginServiceImpl(
 
   @Slow
   private fun initializeUsers() {
-    val activeUserString = loginApplicationSettings.loadActiveUser() ?: return
+    val activeUserString =
+      loginApplicationSettings.loadActiveUser()
+        ?: oldLoginApplicationSettings.getActiveUser()
+        ?: return
     val user2OAuthDataStoreMap: Map<String, Map<LoginFeature, OAuthDataStore>> =
-      loginApplicationSettings.loadOAuthData()
+      loginApplicationSettings.loadOAuthData().takeUnless { it.isEmpty() }
+        ?: restoreOAuthDataFromV1(oldLoginApplicationSettings)
 
     user2OAuthDataStoreMap.forEach { (user, feature2OAuthData) ->
       feature2OAuthData
diff --git a/google-login-plugin/src/com/google/gct/login2/V2MigrationUtils.kt b/google-login-plugin/src/com/google/gct/login2/V2MigrationUtils.kt
new file mode 100644
index 0000000..3c8d234
--- /dev/null
+++ b/google-login-plugin/src/com/google/gct/login2/V2MigrationUtils.kt
@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * 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.
+ */
+package com.google.gct.login2
+
+import com.google.gct.login.GoogleLoginPrefs
+import com.google.gct.login.common.OAuthData
+import com.google.gct.login2.common.OAuthDataStore
+
+/**
+ * Returns previously saved OAuth data from the legacy login flow.
+ *
+ * This ensures a smooth migration experience from the legacy login to our new login.
+ */
+fun restoreOAuthDataFromV1(
+  loginPrefs: GoogleLoginPrefs
+): Map<String, Map<LoginFeature, OAuthDataStore>> {
+  val restored = loginPrefs.extractOAuthData()
+
+  // We simply delete the old credentials at this point to prevent future incorrect data migrations.
+  loginPrefs.removeAllUsers()
+  return restored
+}
+
+private fun GoogleLoginPrefs.extractOAuthData(): Map<String, Map<LoginFeature, OAuthDataStore>> {
+  val allUsers = getStoredUsers()
+
+  return allUsers
+    .mapNotNull { user ->
+      val oAuthDataStore = loadOAuthData(user).toOAuthDataStore() ?: return@mapNotNull null
+      val map =
+        oAuthDataStore.allowedFeatures.associateWith {
+          oAuthDataStore.copy(allowedFeatures = setOf(it))
+        }
+
+      user to map
+    }
+    .toMap()
+}
+
+private fun OAuthData.toOAuthDataStore(): OAuthDataStore? {
+  storedEmail ?: return null
+
+  return OAuthDataStore(
+    userEmail = storedEmail,
+    refreshToken = refreshToken,
+    allowedFeatures = extractFeatures(storedScopes),
+  )
+}
+
+private fun extractFeatures(scopeSet: Set<String>): Set<LoginFeature> {
+  return LoginFeature.EP_NAME.extensionList
+    .filter { loginFeature -> scopeSet.containsAll(loginFeature.oAuthScopes) }
+    .toSet()
+}
diff --git a/google-login-plugin/testSrc/com/google/gct/login2/V2MigrationUtilsKtTest.kt b/google-login-plugin/testSrc/com/google/gct/login2/V2MigrationUtilsKtTest.kt
new file mode 100644
index 0000000..abefd99
--- /dev/null
+++ b/google-login-plugin/testSrc/com/google/gct/login2/V2MigrationUtilsKtTest.kt
@@ -0,0 +1,246 @@
+package com.google.gct.login2
+
+import com.android.flags.junit.FlagRule
+import com.android.testutils.waitForCondition
+import com.android.tools.idea.flags.StudioFlags
+import com.google.common.truth.Truth.assertThat
+import com.google.gct.login.GoogleLoginPrefsRule
+import com.google.gct.login.common.OAuthData
+import com.google.gct.login2.common.OAuthDataStore
+import com.google.gct.login2.settings.GoogleLoginApplicationSettings
+import com.intellij.openapi.util.Disposer
+import com.intellij.testFramework.ApplicationRule
+import com.intellij.testFramework.DisposableRule
+import kotlin.time.Duration.Companion.seconds
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.RuleChain
+
+class V2MigrationUtilsKtTest {
+  private val applicationRule = ApplicationRule()
+  private val loginFeatureRule = LoginFeatureRule()
+  private val disposableRule = DisposableRule()
+  private val prefsRule = GoogleLoginPrefsRule()
+
+  private lateinit var settings: GoogleLoginApplicationSettings
+
+  @get:Rule
+  val rules =
+    RuleChain.outerRule(FlagRule(StudioFlags.ENABLE_SETTINGS_ACCOUNT_UI, true))
+      .around(applicationRule)
+      .around(prefsRule)
+      .around(loginFeatureRule)
+      .around(disposableRule)
+
+  @Before
+  fun setUp() {
+    settings = GoogleLoginApplicationSettings()
+  }
+
+  @Test
+  fun `load from v2`() {
+    fun setUpOldOAuthStorage() {
+      prefsRule.prefs.saveActiveUser("foo@bar.com")
+      prefsRule.prefs.saveOAuthData(
+        OAuthData(
+          "test-access-token",
+          "test-refresh-token",
+          "foo@bar.com",
+          loginFeatureRule.FEATURE1.oAuthScopes.toSortedSet(),
+          100,
+        )
+      )
+    }
+
+    fun setUpNewOAuthStorage() {
+      settings.saveOAuthData(
+        OAuthDataStore(
+          userEmail = "test_user@gmail.com",
+          refreshToken = "refresh_token",
+          setOf(loginFeatureRule.FEATURE1),
+        )
+      )
+      settings.saveActiveUser(TEST_EMAIL)
+
+      assertThat(settings.loadOAuthData())
+        .containsExactly(
+          TEST_EMAIL,
+          mapOf(
+            loginFeatureRule.FEATURE1 to
+              OAuthDataStore(
+                userEmail = TEST_EMAIL,
+                refreshToken = "refresh_token",
+                setOf(loginFeatureRule.FEATURE1),
+              )
+          ),
+        )
+    }
+
+    setUpOldOAuthStorage()
+    setUpNewOAuthStorage()
+
+    val loginService = createGoogleLoginService()
+    with(loginService.allUsersFlow.value) {
+      assertThat(size).isEqualTo(1)
+      assertThat(entries.single().key).isEqualTo(TEST_EMAIL)
+      assertThat(entries.single().value.getAllowedFeatures())
+        .containsExactlyElementsIn(setOf(loginFeatureRule.FEATURE1))
+    }
+  }
+
+  @Test
+  fun `load from v1, single user, single feature`() {
+    fun setUpOldOAuthStorage() {
+      prefsRule.prefs.saveActiveUser("foo@bar.com")
+      prefsRule.prefs.saveOAuthData(
+        OAuthData(
+          "test-access-token",
+          "test-refresh-token",
+          "foo@bar.com",
+          loginFeatureRule.FEATURE1.oAuthScopes.toSortedSet(),
+          100,
+        )
+      )
+    }
+
+    setUpOldOAuthStorage()
+    val loginService = createGoogleLoginService()
+
+    val (email, user) = loginService.allUsersFlow.value.entries.single()
+    assertThat(email).isEqualTo("foo@bar.com")
+    assertThat(loginService.activeUserFlow.value).isEqualTo(user)
+    assertThat(user.email).isEqualTo("foo@bar.com")
+    assertThat(user.feature2LoginStateMap.keys).containsAllIn(setOf(loginFeatureRule.FEATURE1))
+    assertThat(user.feature2LoginStateMap[loginFeatureRule.FEATURE1]!!.getAllowedFeatures())
+      .isEqualTo(setOf(loginFeatureRule.FEATURE1))
+  }
+
+  @Test
+  fun `old data is deleted after v1 to v2 migration`() {
+    fun setUpOldOAuthStorage() {
+      prefsRule.prefs.saveActiveUser("foo@bar.com")
+      prefsRule.prefs.saveOAuthData(
+        OAuthData(
+          "test-access-token",
+          "test-refresh-token",
+          "foo@bar.com",
+          (loginFeatureRule.FEATURE1.oAuthScopes +
+              loginFeatureRule.FEATURE2.oAuthScopes +
+              loginFeatureRule.ENFORCED.oAuthScopes)
+            .toSortedSet(),
+          100,
+        )
+      )
+
+      assertThat(prefsRule.prefs.getActiveUser()).isEqualTo("foo@bar.com")
+      assertThat(prefsRule.prefs.getStoredUsers()).containsExactly("foo@bar.com")
+    }
+
+    setUpOldOAuthStorage()
+    val loginService = createGoogleLoginService()
+
+    val (email, user) = loginService.allUsersFlow.value.entries.single()
+    assertThat(email).isEqualTo("foo@bar.com")
+    assertThat(loginService.activeUserFlow.value).isEqualTo(user)
+    assertThat(user.email).isEqualTo("foo@bar.com")
+    assertThat(user.feature2LoginStateMap.keys)
+      .containsAllIn(
+        setOf(loginFeatureRule.FEATURE1, loginFeatureRule.FEATURE2, loginFeatureRule.ENFORCED)
+      )
+
+    // check old data is deleted
+    assertThat(prefsRule.prefs.getActiveUser()).isNull()
+    assertThat(prefsRule.prefs.getStoredUsers()).isEmpty()
+  }
+
+  @Test
+  fun `load from v1, single user, multiple features`() {
+    fun setUpOldOAuthStorage() {
+      prefsRule.prefs.saveActiveUser("foo@bar.com")
+      prefsRule.prefs.saveOAuthData(
+        OAuthData(
+          "test-access-token",
+          "test-refresh-token",
+          "foo@bar.com",
+          (loginFeatureRule.FEATURE1.oAuthScopes + loginFeatureRule.ENFORCED.oAuthScopes)
+            .toSortedSet(),
+          100,
+        )
+      )
+    }
+
+    setUpOldOAuthStorage()
+    val loginService = createGoogleLoginService()
+
+    val (email, user) = loginService.allUsersFlow.value.entries.single()
+    assertThat(email).isEqualTo("foo@bar.com")
+    assertThat(loginService.activeUserFlow.value).isEqualTo(user)
+    assertThat(user.email).isEqualTo("foo@bar.com")
+    assertThat(user.feature2LoginStateMap.keys)
+      .containsAllIn(setOf(loginFeatureRule.FEATURE1, loginFeatureRule.ENFORCED))
+    assertThat(user.feature2LoginStateMap[loginFeatureRule.FEATURE1]!!.getAllowedFeatures())
+      .isEqualTo(setOf(loginFeatureRule.FEATURE1, loginFeatureRule.ENFORCED))
+    assertThat(user.feature2LoginStateMap[loginFeatureRule.ENFORCED]!!.getAllowedFeatures())
+      .isEqualTo(setOf(loginFeatureRule.FEATURE1, loginFeatureRule.ENFORCED))
+  }
+
+  @Test
+  fun `load from v1, multiple users`() {
+    fun setUpOldOAuthStorage() {
+      prefsRule.prefs.saveActiveUser("foo@bar.com")
+      prefsRule.prefs.saveOAuthData(
+        OAuthData(
+          "test-access-token1",
+          "test-refresh-token1",
+          "foo@bar.com",
+          loginFeatureRule.FEATURE1.oAuthScopes.toSortedSet(),
+          100,
+        )
+      )
+      prefsRule.prefs.saveOAuthData(
+        OAuthData(
+          "test-access-token2",
+          "test-refresh-token2",
+          TEST_EMAIL,
+          loginFeatureRule.FEATURE2.oAuthScopes.toSortedSet(),
+          100,
+        )
+      )
+    }
+
+    setUpOldOAuthStorage()
+    val loginService = createGoogleLoginService()
+    with(loginService.allUsersFlow.value) {
+      assertThat(size).isEqualTo(2)
+      assertThat(keys).containsExactly("foo@bar.com", TEST_EMAIL)
+
+      val user1 = loginService.allUsersFlow.value["foo@bar.com"]!!
+      assertThat(loginService.activeUserFlow.value).isEqualTo(user1)
+      assertThat(user1.email).isEqualTo("foo@bar.com")
+      assertThat(user1.feature2LoginStateMap.keys).containsAllIn(setOf(loginFeatureRule.FEATURE1))
+      assertThat(user1.feature2LoginStateMap[loginFeatureRule.FEATURE1]!!.getAllowedFeatures())
+        .isEqualTo(setOf(loginFeatureRule.FEATURE1))
+
+      val user2 = loginService.allUsersFlow.value[TEST_EMAIL]!!
+      assertThat(user2.email).isEqualTo(TEST_EMAIL)
+      assertThat(user2.feature2LoginStateMap.keys).containsAllIn(setOf(loginFeatureRule.FEATURE2))
+      assertThat(user2.feature2LoginStateMap[loginFeatureRule.FEATURE2]!!.getAllowedFeatures())
+        .isEqualTo(setOf(loginFeatureRule.FEATURE2))
+    }
+  }
+
+  private fun createGoogleLoginService(): GoogleLoginServiceImpl {
+    return GoogleLoginServiceImpl(
+        FakeGoogleAccountClient(),
+        FakeUiFacade(),
+        FakeOAuthServer(),
+        settings,
+        prefsRule.prefs,
+      )
+      .apply {
+        Disposer.register(disposableRule.disposable, this)
+        waitForCondition(1.seconds) { isInitialized }
+      }
+  }
+}