Integrated FactoryResetter with DevicePolicySafetyChecker.

Test: atest FrameworksMockingServicesTests:FactoryResetterTest \
            CtsDevicePolicyManagerTestCases:DeviceOwnerTest#testPolicySafetyCheckerIntegration

Bug: 171603586

Change-Id: Ib046d168c3380bbe4a3923530860ba0acb8fb6ae
diff --git a/core/java/android/app/admin/DevicePolicySafetyChecker.java b/core/java/android/app/admin/DevicePolicySafetyChecker.java
index 1f8a933..b1a80c5 100644
--- a/core/java/android/app/admin/DevicePolicySafetyChecker.java
+++ b/core/java/android/app/admin/DevicePolicySafetyChecker.java
@@ -18,6 +18,8 @@
 import android.annotation.NonNull;
 import android.app.admin.DevicePolicyManager.DevicePolicyOperation;
 
+import com.android.internal.os.IResultReceiver;
+
 /**
  * Interface responsible to check if a {@link DevicePolicyManager} API can be safely executed.
  *
@@ -28,9 +30,7 @@
     /**
      * Returns whether the given {@code operation} can be safely executed at the moment.
      */
-    default boolean isDevicePolicyOperationSafe(@DevicePolicyOperation int operation) {
-        return true;
-    }
+    boolean isDevicePolicyOperationSafe(@DevicePolicyOperation int operation);
 
     /**
      * Returns a new exception for when the given {@code operation} cannot be safely executed.
@@ -39,4 +39,13 @@
     default UnsafeStateException newUnsafeStateException(@DevicePolicyOperation int operation) {
         return new UnsafeStateException(operation);
     }
+
+    /**
+     * Called when a request was made to factory reset the device, so it can be delayed if it's not
+     * safe to proceed.
+     *
+     * @param callback callback whose {@code send()} method must be called when it's safe to factory
+     * reset.
+     */
+    void onFactoryReset(IResultReceiver callback);
 }
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/CallerIdentity.java b/services/devicepolicy/java/com/android/server/devicepolicy/CallerIdentity.java
index b6b4d8a..4e422b2 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/CallerIdentity.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/CallerIdentity.java
@@ -25,7 +25,7 @@
  * All parameters are verified on object creation unless the component name is null and the
  * caller is a delegate.
  */
-class CallerIdentity {
+final class CallerIdentity {
 
     private final int mUid;
     @Nullable
@@ -51,7 +51,7 @@
         return UserHandle.getUserHandleForUid(mUid);
     }
 
-    @Nullable  public String getPackageName() {
+    @Nullable public String getPackageName() {
         return mPackageName;
     }
 
@@ -66,4 +66,16 @@
     public boolean hasPackage() {
         return mPackageName != null;
     }
+
+    @Override
+    public String toString() {
+        StringBuilder builder = new StringBuilder("CallerIdentity[uid=").append(mUid);
+        if (mPackageName != null) {
+            builder.append(", pkg=").append(mPackageName);
+        }
+        if (mComponentName != null) {
+            builder.append(", cmp=").append(mComponentName.flattenToShortString());
+        }
+        return builder.append("]").toString();
+    }
 }
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 7b4c1be..f39813a 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -983,8 +983,20 @@
 
     @Override
     public void setDevicePolicySafetyChecker(DevicePolicySafetyChecker safetyChecker) {
-        Slog.i(LOG_TAG, "Setting DevicePolicySafetyChecker as " + safetyChecker);
+        CallerIdentity callerIdentity = getCallerIdentity();
+        Preconditions.checkCallAuthorization(mIsAutomotive || isAdb(callerIdentity), "can only set "
+                + "DevicePolicySafetyChecker on automotive builds or from ADB (but caller is %s)",
+                callerIdentity);
+        setDevicePolicySafetyCheckerUnchecked(safetyChecker);
+    }
+
+    /**
+     * Used by {@code setDevicePolicySafetyChecker()} above and {@link OneTimeSafetyChecker}.
+     */
+    void setDevicePolicySafetyCheckerUnchecked(DevicePolicySafetyChecker safetyChecker) {
+        Slog.i(LOG_TAG, String.format("Setting DevicePolicySafetyChecker as %s", safetyChecker));
         mSafetyChecker = safetyChecker;
+        mInjector.setDevicePolicySafetyChecker(safetyChecker);
     }
 
     /**
@@ -1021,8 +1033,8 @@
     public void setNextOperationSafety(@DevicePolicyOperation int operation, boolean safe) {
         Preconditions.checkCallAuthorization(
                 hasCallingOrSelfPermission(permission.MANAGE_DEVICE_ADMINS));
-        Slog.i(LOG_TAG, "setNextOperationSafety(" + DevicePolicyManager.operationToString(operation)
-                + ", " + safe + ")");
+        Slog.i(LOG_TAG, String.format("setNextOperationSafety(%s, %b)",
+                DevicePolicyManager.operationToString(operation), safe));
         mSafetyChecker = new OneTimeSafetyChecker(this, operation, safe);
     }
 
@@ -1034,6 +1046,8 @@
 
         public final Context mContext;
 
+        private @Nullable DevicePolicySafetyChecker mSafetyChecker;
+
         Injector(Context context) {
             mContext = context;
         }
@@ -1278,7 +1292,8 @@
         void recoverySystemRebootWipeUserData(boolean shutdown, String reason, boolean force,
                 boolean wipeEuicc, boolean wipeExtRequested, boolean wipeResetProtectionData)
                         throws IOException {
-            FactoryResetter.newBuilder(mContext).setReason(reason).setShutdown(shutdown)
+            FactoryResetter.newBuilder(mContext).setSafetyChecker(mSafetyChecker)
+                    .setReason(reason).setShutdown(shutdown)
                     .setForce(force).setWipeEuicc(wipeEuicc)
                     .setWipeAdoptableStorage(wipeExtRequested)
                     .setWipeFactoryResetProtection(wipeResetProtectionData)
@@ -1433,6 +1448,10 @@
         public boolean isChangeEnabled(long changeId, String packageName, int userId) {
             return CompatChanges.isChangeEnabled(changeId, packageName, UserHandle.of(userId));
         }
+
+        void setDevicePolicySafetyChecker(DevicePolicySafetyChecker safetyChecker) {
+            mSafetyChecker = safetyChecker;
+        }
     }
 
     /**
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/FactoryResetter.java b/services/devicepolicy/java/com/android/server/devicepolicy/FactoryResetter.java
index a0cf7a3..564950b 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/FactoryResetter.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/FactoryResetter.java
@@ -17,14 +17,18 @@
 package com.android.server.devicepolicy;
 
 import android.annotation.Nullable;
+import android.app.admin.DevicePolicySafetyChecker;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.os.Bundle;
 import android.os.RecoverySystem;
+import android.os.RemoteException;
 import android.os.UserManager;
 import android.os.storage.StorageManager;
 import android.service.persistentdata.PersistentDataBlockManager;
-import android.util.Log;
+import android.util.Slog;
 
+import com.android.internal.os.IResultReceiver;
 import com.android.internal.util.Preconditions;
 
 import java.io.IOException;
@@ -38,6 +42,7 @@
     private static final String TAG = FactoryResetter.class.getSimpleName();
 
     private final Context mContext;
+    private final @Nullable DevicePolicySafetyChecker mSafetyChecker;
     private final @Nullable String mReason;
     private final boolean mShutdown;
     private final boolean mForce;
@@ -45,18 +50,40 @@
     private final boolean mWipeAdoptableStorage;
     private final boolean mWipeFactoryResetProtection;
 
-
     /**
      * Factory reset the device according to the builder's arguments.
      */
     public void factoryReset() throws IOException {
-        Log.i(TAG, String.format("factoryReset(): reason=%s, shutdown=%b, force=%b, wipeEuicc=%b"
-                + ", wipeAdoptableStorage=%b, wipeFRP=%b", mReason, mShutdown, mForce, mWipeEuicc,
-                mWipeAdoptableStorage, mWipeFactoryResetProtection));
-
         Preconditions.checkCallAuthorization(mContext.checkCallingOrSelfPermission(
                 android.Manifest.permission.MASTER_CLEAR) == PackageManager.PERMISSION_GRANTED);
 
+        if (mSafetyChecker == null) {
+            factoryResetInternalUnchecked();
+            return;
+        }
+
+        IResultReceiver receiver = new IResultReceiver.Stub() {
+            @Override
+            public void send(int resultCode, Bundle resultData) throws RemoteException {
+                Slog.i(TAG, String.format("Factory reset confirmed by %s, proceeding",
+                        mSafetyChecker));
+                try {
+                    factoryResetInternalUnchecked();
+                } catch (IOException e) {
+                    // Shouldn't happen
+                    Slog.wtf(TAG, "IOException calling underlying systems", e);
+                }
+            }
+        };
+        Slog.i(TAG, String.format("Delaying factory reset until %s confirms", mSafetyChecker));
+        mSafetyChecker.onFactoryReset(receiver);
+    }
+
+    private void factoryResetInternalUnchecked() throws IOException {
+        Slog.i(TAG, String.format("factoryReset(): reason=%s, shutdown=%b, force=%b, wipeEuicc=%b, "
+                + "wipeAdoptableStorage=%b, wipeFRP=%b", mReason, mShutdown, mForce, mWipeEuicc,
+                mWipeAdoptableStorage, mWipeFactoryResetProtection));
+
         UserManager um = mContext.getSystemService(UserManager.class);
         if (!mForce && um.hasUserRestriction(UserManager.DISALLOW_FACTORY_RESET)) {
             throw new SecurityException("Factory reset is not allowed for this user.");
@@ -66,15 +93,15 @@
             PersistentDataBlockManager manager = mContext
                     .getSystemService(PersistentDataBlockManager.class);
             if (manager != null) {
-                Log.w(TAG, "Wiping factory reset protection");
+                Slog.w(TAG, "Wiping factory reset protection");
                 manager.wipe();
             } else {
-                Log.w(TAG, "No need to wipe factory reset protection");
+                Slog.w(TAG, "No need to wipe factory reset protection");
             }
         }
 
         if (mWipeAdoptableStorage) {
-            Log.w(TAG, "Wiping adoptable storage");
+            Slog.w(TAG, "Wiping adoptable storage");
             StorageManager sm = mContext.getSystemService(StorageManager.class);
             sm.wipeAdoptableDisks();
         }
@@ -84,8 +111,9 @@
 
     private FactoryResetter(Builder builder) {
         mContext = builder.mContext;
-        mShutdown = builder.mShutdown;
+        mSafetyChecker = builder.mSafetyChecker;
         mReason = builder.mReason;
+        mShutdown = builder.mShutdown;
         mForce = builder.mForce;
         mWipeEuicc = builder.mWipeEuicc;
         mWipeAdoptableStorage = builder.mWipeAdoptableStorage;
@@ -105,6 +133,7 @@
     public static final class Builder {
 
         private final Context mContext;
+        private @Nullable DevicePolicySafetyChecker mSafetyChecker;
         private @Nullable String mReason;
         private boolean mShutdown;
         private boolean mForce;
@@ -117,6 +146,15 @@
         }
 
         /**
+         * Sets a {@link DevicePolicySafetyChecker} object that will be used to delay the
+         * factory reset when it's not safe to do so.
+         */
+        public Builder setSafetyChecker(@Nullable DevicePolicySafetyChecker safetyChecker) {
+            mSafetyChecker = safetyChecker;
+            return this;
+        }
+
+        /**
          * Sets the (non-null) reason for the factory reset that is visible in the logs
          */
         public Builder setReason(String reason) {
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/OneTimeSafetyChecker.java b/services/devicepolicy/java/com/android/server/devicepolicy/OneTimeSafetyChecker.java
index b0f8bfb..f7a8261 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/OneTimeSafetyChecker.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/OneTimeSafetyChecker.java
@@ -21,6 +21,8 @@
 import android.app.admin.DevicePolicySafetyChecker;
 import android.util.Slog;
 
+import com.android.internal.os.IResultReceiver;
+
 import java.util.Objects;
 
 //TODO(b/172376923): add unit tests
@@ -61,7 +63,12 @@
         }
         Slog.i(TAG, "isDevicePolicyOperationSafe(" + name + "): returning " + safe
                 + " and restoring DevicePolicySafetyChecker to " + mRealSafetyChecker);
-        mService.setDevicePolicySafetyChecker(mRealSafetyChecker);
+        mService.setDevicePolicySafetyCheckerUnchecked(mRealSafetyChecker);
         return safe;
     }
+
+    @Override
+    public void onFactoryReset(IResultReceiver callback) {
+        throw new UnsupportedOperationException();
+    }
 }
diff --git a/services/tests/mockingservicestests/src/com/android/server/devicepolicy/FactoryResetterTest.java b/services/tests/mockingservicestests/src/com/android/server/devicepolicy/FactoryResetterTest.java
index 829e312..c4d14f9 100644
--- a/services/tests/mockingservicestests/src/com/android/server/devicepolicy/FactoryResetterTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/devicepolicy/FactoryResetterTest.java
@@ -26,6 +26,7 @@
 import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertThrows;
 
+import android.app.admin.DevicePolicySafetyChecker;
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.os.RecoverySystem;
@@ -35,6 +36,8 @@
 import android.service.persistentdata.PersistentDataBlockManager;
 import android.util.Log;
 
+import com.android.internal.os.IResultReceiver;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -58,6 +61,7 @@
     private @Mock StorageManager mSm;
     private @Mock PersistentDataBlockManager mPdbm;
     private @Mock UserManager mUm;
+    private @Mock DevicePolicySafetyChecker mSafetyChecker;
 
     @Before
     public void startSession() {
@@ -69,7 +73,7 @@
 
         when(mContext.getSystemService(any(Class.class))).thenAnswer((inv) -> {
             Log.d(TAG, "Mocking " + inv);
-            Class serviceClass = (Class) inv.getArguments()[0];
+            Class<?> serviceClass = (Class<?>) inv.getArguments()[0];
             if (serviceClass.equals(PersistentDataBlockManager.class)) return mPdbm;
             if (serviceClass.equals(StorageManager.class)) return mSm;
             if (serviceClass.equals(UserManager.class)) return mUm;
@@ -194,6 +198,44 @@
         verifyRebootWipeUserDataAllArgsCalled();
     }
 
+    @Test
+    public void testFactoryReset_minimumArgs_safetyChecker_neverReplied() throws Exception {
+        allowFactoryReset();
+
+        FactoryResetter.newBuilder(mContext).setSafetyChecker(mSafetyChecker).build()
+                .factoryReset();
+
+        verifyWipeAdoptableStorageNotCalled();
+        verifyWipeFactoryResetProtectionNotCalled();
+        verifyRebootWipeUserDataNotCalled();
+    }
+
+    @Test
+    public void testFactoryReset_allArgs_safetyChecker_replied() throws Exception {
+        allowFactoryReset();
+
+        doAnswer((inv) -> {
+            Log.d(TAG, "Mocking " + inv);
+            IResultReceiver receiver = (IResultReceiver) inv.getArguments()[0];
+            receiver.send(0, null);
+            return null;
+        }).when(mSafetyChecker).onFactoryReset(any());
+
+        FactoryResetter.newBuilder(mContext)
+                .setSafetyChecker(mSafetyChecker)
+                .setReason(REASON)
+                .setForce(true)
+                .setShutdown(true)
+                .setWipeEuicc(true)
+                .setWipeAdoptableStorage(true)
+                .setWipeFactoryResetProtection(true)
+                .build().factoryReset();
+
+        verifyWipeAdoptableStorageCalled();
+        verifyWipeFactoryResetProtectionCalled();
+        verifyRebootWipeUserDataAllArgsCalled();
+    }
+
     private void revokeMasterClearPermission() {
         when(mContext.checkCallingOrSelfPermission(android.Manifest.permission.MASTER_CLEAR))
                 .thenReturn(PackageManager.PERMISSION_DENIED);