Add ResetPasswordWithToken migration code.

Flag: android.app.admin.flags.reset_password_with_token_coexistence
Change-Id: Ibbd3c15d884d20aacc8a007f1ad7721c9d42ca47
Test: Manually tested. Generated a password token through TestDPC, enabled respective flag and rebooted. While rebooting, I've checked that the migration code was executed through logcat. After rebooting, I've successfully created a password with pre-existing token. After rebooting once more, I've checked that it didn't try to migrate again through logs.
Bug: 359187209
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 57a7b93..407a5a6 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -3515,6 +3515,48 @@
         return true;
     }
 
+    @GuardedBy("getLockObject()")
+    private boolean maybeMigrateResetPasswordTokenLocked(String backupId) {
+        if (!Flags.resetPasswordWithTokenCoexistence()) {
+            Slog.i(LOG_TAG, "ResetPasswordWithToken not migrated because coexistence "
+                    + "support is not enabled.");
+            return false;
+        }
+        if (mOwners.isResetPasswordWithTokenMigrated()) {
+            // TODO(b/359187209): Remove log after Flags.resetPasswordWithTokenCoexistence full
+            //  rollout.
+            Slog.v(LOG_TAG, "ResetPasswordWithToken was previously migrated to "
+                    + "policy engine.");
+            return false;
+        }
+
+        Slog.i(LOG_TAG, "Migrating ResetPasswordWithToken to policy engine");
+
+        // Create backup if none exists
+        mDevicePolicyEngine.createBackup(backupId);
+        try {
+            iterateThroughDpcAdminsLocked((admin, enforcingAdmin) -> {
+                int userId = enforcingAdmin.getUserId();
+                DevicePolicyData policy = getUserData(userId);
+                if (policy.mPasswordTokenHandle != 0) {
+                    Slog.i(LOG_TAG, "Setting RESET_PASSWORD_TOKEN policy");
+                    mDevicePolicyEngine.setLocalPolicy(
+                            PolicyDefinition.RESET_PASSWORD_TOKEN,
+                            enforcingAdmin,
+                            new LongPolicyValue(policy.mPasswordTokenHandle),
+                            userId);
+                }
+            });
+        } catch (Exception e) {
+            Slog.wtf(LOG_TAG,
+                    "Failed to migrate ResetPasswordWithToken to policy engine", e);
+        }
+
+        Slog.i(LOG_TAG, "Marking ResetPasswordWithToken migration complete");
+        mOwners.markResetPasswordWithTokenMigrated();
+        return true;
+    }
+
     /** Register callbacks for statsd pulled atoms. */
     private void registerStatsCallbacks() {
         final StatsManager statsManager = mContext.getSystemService(StatsManager.class);
@@ -19342,6 +19384,8 @@
                     PolicyDefinition.RESET_PASSWORD_TOKEN,
                     enforcingAdmin,
                     userId);
+            // TODO(b/369152176): Address difference in behavior regarding addEscrowToken when
+            //  compared with the else branch.
             long tokenHandle = addEscrowToken(
                     token, currentTokenHandle == null ? 0 : currentTokenHandle, userId);
             if (tokenHandle == 0) {
@@ -24280,12 +24324,21 @@
         maybeMigrateSecurityLoggingPolicyLocked();
         // ID format: <sdk-int>.<auto_increment_id>.<descriptions>'
         String unmanagedBackupId = "35.1.unmanaged-mode";
-        boolean migrated = false;
-        migrated = migrated | maybeMigrateRequiredPasswordComplexityLocked(unmanagedBackupId);
-        migrated = migrated | maybeMigrateSuspendedPackagesLocked(unmanagedBackupId);
-        if (migrated) {
+        boolean unmanagedMigrated = false;
+        unmanagedMigrated =
+                unmanagedMigrated | maybeMigrateRequiredPasswordComplexityLocked(unmanagedBackupId);
+        unmanagedMigrated =
+                unmanagedMigrated | maybeMigrateSuspendedPackagesLocked(unmanagedBackupId);
+        if (unmanagedMigrated) {
             Slogf.i(LOG_TAG, "Backup made: " + unmanagedBackupId);
         }
+
+        String supervisionBackupId = "36.2.supervision-support";
+        boolean supervisionMigrated = maybeMigrateResetPasswordTokenLocked(supervisionBackupId);
+        if (supervisionMigrated) {
+            Slogf.i(LOG_TAG, "Backup made: " + supervisionBackupId);
+        }
+
         // Additional migration steps should repeat the pattern above with a new backupId.
     }
 
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/Owners.java b/services/devicepolicy/java/com/android/server/devicepolicy/Owners.java
index 3f9605a..b3c8408 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/Owners.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/Owners.java
@@ -669,6 +669,19 @@
         }
     }
 
+    void markResetPasswordWithTokenMigrated() {
+        synchronized (mData) {
+            mData.mResetPasswordWithTokenMigrated = true;
+            mData.writeDeviceOwner();
+        }
+    }
+
+    boolean isResetPasswordWithTokenMigrated() {
+        synchronized (mData) {
+            return mData.mResetPasswordWithTokenMigrated;
+        }
+    }
+
     @GuardedBy("mData")
     void pushToAppOpsLocked() {
         if (!mSystemReady) {
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/OwnersData.java b/services/devicepolicy/java/com/android/server/devicepolicy/OwnersData.java
index 87fd002..10e43d9 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/OwnersData.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/OwnersData.java
@@ -91,6 +91,8 @@
     private static final String ATTR_REQUIRED_PASSWORD_COMPLEXITY_MIGRATED =
             "passwordComplexityMigrated";
     private static final String ATTR_SUSPENDED_PACKAGES_MIGRATED = "suspendedPackagesMigrated";
+    private static final String ATTR_RESET_PASSWORD_WITH_TOKEN_MIGRATED =
+            "resetPasswordWithTokenMigrated";
     private static final String ATTR_MIGRATED_POST_UPGRADE = "migratedPostUpgrade";
 
     // Internal state for the device owner package.
@@ -122,6 +124,7 @@
     boolean mSecurityLoggingMigrated = false;
     boolean mRequiredPasswordComplexityMigrated = false;
     boolean mSuspendedPackagesMigrated = false;
+    boolean mResetPasswordWithTokenMigrated = false;
 
     boolean mPoliciesMigratedPostUpdate = false;
 
@@ -417,7 +420,10 @@
                         mSuspendedPackagesMigrated);
 
             }
-
+            if (Flags.resetPasswordWithTokenCoexistence()) {
+                out.attributeBoolean(null, ATTR_RESET_PASSWORD_WITH_TOKEN_MIGRATED,
+                        mResetPasswordWithTokenMigrated);
+            }
             out.endTag(null, TAG_POLICY_ENGINE_MIGRATION);
 
         }
@@ -488,6 +494,9 @@
                     mSuspendedPackagesMigrated = Flags.unmanagedModeMigration()
                             && parser.getAttributeBoolean(null,
                                     ATTR_SUSPENDED_PACKAGES_MIGRATED, false);
+                    mResetPasswordWithTokenMigrated = Flags.resetPasswordWithTokenCoexistence()
+                            && parser.getAttributeBoolean(null,
+                            ATTR_RESET_PASSWORD_WITH_TOKEN_MIGRATED, false);
 
                     break;
                 default: