Disable package when finalization is complete

Disable the DLC package once we have determined we should be finalized.

Because the package is disabled and will not run, the system service has
to also persist a finalized state to know not to re-enable the DLC on
boot. This is done by wring to a simple XML file in global device state
once the DLC has determined that everything is finished.

Bug: 279517666
Test: atest DeviceLockControllerRoboTests
Change-Id: I30c29b27586c13c84901ea8b81f0542bfa543b45
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/SystemDeviceLockManager.java b/DeviceLockController/src/com/android/devicelockcontroller/SystemDeviceLockManager.java
index 39897a3..6df7655 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/SystemDeviceLockManager.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/SystemDeviceLockManager.java
@@ -93,4 +93,15 @@
      */
     void disableKioskKeepalive(@CallbackExecutor Executor executor,
             @NonNull OutcomeReceiver<Void, Exception> callback);
+
+    /**
+     * Set whether device is finalized so that system service knows when to keep the Device Lock
+     * Controller enabled.
+     *
+     * @param finalized true if device is finalized and DLC should not be enabled.
+     * @param executor the {@link Executor} on which to invoke the callback.
+     * @param callback callback this returns either success or an exception.
+     */
+    void setDeviceFinalized(boolean finalized, @CallbackExecutor Executor executor,
+            @NonNull OutcomeReceiver<Void, Exception> callback);
 }
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/SystemDeviceLockManagerImpl.java b/DeviceLockController/src/com/android/devicelockcontroller/SystemDeviceLockManagerImpl.java
index 7b9599f..e376aae 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/SystemDeviceLockManagerImpl.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/SystemDeviceLockManagerImpl.java
@@ -219,4 +219,21 @@
             executor.execute(() -> callback.onError(new RuntimeException(e)));
         }
     }
+
+    @Override
+    @RequiresPermission(MANAGE_DEVICE_LOCK_SERVICE_FROM_CONTROLLER)
+    public void setDeviceFinalized(boolean finalized, Executor executor,
+            @NonNull OutcomeReceiver<Void, Exception> callback) {
+        Objects.requireNonNull(executor);
+        Objects.requireNonNull(callback);
+
+        try {
+            mIDeviceLockService.setDeviceFinalized(finalized,
+                    new RemoteCallback(result -> executor.execute(() -> {
+                        callback.onResult(null /* result */);
+                    }), new Handler(Looper.getMainLooper())));
+        } catch (RemoteException e) {
+            executor.execute(() -> callback.onError(new RuntimeException(e)));
+        }
+    }
 }
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/policy/FinalizationControllerImpl.java b/DeviceLockController/src/com/android/devicelockcontroller/policy/FinalizationControllerImpl.java
index 87a8790..1a5db89 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/policy/FinalizationControllerImpl.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/policy/FinalizationControllerImpl.java
@@ -23,12 +23,16 @@
 import static com.android.devicelockcontroller.provision.worker.ReportDeviceLockProgramCompleteWorker.REPORT_DEVICE_LOCK_PROGRAM_COMPLETE_WORK_NAME;
 
 import android.annotation.IntDef;
+import android.app.AlarmManager;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.os.OutcomeReceiver;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
 import androidx.annotation.WorkerThread;
+import androidx.concurrent.futures.CallbackToFutureAdapter;
 import androidx.work.Constraints;
 import androidx.work.ExistingWorkPolicy;
 import androidx.work.ListenableWorker;
@@ -37,10 +41,11 @@
 import androidx.work.Operation;
 import androidx.work.WorkManager;
 
+import com.android.devicelockcontroller.SystemDeviceLockManager;
+import com.android.devicelockcontroller.SystemDeviceLockManagerImpl;
 import com.android.devicelockcontroller.provision.grpc.DeviceFinalizeClient.ReportDeviceProgramCompleteResponse;
 import com.android.devicelockcontroller.provision.worker.ReportDeviceLockProgramCompleteWorker;
 import com.android.devicelockcontroller.receivers.FinalizationBootCompletedReceiver;
-import com.android.devicelockcontroller.receivers.LockedBootCompletedReceiver;
 import com.android.devicelockcontroller.storage.GlobalParametersClient;
 import com.android.devicelockcontroller.util.LogUtil;
 
@@ -88,8 +93,9 @@
 
     /** Dispatch queue to guarantee state changes occur sequentially */
     private final FinalizationStateDispatchQueue mDispatchQueue;
-    private final Executor mLightweightExecutor;
+    private final Executor mBgExecutor;
     private final Context mContext;
+    private final SystemDeviceLockManager mSystemDeviceLockManager;
     private final Class<? extends ListenableWorker> mReportDeviceFinalizedWorkerClass;
     /** Future for after initial finalization state is set from disk */
     private final ListenableFuture<Void> mStateInitializedFuture;
@@ -98,27 +104,30 @@
         this(context,
                 new FinalizationStateDispatchQueue(),
                 Executors.newCachedThreadPool(),
-                ReportDeviceLockProgramCompleteWorker.class);
+                ReportDeviceLockProgramCompleteWorker.class,
+                SystemDeviceLockManagerImpl.getInstance());
     }
 
     @VisibleForTesting
     public FinalizationControllerImpl(
             Context context,
             FinalizationStateDispatchQueue dispatchQueue,
-            Executor lightweightExecutor,
-            Class<? extends ListenableWorker> reportDeviceFinalizedWorkerClass) {
+            Executor bgExecutor,
+            Class<? extends ListenableWorker> reportDeviceFinalizedWorkerClass,
+            SystemDeviceLockManager systemDeviceLockManager) {
         mContext = context;
         mDispatchQueue = dispatchQueue;
         mDispatchQueue.init(this::onStateChanged);
-        mLightweightExecutor = lightweightExecutor;
+        mBgExecutor = bgExecutor;
         mReportDeviceFinalizedWorkerClass = reportDeviceFinalizedWorkerClass;
+        mSystemDeviceLockManager = systemDeviceLockManager;
 
         // Set the initial state
         ListenableFuture<Integer> initialStateFuture =
                 GlobalParametersClient.getInstance().getFinalizationState();
         mStateInitializedFuture = Futures.transformAsync(initialStateFuture,
                 mDispatchQueue::enqueueStateChange,
-                mLightweightExecutor);
+                mBgExecutor);
     }
 
     @VisibleForTesting
@@ -130,7 +139,7 @@
     public ListenableFuture<Void> notifyRestrictionsCleared() {
         return Futures.transformAsync(mStateInitializedFuture,
                 unused -> mDispatchQueue.enqueueStateChange(FINALIZED_UNREPORTED),
-                mLightweightExecutor);
+                mBgExecutor);
     }
 
     @Override
@@ -139,7 +148,7 @@
         if (response.isSuccessful()) {
             return Futures.transformAsync(mStateInitializedFuture,
                     unused -> mDispatchQueue.enqueueStateChange(FINALIZED),
-                    mLightweightExecutor);
+                    mBgExecutor);
         } else {
             // TODO(279517666): Determine how to handle an unrecoverable failure
             // response from the server
@@ -151,6 +160,8 @@
     @WorkerThread
     private ListenableFuture<Void> onStateChanged(@FinalizationState int oldState,
             @FinalizationState int newState) {
+        final ListenableFuture<Void> persistStateFuture =
+                GlobalParametersClient.getInstance().setFinalizationState(newState);
         if (oldState == UNFINALIZED) {
             // Enable boot receiver to check finalization state on disk
             PackageManager pm = mContext.getPackageManager();
@@ -162,20 +173,21 @@
         }
         switch (newState) {
             case UNFINALIZED:
-                // no-op
-                break;
+                return persistStateFuture;
             case FINALIZED_UNREPORTED:
                 requestWorkToReportFinalized();
-                break;
+                return persistStateFuture;
             case FINALIZED:
-                disableEntireApplication();
-                break;
+                // Ensure disabling only happens after state is written to disk in case we somehow
+                // exit the disabled state and need to disable again.
+                return Futures.transformAsync(persistStateFuture,
+                        unused -> disableEntireApplication(),
+                        mBgExecutor);
             case UNINITIALIZED:
                 throw new IllegalArgumentException("Tried to set state back to uninitialized!");
             default:
                 throw new IllegalArgumentException("Unknown state " + newState);
         }
-        return GlobalParametersClient.getInstance().setFinalizationState(newState);
     }
 
     /**
@@ -215,15 +227,43 @@
      *
      * This will remove any work, alarms, receivers, etc., and this application should never run
      * on the device again after this point.
+     *
+     * This method returns a future but it is a bit of an odd case as the application itself
+     * may end up disabled before/after the future is handled depending on when package manager
+     * enforces the application is disabled.
+     *
+     * @return future for when this is done
      */
-    private void disableEntireApplication() {
-        PackageManager pm = mContext.getPackageManager();
-        pm.setComponentEnabledSetting(
-                new ComponentName(mContext,
-                        LockedBootCompletedReceiver.class),
-                PackageManager.COMPONENT_ENABLED_STATE_DISABLED,
-                0 /* flags */);
-        // TODO(279517666): Disable application and persist a boolean so that DeviceLockService
-        // does not re-enable application on boot / user switch
+    private ListenableFuture<Void> disableEntireApplication() {
+        WorkManager workManager = WorkManager.getInstance(mContext);
+        workManager.cancelAllWork();
+        AlarmManager alarmManager = mContext.getSystemService(AlarmManager.class);
+        alarmManager.cancelAll();
+        ListenableFuture<Void> disableApplicationFuture = CallbackToFutureAdapter.getFuture(
+                completer -> {
+                        mSystemDeviceLockManager.setDeviceFinalized(true, mBgExecutor,
+                                new OutcomeReceiver<>() {
+                                    @Override
+                                    public void onResult(Void result) {
+                                        // This kills and disables the app
+                                        PackageManager pm = mContext.getPackageManager();
+                                        pm.setApplicationEnabledSetting(
+                                                mContext.getPackageName(),
+                                                PackageManager.COMPONENT_ENABLED_STATE_DISABLED,
+                                                0 /* flags */);
+                                        completer.set(null);
+                                    }
+
+                                    @Override
+                                    public void onError(@NonNull Exception error) {
+                                        LogUtil.e(TAG, "Failed to set device finalized in"
+                                                + "system service.", error);
+                                        completer.setException(error);
+                                    }
+                                });
+                    return "Disable application future";
+                }
+        );
+        return disableApplicationFuture;
     }
 }
diff --git a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/FinalizationControllerImplTest.java b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/FinalizationControllerImplTest.java
index f5e3cb4..8788589 100644
--- a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/FinalizationControllerImplTest.java
+++ b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/FinalizationControllerImplTest.java
@@ -19,12 +19,11 @@
 import static com.android.devicelockcontroller.policy.FinalizationControllerImpl.FinalizationState.FINALIZED;
 import static com.android.devicelockcontroller.policy.FinalizationControllerImpl.FinalizationState.FINALIZED_UNREPORTED;
 import static com.android.devicelockcontroller.provision.worker.ReportDeviceLockProgramCompleteWorker.REPORT_DEVICE_LOCK_PROGRAM_COMPLETE_WORK_NAME;
-
 import static com.google.common.truth.Truth.assertThat;
 
-import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.os.OutcomeReceiver;
 
 import androidx.annotation.NonNull;
 import androidx.test.core.app.ApplicationProvider;
@@ -34,8 +33,8 @@
 import androidx.work.WorkerParameters;
 import androidx.work.testing.WorkManagerTestInitHelper;
 
+import com.android.devicelockcontroller.SystemDeviceLockManager;
 import com.android.devicelockcontroller.provision.grpc.DeviceFinalizeClient.ReportDeviceProgramCompleteResponse;
-import com.android.devicelockcontroller.receivers.LockedBootCompletedReceiver;
 import com.android.devicelockcontroller.storage.GlobalParametersClient;
 
 import com.google.common.util.concurrent.ExecutionSequencer;
@@ -57,11 +56,12 @@
 
     private static final int TIMEOUT_MS = 1000;
 
+    private SystemDeviceLockManager mSystemDeviceLockManager = new TestSystemDeviceLockManager();
     private Context mContext;
     private FinalizationControllerImpl mFinalizationController;
     private FinalizationStateDispatchQueue mDispatchQueue;
     private ExecutionSequencer mExecutionSequencer = ExecutionSequencer.create();
-    private Executor mLightweightExecutor = Executors.newCachedThreadPool();
+    private Executor mBgExecutor = Executors.newCachedThreadPool();
     private GlobalParametersClient mGlobalParametersClient;
 
     @Before
@@ -110,10 +110,8 @@
 
         // THEN the application is disabled and the disk value is set to finalized
         PackageManager pm = mContext.getPackageManager();
-        assertThat(pm.getComponentEnabledSetting(
-                new ComponentName(mContext, LockedBootCompletedReceiver.class)))
+        assertThat(pm.getApplicationEnabledSetting(mContext.getPackageName()))
                 .isEqualTo(PackageManager.COMPONENT_ENABLED_STATE_DISABLED);
-        // TODO(279517666): Assert checks that application itself is disabled when implemented
         assertThat(mGlobalParametersClient.getFinalizationState().get()).isEqualTo(FINALIZED);
     }
 
@@ -138,7 +136,52 @@
 
     private FinalizationControllerImpl makeFinalizationController() {
         return new FinalizationControllerImpl(
-                mContext, mDispatchQueue, mLightweightExecutor, TestWorker.class);
+                mContext, mDispatchQueue, mBgExecutor, TestWorker.class, mSystemDeviceLockManager);
+    }
+
+    private static final class TestSystemDeviceLockManager implements SystemDeviceLockManager {
+
+        @Override
+        public void addFinancedDeviceKioskRole(@NonNull String packageName, Executor executor,
+                @NonNull OutcomeReceiver<Void, Exception> callback) {
+
+        }
+
+        @Override
+        public void removeFinancedDeviceKioskRole(@NonNull String packageName, Executor executor,
+                @NonNull OutcomeReceiver<Void, Exception> callback) {
+
+        }
+
+        @Override
+        public void setExemptFromActivityBackgroundStartRestriction(boolean exempt,
+                Executor executor, @NonNull OutcomeReceiver<Void, Exception> callback) {
+
+        }
+
+        @Override
+        public void setExemptFromHibernation(String packageName, boolean exempt, Executor executor,
+                @NonNull OutcomeReceiver<Void, Exception> callback) {
+
+        }
+
+        @Override
+        public void enableKioskKeepalive(String packageName, Executor executor,
+                @NonNull OutcomeReceiver<Void, Exception> callback) {
+
+        }
+
+        @Override
+        public void disableKioskKeepalive(Executor executor,
+                @NonNull OutcomeReceiver<Void, Exception> callback) {
+
+        }
+
+        @Override
+        public void setDeviceFinalized(boolean finalized, Executor executor,
+                @NonNull OutcomeReceiver<Void, Exception> callback) {
+            executor.execute(() -> callback.onResult(null));
+        }
     }
 
     /**
diff --git a/framework/java/android/devicelock/IDeviceLockService.aidl b/framework/java/android/devicelock/IDeviceLockService.aidl
index 67eec5e..e8ad6df 100644
--- a/framework/java/android/devicelock/IDeviceLockService.aidl
+++ b/framework/java/android/devicelock/IDeviceLockService.aidl
@@ -96,4 +96,9 @@
      * Disable kiosk keepalive.
      */
      void disableKioskKeepalive(in RemoteCallback remoteCallback);
+
+    /**
+     * Set device finalized.
+     */
+     void setDeviceFinalized(in boolean finalized, in RemoteCallback remoteCallback);
 }
diff --git a/service/Android.bp b/service/Android.bp
index ff7eb5c..fbae37b 100644
--- a/service/Android.bp
+++ b/service/Android.bp
@@ -38,6 +38,7 @@
     ],
     static_libs: [
         "devicelockcontroller-interface",
+        "devicelockcontroller-common-lib",
     ],
     apex_available: [
         "com.android.devicelock",
diff --git a/service/java/com/android/server/devicelock/DeviceLockPersistentStore.java b/service/java/com/android/server/devicelock/DeviceLockPersistentStore.java
new file mode 100644
index 0000000..d53c7f6
--- /dev/null
+++ b/service/java/com/android/server/devicelock/DeviceLockPersistentStore.java
@@ -0,0 +1,162 @@
+/*
+ * 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.server.devicelock;
+
+import android.annotation.WorkerThread;
+import android.os.Environment;
+import android.util.AtomicFile;
+import android.util.Slog;
+import android.util.Xml;
+
+import com.android.devicelockcontroller.util.ThreadUtils;
+
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+import org.xmlpull.v1.XmlSerializer;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+
+/**
+ * Class that manages persisting device state data for the system service.
+ */
+public final class DeviceLockPersistentStore {
+    private static final String TAG = DeviceLockPersistentStore.class.getSimpleName();
+    private static final String SYSTEM_DIR = "system";
+    private static final String DEVICE_LOCK_DIR = "device_lock";
+    private static final String DEVICE_STATE_FILE = "device_state.xml";
+    private static final String TAG_DEVICE_STATE = "device_state";
+    private static final String ATTR_IS_DEVICE_FINALIZED = "is_device_finalized";
+
+    private final Executor mBgExecutor = Executors.newSingleThreadExecutor();
+    private final File mFile;
+
+    DeviceLockPersistentStore() {
+        final File systemDir = new File(Environment.getDataDirectory(), SYSTEM_DIR);
+        final File deviceLockDir = new File(systemDir, DEVICE_LOCK_DIR);
+        if (!deviceLockDir.exists()) {
+            final boolean madeDirs = deviceLockDir.mkdirs();
+            if (!madeDirs) {
+                Slog.e(TAG, "Failed to make directory " + deviceLockDir.getAbsolutePath());
+            }
+        }
+        mFile = new File(deviceLockDir, DEVICE_STATE_FILE);
+    }
+
+    /**
+     * Schedule a write of the finalized state.
+     *
+     * @param finalized true if device is fully finalized
+     */
+    public void scheduleWrite(boolean finalized) {
+        mBgExecutor.execute(() -> writeState(finalized));
+    }
+
+    /**
+     * Read the finalized state from disk
+     *
+     * @param callback callback for when state is read
+     * @param callbackExecutor executor to run callback on
+     */
+    public void readFinalizedState(DeviceStateCallback callback, Executor callbackExecutor) {
+        mBgExecutor.execute(() -> {
+            final boolean isFinalized = readState();
+            callbackExecutor.execute(() -> callback.onDeviceStateRead(isFinalized));
+        });
+    }
+
+    @WorkerThread
+    private void writeState(boolean finalized) {
+        ThreadUtils.assertWorkerThread("writeState");
+        synchronized (this) {
+            AtomicFile atomicFile = new AtomicFile(mFile);
+
+            try (FileOutputStream fileOutputStream = atomicFile.startWrite()) {
+                try {
+                    XmlSerializer serializer = Xml.newSerializer();
+                    serializer.setOutput(fileOutputStream, Xml.Encoding.UTF_8.name());
+                    serializer.startDocument(Xml.Encoding.UTF_8.name(), true /* standalone */);
+                    writeToXml(serializer, finalized);
+                    serializer.endDocument();
+                    fileOutputStream.flush();
+                    atomicFile.finishWrite(fileOutputStream);
+                } catch (IOException e) {
+                    Slog.e(TAG, "Failed to write to XML", e);
+                    atomicFile.failWrite(fileOutputStream);
+                }
+            } catch (IOException e) {
+                Slog.e(TAG, "Failed to start write", e);
+            }
+        }
+    }
+
+    private void writeToXml(XmlSerializer serializer, boolean finalized) throws IOException {
+        serializer.startTag(null /* namespace */, TAG_DEVICE_STATE);
+        serializer.attribute(null /* namespace */,
+                ATTR_IS_DEVICE_FINALIZED, Boolean.toString(finalized));
+        serializer.endTag(null /* namespace */, TAG_DEVICE_STATE);
+    }
+
+    @WorkerThread
+    private boolean readState() {
+        ThreadUtils.assertWorkerThread("readState");
+        synchronized (this) {
+            if (!mFile.exists()) {
+                return false;
+            }
+            AtomicFile atomicFile = new AtomicFile(mFile);
+
+            try (FileInputStream inputStream = atomicFile.openRead()) {
+                XmlPullParser parser = Xml.newPullParser();
+                parser.setInput(inputStream, Xml.Encoding.UTF_8.name());
+                return getStateFromXml(parser);
+            } catch (XmlPullParserException | IOException e) {
+                Slog.e(TAG, "Failed to read XML", e);
+                return false;
+            }
+        }
+    }
+
+    private boolean getStateFromXml(XmlPullParser parser)
+            throws XmlPullParserException, IOException {
+        while (parser.getEventType() != XmlPullParser.START_TAG
+                || !TAG_DEVICE_STATE.equals(parser.getName())) {
+            if (parser.getEventType() == XmlPullParser.END_DOCUMENT) {
+                throw new XmlPullParserException("Malformed XML. Unable to find start of tag.");
+            }
+            parser.next();
+        }
+        return Boolean.parseBoolean(
+                parser.getAttributeValue(null /* namespace */, ATTR_IS_DEVICE_FINALIZED));
+    }
+
+    /**
+     * Callback for when state is read from disk.
+     */
+    interface DeviceStateCallback {
+        /**
+         * Callback for when state is finished reading from disk.
+         *
+         * @param isFinalized whether device is finalized
+         */
+        void onDeviceStateRead(boolean isFinalized);
+    }
+}
diff --git a/service/java/com/android/server/devicelock/DeviceLockService.java b/service/java/com/android/server/devicelock/DeviceLockService.java
index f4dbb96..c8fe75c 100644
--- a/service/java/com/android/server/devicelock/DeviceLockService.java
+++ b/service/java/com/android/server/devicelock/DeviceLockService.java
@@ -76,7 +76,7 @@
         Objects.requireNonNull(to);
         Slog.d(TAG, "onUserSwitching from: " + from + " to: " + to);
         final UserHandle userHandle = to.getUserHandle();
-        mImpl.setDeviceLockControllerPackageDefaultEnabledState(userHandle);
+        mImpl.enableDeviceLockControllerIfNeeded(userHandle);
         mImpl.onUserSwitching(userHandle);
     }
 
diff --git a/service/java/com/android/server/devicelock/DeviceLockServiceImpl.java b/service/java/com/android/server/devicelock/DeviceLockServiceImpl.java
index a83e3c1..7e6b9f6 100644
--- a/service/java/com/android/server/devicelock/DeviceLockServiceImpl.java
+++ b/service/java/com/android/server/devicelock/DeviceLockServiceImpl.java
@@ -85,6 +85,8 @@
     private final ArrayMap<Integer, KioskKeepaliveServiceConnection>
             mKioskKeepaliveServiceConnections;
 
+    private final DeviceLockPersistentStore mPersistentStore;
+
     // The following should be a SystemApi on AppOpsManager.
     private static final String OPSTR_SYSTEM_EXEMPT_FROM_ACTIVITY_BG_START_RESTRICTION =
             "android:system_exempt_from_activity_bg_start_restriction";
@@ -174,6 +176,8 @@
 
         mPackageUtils = new DeviceLockControllerPackageUtils(context);
 
+        mPersistentStore = new DeviceLockPersistentStore();
+
         final StringBuilder errorMessage = new StringBuilder();
         mServiceInfo = mPackageUtils.findService(errorMessage);
 
@@ -182,8 +186,7 @@
         }
 
         if (!mServiceInfo.applicationInfo.enabled) {
-            Slog.w(TAG, "Device Lock Controller is disabled");
-            setDeviceLockControllerPackageDefaultEnabledState(UserHandle.SYSTEM);
+            enableDeviceLockControllerIfNeeded(UserHandle.SYSTEM);
         }
 
         final ComponentName componentName = new ComponentName(mServiceInfo.packageName,
@@ -197,7 +200,15 @@
                 Context.RECEIVER_EXPORTED);
     }
 
-    void setDeviceLockControllerPackageDefaultEnabledState(@NonNull UserHandle userHandle) {
+    void enableDeviceLockControllerIfNeeded(@NonNull UserHandle userHandle) {
+        mPersistentStore.readFinalizedState(isFinalized -> {
+            if (!isFinalized) {
+                setDeviceLockControllerPackageDefaultEnabledState(userHandle);
+            }
+        }, mContext.getMainExecutor());
+    }
+
+    private void setDeviceLockControllerPackageDefaultEnabledState(UserHandle userHandle) {
         final String controllerPackageName = mServiceInfo.packageName;
 
         Context controllerContext;
@@ -695,4 +706,13 @@
         result.putBoolean(KEY_REMOTE_CALLBACK_RESULT, serviceConnection != null);
         remoteCallback.sendResult(result);
     }
+
+    @Override
+    public void setDeviceFinalized(boolean finalized, @NonNull RemoteCallback remoteCallback) {
+        mPersistentStore.scheduleWrite(finalized);
+
+        final Bundle result = new Bundle();
+        result.putBoolean(KEY_REMOTE_CALLBACK_RESULT, true);
+        remoteCallback.sendResult(result);
+    }
 }