[Role Logic Move] Add migration logic for isRoleFallbackEnabled

Migrate the isRoleFallbackEnabled preference to system server as a part
of the role logic move project. We only do the migration if it's V+. On
the system server side it will check the existing preference in
SharedPreference and migrate for all roles for the user.

API-Coverage-Bug: 314281553
Bug: 302563864
Test: RolePersistenceTest and RoleManagerTest
Change-Id: Ib24ebb8211359a652e52007fecf79a2390575e9c

14 files changed, 289 insertions, 48 deletions

diff --git a/PermissionController/role-controller/java/com/android/role/controller/service/RoleControllerServiceImpl.java b/PermissionController/role-controller/java/com/android/role/controller/service/RoleControllerServiceImpl.java
index d3c099d5e..2a6010c4d 100644
--- a/PermissionController/role-controller/java/com/android/role/controller/service/RoleControllerServiceImpl.java
+++ b/PermissionController/role-controller/java/com/android/role/controller/service/RoleControllerServiceImpl.java
@@ -32,6 +32,7 @@ import androidx.annotation.WorkerThread;
 import com.android.role.controller.model.Role;
 import com.android.role.controller.model.Roles;
 import com.android.role.controller.util.CollectionUtils;
+import com.android.role.controller.util.LegacyRoleFallbackEnabledUtils;
 import com.android.role.controller.util.PackageUtils;
 import com.android.role.controller.util.UserUtils;
 
@@ -468,6 +469,13 @@ public class RoleControllerServiceImpl extends RoleControllerService {
         return role.isVisibleAsUser(mUser, mContext);
     }
 
+    @Override
+    @NonNull
+    public List<String> onGetLegacyFallbackDisabledRoles() {
+        return LegacyRoleFallbackEnabledUtils.getFallbackDisabledRoles(mUser, mContext);
+    }
+
+
     private static boolean checkFlags(int flags, int allowedFlags) {
         if ((flags & allowedFlags) != flags) {
             Log.e(LOG_TAG, "flags is invalid, flags: 0x" + Integer.toHexString(flags)
diff --git a/PermissionController/role-controller/java/com/android/role/controller/util/LegacyRoleFallbackEnabledUtils.java b/PermissionController/role-controller/java/com/android/role/controller/util/LegacyRoleFallbackEnabledUtils.java
new file mode 100644
index 000000000..5be10a26a
--- /dev/null
+++ b/PermissionController/role-controller/java/com/android/role/controller/util/LegacyRoleFallbackEnabledUtils.java
@@ -0,0 +1,107 @@
+/*
+ * Copyright (C) 2023 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.android.role.controller.util;
+
+import android.annotation.TargetApi;
+import android.content.Context;
+import android.content.SharedPreferences;
+import android.content.pm.PackageManager;
+import android.os.Build;
+import android.os.UserHandle;
+
+import androidx.annotation.NonNull;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class LegacyRoleFallbackEnabledUtils {
+    /**
+     * Name of generic shared preferences file.
+     */
+    private static final String PREFERENCES_FILE = "preferences";
+
+    /**
+     * Key in the generic shared preferences that stores if the user manually selected the "none"
+     * role holder for a role.
+     */
+    private static final String IS_NONE_ROLE_HOLDER_SELECTED_KEY = "is_none_role_holder_selected:";
+
+    /**
+     * Get a device protected storage based shared preferences. Avoid storing sensitive data in it.
+     *
+     * @param context the context to get the shared preferences
+     * @return a device protected storage based shared preferences
+     */
+    @NonNull
+    @TargetApi(Build.VERSION_CODES.VANILLA_ICE_CREAM)
+    private static SharedPreferences getSharedPreferences(@NonNull UserHandle user,
+            @NonNull Context context) {
+        String packageName = context.getPackageManager().getPermissionControllerPackageName();
+        try {
+            context = context.createPackageContextAsUser(packageName, 0, user);
+        } catch (PackageManager.NameNotFoundException e) {
+            throw new RuntimeException(e);
+        }
+        if (!context.isDeviceProtectedStorage()) {
+            context = context.createDeviceProtectedStorageContext();
+        }
+        return context.getSharedPreferences(PREFERENCES_FILE, Context.MODE_PRIVATE);
+    }
+
+    /**
+     * Get all role names with fallback disabled, which means their "none" are set to true.
+     *
+     * @return A list of role names with fallback disabled.
+     */
+    public static List<String> getFallbackDisabledRoles(@NonNull UserHandle user,
+            @NonNull Context context) {
+        List<String> fallbackDisabledRoles = new ArrayList<>();
+        SharedPreferences sharedPreferences = getSharedPreferences(user, context);
+        for (String key : sharedPreferences.getAll().keySet()) {
+            if (key.startsWith(IS_NONE_ROLE_HOLDER_SELECTED_KEY)
+                    && sharedPreferences.getBoolean(key, false)) {
+                String roleName = key.substring(IS_NONE_ROLE_HOLDER_SELECTED_KEY.length());
+                fallbackDisabledRoles.add(roleName);
+            }
+        }
+        return fallbackDisabledRoles;
+    }
+
+    /**
+     * Check whether the role has the fallback holder enabled.
+     *
+     * @return whether the "none" role holder is not selected
+     */
+    public static boolean isRoleFallbackEnabledAsUser(@NonNull String roleName,
+            @NonNull UserHandle user, @NonNull Context context) {
+        return !getSharedPreferences(user, context)
+                .getBoolean(IS_NONE_ROLE_HOLDER_SELECTED_KEY + roleName, false);
+    }
+
+    /**
+     * Set whether the role has fallback holder enabled.
+     */
+    public static void setRoleFallbackEnabledAsUser(@NonNull String roleName,
+            boolean fallbackEnabled, @NonNull UserHandle user, @NonNull Context context) {
+        String key = IS_NONE_ROLE_HOLDER_SELECTED_KEY + roleName;
+        if (fallbackEnabled) {
+            getSharedPreferences(user, context).edit().remove(key).apply();
+        } else {
+            getSharedPreferences(user, context).edit().putBoolean(key, true).apply();
+        }
+    }
+}
diff --git a/PermissionController/role-controller/java/com/android/role/controller/util/RoleManagerCompat.java b/PermissionController/role-controller/java/com/android/role/controller/util/RoleManagerCompat.java
index 8a6fd579c..ec63528d1 100644
--- a/PermissionController/role-controller/java/com/android/role/controller/util/RoleManagerCompat.java
+++ b/PermissionController/role-controller/java/com/android/role/controller/util/RoleManagerCompat.java
@@ -18,8 +18,8 @@ package com.android.role.controller.util;
 
 import android.app.role.RoleManager;
 import android.content.Context;
-import android.content.SharedPreferences;
 import android.os.UserHandle;
+import android.permission.flags.Flags;
 
 import androidx.annotation.NonNull;
 
@@ -31,16 +31,7 @@ import com.android.role.controller.model.Role;
  */
 public class RoleManagerCompat {
 
-    /**
-     * Key in the generic shared preferences that stores if the user manually selected the "none"
-     * role holder for a role.
-     */
-    private static final String IS_NONE_ROLE_HOLDER_SELECTED_KEY = "is_none_role_holder_selected:";
 
-    /**
-     * Name of generic shared preferences file.
-     */
-    private static final String PREFERENCES_FILE = "preferences";
 
     private RoleManagerCompat() {}
 
@@ -56,47 +47,34 @@ public class RoleManagerCompat {
     }
 
     /**
-     * Get a device protected storage based shared preferences. Avoid storing sensitive data in it.
-     *
-     * @param context the context to get the shared preferences
-     * @return a device protected storage based shared preferences
-     */
-    @NonNull
-    private static SharedPreferences getDeviceProtectedSharedPreferences(@NonNull Context context) {
-        if (!context.isDeviceProtectedStorage()) {
-            context = context.createDeviceProtectedStorageContext();
-        }
-        return context.getSharedPreferences(PREFERENCES_FILE, Context.MODE_PRIVATE);
-    }
-
-    /**
      * Check whether the role has the fallback holder enabled.
      *
      * @return whether the "none" role holder is not selected
      */
     public static boolean isRoleFallbackEnabledAsUser(@NonNull Role role, @NonNull UserHandle user,
             @NonNull Context context) {
-        Context userContext = UserUtils.getUserContext(context, user);
-        boolean isNoneHolderSelected = getDeviceProtectedSharedPreferences(userContext)
-                .getBoolean(IS_NONE_ROLE_HOLDER_SELECTED_KEY + role.getName(), false);
-        return !isNoneHolderSelected;
+        if (SdkLevel.isAtLeastV() && Flags.systemServerRoleControllerEnabled()) {
+            Context userContext = UserUtils.getUserContext(context, user);
+            RoleManager userRoleManager = userContext.getSystemService(RoleManager.class);
+            return userRoleManager.isRoleFallbackEnabled(role.getName());
+        } else {
+            return LegacyRoleFallbackEnabledUtils.isRoleFallbackEnabledAsUser(role.getName(), user,
+                    context);
+        }
     }
 
     /**
      * Set whether the role has fallback holder enabled.
-     *
      */
     public static void setRoleFallbackEnabledAsUser(@NonNull Role role,
             boolean fallbackEnabled, @NonNull UserHandle user, @NonNull Context context) {
-        Context userContext = UserUtils.getUserContext(context, user);
-        if (fallbackEnabled) {
-            getDeviceProtectedSharedPreferences(userContext).edit()
-                    .remove(IS_NONE_ROLE_HOLDER_SELECTED_KEY + role.getName())
-                    .apply();
+        if (SdkLevel.isAtLeastV() && Flags.systemServerRoleControllerEnabled()) {
+            Context userContext = UserUtils.getUserContext(context, user);
+            RoleManager userRoleManager = userContext.getSystemService(RoleManager.class);
+            userRoleManager.setRoleFallbackEnabled(role.getName(), fallbackEnabled);
         } else {
-            getDeviceProtectedSharedPreferences(userContext).edit()
-                    .putBoolean(IS_NONE_ROLE_HOLDER_SELECTED_KEY + role.getName(), true)
-                    .apply();
+            LegacyRoleFallbackEnabledUtils.setRoleFallbackEnabledAsUser(role.getName(),
+                    fallbackEnabled, user, context);
         }
     }
 }
diff --git a/framework-s/api/system-current.txt b/framework-s/api/system-current.txt
index 1111470d2..b5c379937 100644
--- a/framework-s/api/system-current.txt
+++ b/framework-s/api/system-current.txt
@@ -10,6 +10,7 @@ package android.app.role {
     method @Deprecated @WorkerThread public abstract boolean onAddRoleHolder(@NonNull String, @NonNull String, int);
     method @Deprecated @Nullable public final android.os.IBinder onBind(@Nullable android.content.Intent);
     method @Deprecated @WorkerThread public abstract boolean onClearRoleHolders(@NonNull String, int);
+    method @Deprecated @FlaggedApi("android.permission.flags.system_server_role_controller_enabled") @NonNull public java.util.List<java.lang.String> onGetLegacyFallbackDisabledRoles();
     method @Deprecated @WorkerThread public abstract boolean onGrantDefaultRoles();
     method @Deprecated public abstract boolean onIsApplicationQualifiedForRole(@NonNull String, @NonNull String);
     method @Deprecated public boolean onIsApplicationVisibleForRole(@NonNull String, @NonNull String);
diff --git a/framework-s/java/android/app/role/IRoleController.aidl b/framework-s/java/android/app/role/IRoleController.aidl
index 8a43d7fa9..948915f8d 100644
--- a/framework-s/java/android/app/role/IRoleController.aidl
+++ b/framework-s/java/android/app/role/IRoleController.aidl
@@ -40,4 +40,6 @@ oneway interface IRoleController {
             in RemoteCallback callback);
 
     void isRoleVisible(in String roleName, in RemoteCallback callback);
+
+    void getLegacyFallbackDisabledRoles(in RemoteCallback callback);
 }
diff --git a/framework-s/java/android/app/role/RoleControllerManager.java b/framework-s/java/android/app/role/RoleControllerManager.java
index 3b990b315..57da2ccd0 100644
--- a/framework-s/java/android/app/role/RoleControllerManager.java
+++ b/framework-s/java/android/app/role/RoleControllerManager.java
@@ -37,6 +37,7 @@ import com.android.internal.annotations.GuardedBy;
 import com.android.internal.infra.AndroidFuture;
 import com.android.internal.infra.ServiceConnector;
 
+import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
@@ -48,6 +49,12 @@ import java.util.function.Consumer;
  */
 public class RoleControllerManager {
 
+    /**
+     * Bundle key for getting legacy fallback disabled roles
+     */
+    public static final String KEY_LEGACY_FALLBACK_DISABLED_ROLES =
+            "LEGACY_FALLBACK_DISABLED_ROLES";
+
     private static final String LOG_TAG = RoleControllerManager.class.getSimpleName();
 
     private static final long REQUEST_TIMEOUT_MILLIS = 15 * 1000;
@@ -187,8 +194,7 @@ public class RoleControllerManager {
             @RoleManager.ManageHoldersFlags int flags, @NonNull RemoteCallback callback) {
         AndroidFuture<Bundle> operation = mRemoteService.postAsync(service -> {
             AndroidFuture<Bundle> future = new AndroidFuture<>();
-            service.onClearRoleHolders(roleName, flags,
-                    new RemoteCallback(future::complete));
+            service.onClearRoleHolders(roleName, flags, new RemoteCallback(future::complete));
             return future;
         });
         propagateCallback(operation, "onClearRoleHolders", callback);
@@ -227,6 +233,35 @@ public class RoleControllerManager {
         propagateCallback(operation, "isRoleVisible", executor, callback);
     }
 
+    /**
+     * @see RoleControllerService#onGrantDefaultRoles()
+     *
+     * @hide
+     */
+    public void getLegacyFallbackDisabledRoles(@NonNull @CallbackExecutor Executor executor,
+            @NonNull Consumer<List<String>> callback) {
+        mRemoteService.postAsync(service -> {
+            AndroidFuture<Bundle> future = new AndroidFuture<>();
+            service.getLegacyFallbackDisabledRoles(new RemoteCallback(future::complete));
+            return future;
+        }).orTimeout(REQUEST_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS)
+                .whenComplete((res, err) -> executor.execute(() -> {
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        if (err != null) {
+                            Log.e(LOG_TAG, "Error calling getLegacyFallbackDisabledRoles()",
+                                    err);
+                            callback.accept(null);
+                        } else {
+                            callback.accept(res.getStringArrayList(
+                                    KEY_LEGACY_FALLBACK_DISABLED_ROLES));
+                        }
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
+                }));
+    }
+
     private void propagateCallback(AndroidFuture<Bundle> operation, String opName,
             @CallbackExecutor @NonNull Executor executor,
             Consumer<Boolean> destination) {
diff --git a/framework-s/java/android/app/role/RoleControllerService.java b/framework-s/java/android/app/role/RoleControllerService.java
index cf7872913..60a13f7ba 100644
--- a/framework-s/java/android/app/role/RoleControllerService.java
+++ b/framework-s/java/android/app/role/RoleControllerService.java
@@ -17,6 +17,7 @@
 package android.app.role;
 
 import android.Manifest;
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
@@ -31,9 +32,12 @@ import android.os.IBinder;
 import android.os.Process;
 import android.os.RemoteCallback;
 import android.os.UserHandle;
+import android.permission.flags.Flags;
 
 import com.android.internal.util.Preconditions;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.Executor;
 
@@ -175,6 +179,19 @@ public abstract class RoleControllerService extends Service {
                 boolean visible = onIsRoleVisible(roleName);
                 callback.sendResult(visible ? Bundle.EMPTY : null);
             }
+
+            @Override
+            public void getLegacyFallbackDisabledRoles(RemoteCallback callback) {
+                enforceCallerSystemUid("getLegacyFallbackDisabledRoles");
+
+                Objects.requireNonNull(callback, "callback cannot be null");
+
+                List<String> legacyFallbackDisabledRoles = onGetLegacyFallbackDisabledRoles();
+                Bundle result = new Bundle();
+                result.putStringArrayList(RoleControllerManager.KEY_LEGACY_FALLBACK_DISABLED_ROLES,
+                        new ArrayList<>(legacyFallbackDisabledRoles));
+                callback.sendResult(result);
+            }
         };
     }
 
@@ -301,4 +318,15 @@ public abstract class RoleControllerService extends Service {
      * @return whether the role should be visible to user
      */
     public abstract boolean onIsRoleVisible(@NonNull String roleName);
+
+    /**
+     * Get the legacy fallback disabled state.
+     *
+     * @return A list of role names with disabled fallback state.
+     */
+    @FlaggedApi(Flags.FLAG_SYSTEM_SERVER_ROLE_CONTROLLER_ENABLED)
+    @NonNull
+    public List<String> onGetLegacyFallbackDisabledRoles() {
+        throw new UnsupportedOperationException();
+    }
 }
diff --git a/service/java/com/android/role/LocalRoleController.java b/service/java/com/android/role/LocalRoleController.java
index 6548aaa03..03508d1cb 100644
--- a/service/java/com/android/role/LocalRoleController.java
+++ b/service/java/com/android/role/LocalRoleController.java
@@ -16,6 +16,7 @@
 
 package com.android.role;
 
+import android.annotation.CallbackExecutor;
 import android.app.role.RoleControllerService;
 import android.app.role.RoleManager;
 import android.content.Context;
@@ -29,6 +30,7 @@ import androidx.annotation.NonNull;
 
 import com.android.role.controller.service.RoleControllerServiceImpl;
 
+import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
@@ -93,4 +95,10 @@ public class LocalRoleController implements RoleController {
             @NonNull String packageName) {
         return mService.onIsApplicationVisibleForRole(roleName, packageName);
     }
+
+    @Override
+    public void getLegacyFallbackDisabledRoles(@NonNull @CallbackExecutor Executor executor,
+            @NonNull Consumer<List<String>> callback) {
+        throw new UnsupportedOperationException();
+    }
 }
diff --git a/service/java/com/android/role/RemoteRoleController.java b/service/java/com/android/role/RemoteRoleController.java
index cbf164d6d..0920dd56a 100644
--- a/service/java/com/android/role/RemoteRoleController.java
+++ b/service/java/com/android/role/RemoteRoleController.java
@@ -16,6 +16,7 @@
 
 package com.android.role;
 
+import android.annotation.CallbackExecutor;
 import android.app.role.RoleControllerManager;
 import android.app.role.RoleManager;
 import android.content.Context;
@@ -26,6 +27,7 @@ import androidx.annotation.NonNull;
 
 import com.android.permission.util.ForegroundThread;
 
+import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
@@ -73,4 +75,10 @@ public class RemoteRoleController implements RoleController {
             @NonNull String packageName) {
         throw new UnsupportedOperationException();
     }
+
+    @Override
+    public void getLegacyFallbackDisabledRoles(@NonNull @CallbackExecutor Executor executor,
+            @NonNull Consumer<List<String>> callback) {
+        mRoleControllerManager.getLegacyFallbackDisabledRoles(executor, callback);
+    }
 }
diff --git a/service/java/com/android/role/RoleController.java b/service/java/com/android/role/RoleController.java
index 8a7ef8646..6ebbdfd45 100644
--- a/service/java/com/android/role/RoleController.java
+++ b/service/java/com/android/role/RoleController.java
@@ -21,6 +21,7 @@ import android.annotation.NonNull;
 import android.app.role.RoleManager;
 import android.os.RemoteCallback;
 
+import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
@@ -58,4 +59,10 @@ public interface RoleController {
      * @see android.app.role.RoleControllerManager#isApplicationVisibleForRole
      */
     boolean isApplicationVisibleForRole(@NonNull String roleName, @NonNull String packageName);
+
+    /**
+     * @see android.app.role.RoleControllerManager#getLegacyFallbackDisabledRoles
+     */
+    void getLegacyFallbackDisabledRoles(@NonNull @CallbackExecutor Executor executor,
+            @NonNull Consumer<List<String>> callback);
 }
diff --git a/service/java/com/android/role/RoleService.java b/service/java/com/android/role/RoleService.java
index c1a3ea4d9..2524627eb 100644
--- a/service/java/com/android/role/RoleService.java
+++ b/service/java/com/android/role/RoleService.java
@@ -252,9 +252,40 @@ public class RoleService extends SystemService implements RoleUserState.Callback
 
     @Override
     public void onUserStarting(@NonNull TargetUser user) {
+        if (SdkLevel.isAtLeastV() && Flags.systemServerRoleControllerEnabled()) {
+            upgradeLegacyFallbackEnabledRolesIfNeeded(user);
+        }
+
         maybeGrantDefaultRolesSync(user.getUserHandle().getIdentifier());
     }
 
+    private void upgradeLegacyFallbackEnabledRolesIfNeeded(@NonNull TargetUser user) {
+        int userId = user.getUserHandle().getIdentifier();
+        RoleUserState userState = getOrCreateUserState(userId);
+        if (!userState.isVersionUpgradeNeeded()) {
+            return;
+        }
+        List<String> legacyFallbackDisabledRoles = getLegacyFallbackDisabledRolesSync(userId);
+        if (legacyFallbackDisabledRoles == null) {
+            return;
+        }
+        userState.upgradeVersion(legacyFallbackDisabledRoles);
+    }
+
+    @MainThread
+    private List<String> getLegacyFallbackDisabledRolesSync(@UserIdInt int userId) {
+        AndroidFuture<List<String>> future = new AndroidFuture<>();
+        RoleController controller = new RemoteRoleController(UserHandle.of(userId), getContext());
+        controller.getLegacyFallbackDisabledRoles(ForegroundThread.getExecutor(), future::complete);
+        try {
+            return future.get(30, TimeUnit.SECONDS);
+        } catch (InterruptedException | ExecutionException | TimeoutException e) {
+            Log.e(LOG_TAG, "Failed to get the legacy role fallback disabled state for user "
+                    + userId, e);
+            return null;
+        }
+    }
+
     @MainThread
     private void maybeGrantDefaultRolesSync(@UserIdInt int userId) {
         AndroidFuture<Void> future = maybeGrantDefaultRolesInternal(userId);
@@ -493,8 +524,7 @@ public class RoleService extends SystemService implements RoleUserState.Callback
             Preconditions.checkStringNotEmpty(packageName, "packageName cannot be null or empty");
             Objects.requireNonNull(callback, "callback cannot be null");
 
-            getOrCreateController(userId).onAddRoleHolder(roleName, packageName, flags,
-                    callback);
+            getOrCreateController(userId).onAddRoleHolder(roleName, packageName, flags, callback);
         }
 
         @Override
diff --git a/service/java/com/android/role/RoleUserState.java b/service/java/com/android/role/RoleUserState.java
index 201633898..39c1e8c3d 100644
--- a/service/java/com/android/role/RoleUserState.java
+++ b/service/java/com/android/role/RoleUserState.java
@@ -89,7 +89,7 @@ class RoleUserState {
     private ArrayMap<String, ArraySet<String>> mRoles = new ArrayMap<>();
 
     /**
-     *
+     * Role names of the roles with fallback enabled.
      */
     @GuardedBy("mLock")
     @NonNull
@@ -150,6 +150,15 @@ class RoleUserState {
     }
 
     /**
+     * Checks the version and returns whether a version upgrade is needed.
+     */
+    public boolean isVersionUpgradeNeeded() {
+        synchronized (mLock) {
+            return mVersion < VERSION_FALLBACK_STATE_MIGRATED;
+        }
+    }
+
+    /**
      * Get the hash representing the state of packages during the last time initial grants was run.
      *
      * @return the hash representing the state of packages
@@ -228,6 +237,23 @@ class RoleUserState {
     }
 
     /**
+     * Upgrade this user state to the latest version if needed.
+     */
+    public void upgradeVersion(@NonNull List<String> legacyFallbackDisabledRoles) {
+        synchronized (mLock) {
+            if (mVersion < VERSION_FALLBACK_STATE_MIGRATED) {
+                int legacyFallbackDisabledRolesSize = legacyFallbackDisabledRoles.size();
+                for (int i = 0; i < legacyFallbackDisabledRolesSize; i++) {
+                    String roleName = legacyFallbackDisabledRoles.get(i);
+                    mFallbackEnabledRoles.remove(roleName);
+                }
+                mVersion = VERSION_FALLBACK_STATE_MIGRATED;
+                scheduleWriteFileLocked();
+            }
+        }
+    }
+
+    /**
      * Get whether the role is available.
      *
      * @param roleName the name of the role to get the holders for
@@ -269,6 +295,7 @@ class RoleUserState {
         synchronized (mLock) {
             if (!mRoles.containsKey(roleName)) {
                 mRoles.put(roleName, new ArraySet<>());
+                mFallbackEnabledRoles.add(roleName);
                 Log.i(LOG_TAG, "Added new role: " + roleName);
                 scheduleWriteFileLocked();
                 return true;
@@ -297,6 +324,7 @@ class RoleUserState {
                                 + " role: " + roleName + ", holders: " + packageNames);
                     }
                     mRoles.removeAt(i);
+                    mFallbackEnabledRoles.remove(roleName);
                     changed = true;
                 }
             }
@@ -432,12 +460,15 @@ class RoleUserState {
             RolesState roleState = mPersistence.readForUser(UserHandle.of(mUserId));
 
             Map<String, Set<String>> roles;
+            Set<String> fallbackEnabledRoles;
             if (roleState != null) {
                 mVersion = roleState.getVersion();
                 mPackagesHash = roleState.getPackagesHash();
                 roles = roleState.getRoles();
+                fallbackEnabledRoles = roleState.getFallbackEnabledRoles();
             } else {
                 roles = mPlatformHelper.getLegacyRoleState(mUserId);
+                fallbackEnabledRoles = roles.keySet();
             }
             mRoles.clear();
             for (Map.Entry<String, Set<String>> entry : roles.entrySet()) {
@@ -445,12 +476,10 @@ class RoleUserState {
                 ArraySet<String> roleHolders = new ArraySet<>(entry.getValue());
                 mRoles.put(roleName, roleHolders);
             }
-
+            mFallbackEnabledRoles.clear();
+            mFallbackEnabledRoles.addAll(fallbackEnabledRoles);
             if (roleState == null) {
                 scheduleWriteFileLocked();
-            } else {
-                mFallbackEnabledRoles.clear();
-                mFallbackEnabledRoles.addAll(roleState.getFallbackEnabledRoles());
             }
         }
     }
diff --git a/service/java/com/android/role/persistence/RolesState.java b/service/java/com/android/role/persistence/RolesState.java
index 26644c358..a189dd4c2 100644
--- a/service/java/com/android/role/persistence/RolesState.java
+++ b/service/java/com/android/role/persistence/RolesState.java
@@ -22,7 +22,6 @@ import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.annotation.SystemApi.Client;
 import android.permission.flags.Flags;
-import android.util.ArraySet;
 
 import java.util.Map;
 import java.util.Objects;
@@ -68,7 +67,7 @@ public final class RolesState {
      */
     public RolesState(int version, @Nullable String packagesHash,
             @NonNull Map<String, Set<String>> roles) {
-        this(version, packagesHash, roles, new ArraySet<>());
+        this(version, packagesHash, roles, roles.keySet());
     }
 
     /**
diff --git a/tests/apex/java/com/android/role/persistence/RolesPersistenceTest.kt b/tests/apex/java/com/android/role/persistence/RolesPersistenceTest.kt
index 0cf1fa665..6500b3926 100644
--- a/tests/apex/java/com/android/role/persistence/RolesPersistenceTest.kt
+++ b/tests/apex/java/com/android/role/persistence/RolesPersistenceTest.kt
@@ -111,6 +111,7 @@ class RolesPersistenceTest {
 
         assertThat(persistedState).isNull()
     }
+
     private fun getState(): RolesState =
         when (stateVersion) {
             StateVersion.VERSION_UNDEFINED -> stateVersionUndefined