Merge "Add multi client ror support" am: fa7173b8ac

Original change: https://android-review.googlesource.com/c/platform/frameworks/base/+/1483936

MUST ONLY BE SUBMITTED BY AUTOMERGER

Change-Id: Id2c73984af2b0db71794f0f029daac02e9f1529a

core/java/android/os/IRecoverySystem.aidl

@@ -27,7 +27,8 @@ interface IRecoverySystem {
     boolean setupBcb(in String command);
     boolean clearBcb();
     void rebootRecoveryWithCommand(in String command);
-    boolean requestLskf(in String updateToken, in IntentSender sender);
-    boolean clearLskf();
-    boolean rebootWithLskf(in String updateToken, in String reason);
+    boolean requestLskf(in String packageName, in IntentSender sender);
+    boolean clearLskf(in String packageName);
+    boolean isLskfCaptured(in String packageName);
+    boolean rebootWithLskf(in String packageName, in String reason, in boolean slotSwitch);
 }

core/java/android/os/RecoverySystem.java

@@ -632,17 +632,11 @@ public class RecoverySystem {
      * Prepare to apply an unattended update by asking the user for their Lock Screen Knowledge
      * Factor (LSKF). If supplied, the {@code intentSender} will be called when the system is setup
      * and ready to apply the OTA.
-     * <p>
-     * When the system is already prepared for update and this API is called again with the same
-     * {@code updateToken}, it will not call the intent sender nor request the user enter their Lock
-     * Screen Knowledge Factor.
-     * <p>
-     * When this API is called again with a different {@code updateToken}, the prepared-for-update
-     * status is reset and process repeats as though it's the initial call to this method as
-     * described in the first paragraph.
      *
      * @param context the Context to use.
-     * @param updateToken token used to indicate which update was prepared
+     * @param updateToken this parameter is deprecated and won't be used. See details in
+     *                    <a href="http://go/multi-client-ror">http://go/multi-client-ror</a>
+     *                    TODO(xunchang) update the link of document with the public doc.
      * @param intentSender the intent to call when the update is prepared; may be {@code null}
      * @throws IOException if there were any errors setting up unattended update
      * @hide
@@ -655,7 +649,7 @@ public class RecoverySystem {
             throw new NullPointerException("updateToken == null");
         }
         RecoverySystem rs = (RecoverySystem) context.getSystemService(Context.RECOVERY_SERVICE);
-        if (!rs.requestLskf(updateToken, intentSender)) {
+        if (!rs.requestLskf(context.getPackageName(), intentSender)) {
             throw new IOException("preparation for update failed");
         }
     }
@@ -673,18 +667,18 @@ public class RecoverySystem {
     public static void clearPrepareForUnattendedUpdate(@NonNull Context context)
             throws IOException {
         RecoverySystem rs = (RecoverySystem) context.getSystemService(Context.RECOVERY_SERVICE);
-        if (!rs.clearLskf()) {
+        if (!rs.clearLskf(context.getPackageName())) {
             throw new IOException("could not reset unattended update state");
         }
     }
 
     /**
-     * Request that the device reboot and apply the update that has been prepared. The
-     * {@code updateToken} must match what was given for {@link #prepareForUnattendedUpdate} or
-     * this will return {@code false}.
+     * Request that the device reboot and apply the update that has been prepared.
      *
      * @param context the Context to use.
-     * @param updateToken the token used to call {@link #prepareForUnattendedUpdate} before
+     * @param updateToken this parameter is deprecated and won't be used. See details in
+     *                    <a href="http://go/multi-client-ror">http://go/multi-client-ror</a>
+     *                    TODO(xunchang) update the link of document with the public doc.
      * @param reason the reboot reason to give to the {@link PowerManager}
      * @throws IOException if the reboot couldn't proceed because the device wasn't ready for an
      *               unattended reboot or if the {@code updateToken} did not match the previously
@@ -699,7 +693,8 @@ public class RecoverySystem {
             throw new NullPointerException("updateToken == null");
         }
         RecoverySystem rs = (RecoverySystem) context.getSystemService(Context.RECOVERY_SERVICE);
-        if (!rs.rebootWithLskf(updateToken, reason)) {
+        // OTA is the sole user before S, and a slot switch is required for ota update.
+        if (!rs.rebootWithLskf(context.getPackageName(), reason, true)) {
             throw new IOException("system not prepared to apply update");
         }
     }
@@ -1283,16 +1278,15 @@ public class RecoverySystem {
     /**
      * Begins the process of asking the user for the Lock Screen Knowledge Factor.
      *
-     * @param updateToken token that will be used in calls to {@link #rebootAndApply} to ensure
-     *                    that the preparation was for the correct update
+     * @param packageName the package name of the caller who requests resume on reboot
      * @return true if the request was correct
      * @throws IOException if the recovery system service could not be contacted
      */
-    private boolean requestLskf(String updateToken, IntentSender sender) throws IOException {
+    private boolean requestLskf(String packageName, IntentSender sender) throws IOException {
         try {
-            return mService.requestLskf(updateToken, sender);
+            return mService.requestLskf(packageName, sender);
         } catch (RemoteException e) {
-            throw new IOException("could request update");
+            throw new IOException("could request LSKF capture");
         }
     }
 
@@ -1302,22 +1296,37 @@ public class RecoverySystem {
      * @return true if the setup for OTA was cleared
      * @throws IOException if the recovery system service could not be contacted
      */
-    private boolean clearLskf() throws IOException {
+    private boolean clearLskf(String packageName) throws IOException {
         try {
-            return mService.clearLskf();
+            return mService.clearLskf(packageName);
         } catch (RemoteException e) {
             throw new IOException("could not clear LSKF");
         }
     }
 
+    /**
+     * Queries if the resume on reboot has been prepared for a given caller.
+     *
+     * @param packageName the identifier of the caller who requests resume on reboot
+     * @return true if resume on reboot is prepared.
+     * @throws IOException if the recovery system service could not be contacted
+     */
+    private boolean isLskfCaptured(String packageName) throws IOException {
+        try {
+            return mService.isLskfCaptured(packageName);
+        } catch (RemoteException e) {
+            throw new IOException("could not get LSKF capture state");
+        }
+    }
+
     /**
      * Calls the recovery system service to reboot and apply update.
      *
-     * @param updateToken the update token for which the update was prepared
      */
-    private boolean rebootWithLskf(String updateToken, String reason) throws IOException {
+    private boolean rebootWithLskf(String packageName, String reason, boolean slotSwitch)
+            throws IOException {
         try {
-            return mService.rebootWithLskf(updateToken, reason);
+            return mService.rebootWithLskf(packageName, reason, slotSwitch);
         } catch (RemoteException e) {
             throw new IOException("could not reboot for update");
         }

services/core/java/com/android/server/recoverysystem/RecoverySystemService.java

@@ -16,8 +16,10 @@
 
 package com.android.server.recoverysystem;
 
+import android.annotation.IntDef;
 import android.content.Context;
 import android.content.IntentSender;
+import android.content.pm.PackageManager;
 import android.net.LocalSocket;
 import android.net.LocalSocketAddress;
 import android.os.Binder;
@@ -32,6 +34,7 @@ import android.os.ShellCallback;
 import android.os.SystemProperties;
 import android.util.Slog;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.widget.LockSettingsInternal;
 import com.android.internal.widget.RebootEscrowListener;
@@ -46,6 +49,10 @@ import java.io.FileDescriptor;
 import java.io.FileWriter;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * The recovery system service is responsible for coordinating recovery related
@@ -76,9 +83,53 @@ public class RecoverySystemService extends IRecoverySystem.Stub implements Reboo
     private final Injector mInjector;
     private final Context mContext;
 
-    private boolean mPreparedForReboot;
-    private String mUnattendedRebootToken;
-    private IntentSender mPreparedForRebootIntentSender;
+    @GuardedBy("this")
+    private final Map<String, IntentSender> mCallerPendingRequest = new HashMap<>();
+    @GuardedBy("this")
+    private final Set<String> mCallerPreparedForReboot = new HashSet<>();
+
+    /**
+     * Need to prepare for resume on reboot.
+     */
+    private static final int ROR_NEED_PREPARATION = 0;
+    /**
+     * Resume on reboot has been prepared, notify the caller.
+     */
+    private static final int ROR_SKIP_PREPARATION_AND_NOTIFY = 1;
+    /**
+     * Resume on reboot has been requested. Caller won't be notified until the preparation is done.
+     */
+    private static final int ROR_SKIP_PREPARATION_NOT_NOTIFY = 2;
+
+    /**
+     * The caller never requests for resume on reboot, no need for clear.
+     */
+    private static final int ROR_NOT_REQUESTED = 0;
+    /**
+     * Clear the resume on reboot preparation state.
+     */
+    private static final int ROR_REQUESTED_NEED_CLEAR = 1;
+    /**
+     * The caller has requested for resume on reboot. No need for clear since other callers may
+     * exist.
+     */
+    private static final int ROR_REQUESTED_SKIP_CLEAR = 2;
+
+    /**
+     * The action to perform upon new resume on reboot prepare request for a given client.
+     */
+    @IntDef({ ROR_NEED_PREPARATION,
+            ROR_SKIP_PREPARATION_AND_NOTIFY,
+            ROR_SKIP_PREPARATION_NOT_NOTIFY })
+    @interface ResumeOnRebootActionsOnRequest {}
+
+    /**
+     * The action to perform upon resume on reboot clear request for a given client.
+     */
+    @IntDef({ROR_NOT_REQUESTED,
+            ROR_REQUESTED_NEED_CLEAR,
+            ROR_REQUESTED_SKIP_CLEAR})
+    @interface ResumeOnRebootActionsOnClear{}
 
     static class Injector {
         protected final Context mContext;
@@ -286,47 +337,92 @@ public class RecoverySystemService extends IRecoverySystem.Stub implements Reboo
         }
     }
 
-    @Override // Binder call
-    public boolean requestLskf(String updateToken, IntentSender intentSender) {
-        mContext.enforceCallingOrSelfPermission(android.Manifest.permission.RECOVERY, null);
+    private void enforcePermissionForResumeOnReboot() {
+        if (mContext.checkCallingOrSelfPermission(android.Manifest.permission.RECOVERY)
+                != PackageManager.PERMISSION_GRANTED
+                && mContext.checkCallingOrSelfPermission(android.Manifest.permission.REBOOT)
+                        != PackageManager.PERMISSION_GRANTED) {
+            throw new SecurityException("Caller or self must have "
+                    + android.Manifest.permission.RECOVERY + " or "
+                    + android.Manifest.permission.REBOOT + " for resume on reboot.");
+        }
+    }
 
-        if (updateToken == null) {
+    @Override // Binder call
+    public boolean requestLskf(String packageName, IntentSender intentSender) {
+        enforcePermissionForResumeOnReboot();
+
+        if (packageName == null) {
+            Slog.w(TAG, "Missing packageName when requesting lskf.");
             return false;
         }
 
-        // No need to prepare again for the same token.
-        if (mPreparedForReboot && updateToken.equals(mUnattendedRebootToken)) {
-            return true;
+        @ResumeOnRebootActionsOnRequest int action = updateRoRPreparationStateOnNewRequest(
+                packageName, intentSender);
+        switch (action) {
+            case ROR_SKIP_PREPARATION_AND_NOTIFY:
+                // We consider the preparation done if someone else has prepared.
+                sendPreparedForRebootIntentIfNeeded(intentSender);
+                return true;
+            case ROR_SKIP_PREPARATION_NOT_NOTIFY:
+                return true;
+            case ROR_NEED_PREPARATION:
+                final long origId = Binder.clearCallingIdentity();
+                try {
+                    mInjector.getLockSettingsService().prepareRebootEscrow();
+                } finally {
+                    Binder.restoreCallingIdentity(origId);
+                }
+                return true;
+            default:
+                throw new IllegalStateException("Unsupported action type on new request " + action);
+        }
+    }
+
+    // Checks and updates the resume on reboot preparation state.
+    private synchronized @ResumeOnRebootActionsOnRequest int updateRoRPreparationStateOnNewRequest(
+            String packageName, IntentSender intentSender) {
+        if (!mCallerPreparedForReboot.isEmpty()) {
+            if (mCallerPreparedForReboot.contains(packageName)) {
+                Slog.i(TAG, "RoR already has prepared for " + packageName);
+            }
+
+            // Someone else has prepared. Consider the preparation done, and send back the intent.
+            mCallerPreparedForReboot.add(packageName);
+            return ROR_SKIP_PREPARATION_AND_NOTIFY;
         }
 
-        mPreparedForReboot = false;
-        mUnattendedRebootToken = updateToken;
-        mPreparedForRebootIntentSender = intentSender;
-
-        final long origId = Binder.clearCallingIdentity();
-        try {
-            mInjector.getLockSettingsService().prepareRebootEscrow();
-        } finally {
-            Binder.restoreCallingIdentity(origId);
+        boolean needPreparation = mCallerPendingRequest.isEmpty();
+        if (mCallerPendingRequest.containsKey(packageName)) {
+            Slog.i(TAG, "Duplicate RoR preparation request for " + packageName);
         }
-
-        return true;
+        // Update the request with the new intentSender.
+        mCallerPendingRequest.put(packageName, intentSender);
+        return needPreparation ? ROR_NEED_PREPARATION : ROR_SKIP_PREPARATION_NOT_NOTIFY;
     }
 
     @Override
     public void onPreparedForReboot(boolean ready) {
-        if (mUnattendedRebootToken == null) {
-            Slog.w(TAG, "onPreparedForReboot called when mUnattendedRebootToken is null");
-        }
-
-        mPreparedForReboot = ready;
-        if (ready) {
-            sendPreparedForRebootIntentIfNeeded();
+        if (!ready) {
+            return;
         }
+        updateRoRPreparationStateOnPreparedForReboot();
     }
 
-    private void sendPreparedForRebootIntentIfNeeded() {
-        final IntentSender intentSender = mPreparedForRebootIntentSender;
+    private synchronized void updateRoRPreparationStateOnPreparedForReboot() {
+        if (!mCallerPreparedForReboot.isEmpty()) {
+            Slog.w(TAG, "onPreparedForReboot called when some clients have prepared.");
+        }
+
+        // Send intents to notify callers
+        for (Map.Entry<String, IntentSender> entry : mCallerPendingRequest.entrySet()) {
+            sendPreparedForRebootIntentIfNeeded(entry.getValue());
+            mCallerPreparedForReboot.add(entry.getKey());
+        }
+        mCallerPendingRequest.clear();
+    }
+
+    private void sendPreparedForRebootIntentIfNeeded(IntentSender intentSender) {
         if (intentSender != null) {
             try {
                 intentSender.sendIntent(null, 0, null, null, null);
@@ -337,37 +433,61 @@ public class RecoverySystemService extends IRecoverySystem.Stub implements Reboo
     }
 
     @Override // Binder call
-    public boolean clearLskf() {
-        mContext.enforceCallingOrSelfPermission(android.Manifest.permission.RECOVERY, null);
-
-        mPreparedForReboot = false;
-        mUnattendedRebootToken = null;
-        mPreparedForRebootIntentSender = null;
-
-        final long origId = Binder.clearCallingIdentity();
-        try {
-            mInjector.getLockSettingsService().clearRebootEscrow();
-        } finally {
-            Binder.restoreCallingIdentity(origId);
+    public boolean clearLskf(String packageName) {
+        enforcePermissionForResumeOnReboot();
+        if (packageName == null) {
+            Slog.w(TAG, "Missing packageName when clearing lskf.");
+            return false;
         }
 
-        return true;
+        @ResumeOnRebootActionsOnClear int action = updateRoRPreparationStateOnClear(packageName);
+        switch (action) {
+            case ROR_NOT_REQUESTED:
+                Slog.w(TAG, "RoR clear called before preparation for caller " + packageName);
+                return true;
+            case ROR_REQUESTED_SKIP_CLEAR:
+                return true;
+            case ROR_REQUESTED_NEED_CLEAR:
+                final long origId = Binder.clearCallingIdentity();
+                try {
+                    mInjector.getLockSettingsService().clearRebootEscrow();
+                } finally {
+                    Binder.restoreCallingIdentity(origId);
+                }
+                return true;
+            default:
+                throw new IllegalStateException("Unsupported action type on clear " + action);
+        }
+    }
+
+    private synchronized @ResumeOnRebootActionsOnClear int updateRoRPreparationStateOnClear(
+            String packageName) {
+        if (!mCallerPreparedForReboot.contains(packageName) && !mCallerPendingRequest.containsKey(
+                packageName)) {
+            Slog.w(TAG, packageName + " hasn't prepared for resume on reboot");
+            return ROR_NOT_REQUESTED;
+        }
+        mCallerPendingRequest.remove(packageName);
+        mCallerPreparedForReboot.remove(packageName);
+
+        // Check if others have prepared ROR.
+        boolean needClear = mCallerPendingRequest.isEmpty() && mCallerPreparedForReboot.isEmpty();
+        return needClear ? ROR_REQUESTED_NEED_CLEAR : ROR_REQUESTED_SKIP_CLEAR;
     }
 
     @Override // Binder call
-    public boolean rebootWithLskf(String updateToken, String reason) {
-        mContext.enforceCallingOrSelfPermission(android.Manifest.permission.RECOVERY, null);
-
-        if (!mPreparedForReboot) {
-            Slog.i(TAG, "Reboot requested before prepare completed");
-            return false;
-        }
-
-        if (updateToken != null && !updateToken.equals(mUnattendedRebootToken)) {
-            Slog.i(TAG, "Reboot requested after preparation, but with mismatching token");
+    public boolean rebootWithLskf(String packageName, String reason, boolean slotSwitch) {
+        enforcePermissionForResumeOnReboot();
+        if (packageName == null) {
+            Slog.w(TAG, "Missing packageName when rebooting with lskf.");
+            return false;
+        }
+        if (!isLskfCaptured(packageName)) {
             return false;
         }
 
+        // TODO(xunchang) check the slot to boot into, and fail the reboot upon slot mismatch.
+        // TODO(xunchang) write the vbmeta digest along with the escrowKey before reboot.
         if (!mInjector.getLockSettingsService().armRebootEscrow()) {
             Slog.w(TAG, "Failure to escrow key for reboot");
             return false;
@@ -378,6 +498,16 @@ public class RecoverySystemService extends IRecoverySystem.Stub implements Reboo
         return true;
     }
 
+    @Override // Binder call
+    public synchronized boolean isLskfCaptured(String packageName) {
+        enforcePermissionForResumeOnReboot();
+        if (!mCallerPreparedForReboot.contains(packageName)) {
+            Slog.i(TAG, "Reboot requested before prepare completed for caller " + packageName);
+            return false;
+        }
+        return true;
+    }
+
     /**
      * Check if any of the init services is still running. If so, we cannot
      * start a new uncrypt/setup-bcb/clear-bcb service right away; otherwise

services/core/java/com/android/server/recoverysystem/RecoverySystemShellCommand.java

@@ -56,26 +56,31 @@ public class RecoverySystemShellCommand extends ShellCommand {
     }
 
     private int requestLskf() throws RemoteException {
-        String updateToken = getNextArgRequired();
-        boolean success = mService.requestLskf(updateToken, null);
+        String packageName = getNextArgRequired();
+        boolean success = mService.requestLskf(packageName, null);
         PrintWriter pw = getOutPrintWriter();
-        pw.println("Request LSKF status: " + (success ? "success" : "failure"));
+        pw.printf("Request LSKF for packageName: %s, status: %s\n", packageName,
+                success ? "success" : "failure");
         return 0;
     }
 
     private int clearLskf() throws RemoteException {
-        boolean success = mService.clearLskf();
+        String packageName = getNextArgRequired();
+        boolean success = mService.clearLskf(packageName);
         PrintWriter pw = getOutPrintWriter();
-        pw.println("Clear LSKF: " + (success ? "success" : "failure"));
+        pw.printf("Clear LSKF for packageName: %s, status: %s\n", packageName,
+                success ? "success" : "failure");
         return 0;
     }
 
     private int rebootAndApply() throws RemoteException {
-        String updateToken = getNextArgRequired();
+        String packageName = getNextArgRequired();
         String rebootReason = getNextArgRequired();
-        boolean success = mService.rebootWithLskf(updateToken, rebootReason);
+        boolean success = mService.rebootWithLskf(packageName, rebootReason, true);
         PrintWriter pw = getOutPrintWriter();
-        pw.println("Reboot and apply status: " + (success ? "success" : "failure"));
+        // Keep the old message for cts test.
+        pw.printf("%s Reboot and apply status: %s\n", packageName,
+                success ? "success" : "failure");
         return 0;
     }
 

services/tests/servicestests/src/com/android/server/recoverysystem/RecoverySystemServiceTest.java

@@ -21,6 +21,7 @@ import static org.junit.Assert.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doThrow;
@@ -33,6 +34,7 @@ import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.content.IntentSender;
+import android.content.pm.PackageManager;
 import android.os.Handler;
 import android.os.IPowerManager;
 import android.os.IRecoverySystemProgressListener;
@@ -67,6 +69,9 @@ public class RecoverySystemServiceTest {
     private FileWriter mUncryptUpdateFileWriter;
     private LockSettingsInternal mLockSettingsInternal;
 
+    private static final String FAKE_OTA_PACKAGE_NAME = "fake.ota.package";
+    private static final String FAKE_OTHER_PACKAGE_NAME = "fake.other.package";
+
     @Before
     public void setup() {
         mContext = mock(Context.class);
@@ -209,65 +214,99 @@ public class RecoverySystemServiceTest {
 
     @Test(expected = SecurityException.class)
     public void requestLskf_protected() {
-        doThrow(SecurityException.class).when(mContext).enforceCallingOrSelfPermission(
-                eq(android.Manifest.permission.RECOVERY), any());
-        mRecoverySystemService.requestLskf("test", null);
-    }
-
-
-    @Test
-    public void requestLskf_nullToken_failure() {
-        assertThat(mRecoverySystemService.requestLskf(null, null), is(false));
+        when(mContext.checkCallingOrSelfPermission(anyString())).thenReturn(
+                PackageManager.PERMISSION_DENIED);
+        mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, null);
     }
 
     @Test
     public void requestLskf_success() throws Exception {
         IntentSender intentSender = mock(IntentSender.class);
-        assertThat(mRecoverySystemService.requestLskf("test", intentSender), is(true));
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, intentSender),
+                is(true));
         mRecoverySystemService.onPreparedForReboot(true);
         verify(intentSender).sendIntent(any(), anyInt(), any(), any(), any());
     }
 
     @Test
-    public void requestLskf_subsequentRequestClearsPrepared() throws Exception {
+    public void requestLskf_subsequentRequestNotClearPrepared() throws Exception {
         IntentSender intentSender = mock(IntentSender.class);
-        assertThat(mRecoverySystemService.requestLskf("test", intentSender), is(true));
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, intentSender),
+                is(true));
         mRecoverySystemService.onPreparedForReboot(true);
         verify(intentSender).sendIntent(any(), anyInt(), any(), any(), any());
 
-        assertThat(mRecoverySystemService.requestLskf("test2", null), is(true));
-        assertThat(mRecoverySystemService.rebootWithLskf("test", null), is(false));
-        assertThat(mRecoverySystemService.rebootWithLskf("test2", "foobar"), is(false));
-
-        mRecoverySystemService.onPreparedForReboot(true);
-        assertThat(mRecoverySystemService.rebootWithLskf("test2", "foobar"), is(true));
-        verify(intentSender).sendIntent(any(), anyInt(), any(), any(), any());
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, null), is(true));
+        assertThat(mRecoverySystemService.rebootWithLskf(FAKE_OTA_PACKAGE_NAME, "foobar", true),
+                is(true));
         verify(mIPowerManager).reboot(anyBoolean(), eq("foobar"), anyBoolean());
     }
 
-
     @Test
     public void requestLskf_requestedButNotPrepared() throws Exception {
         IntentSender intentSender = mock(IntentSender.class);
-        assertThat(mRecoverySystemService.requestLskf("test", intentSender), is(true));
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, intentSender),
+                is(true));
         verify(intentSender, never()).sendIntent(any(), anyInt(), any(), any(), any());
     }
 
+    @Test
+    public void isLskfCaptured_requestedButNotPrepared() throws Exception {
+        IntentSender intentSender = mock(IntentSender.class);
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, intentSender),
+                is(true));
+        assertThat(mRecoverySystemService.isLskfCaptured(FAKE_OTA_PACKAGE_NAME), is(false));
+    }
+
+    @Test
+    public void isLskfCaptured_Prepared() throws Exception {
+        IntentSender intentSender = mock(IntentSender.class);
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, intentSender),
+                is(true));
+        mRecoverySystemService.onPreparedForReboot(true);
+        verify(intentSender).sendIntent(any(), anyInt(), any(), any(), any());
+        assertThat(mRecoverySystemService.isLskfCaptured(FAKE_OTA_PACKAGE_NAME), is(true));
+    }
+
     @Test(expected = SecurityException.class)
     public void clearLskf_protected() {
-        doThrow(SecurityException.class).when(mContext).enforceCallingOrSelfPermission(
-                eq(android.Manifest.permission.RECOVERY), any());
-        mRecoverySystemService.clearLskf();
+        when(mContext.checkCallingOrSelfPermission(anyString())).thenReturn(
+                PackageManager.PERMISSION_DENIED);
+        mRecoverySystemService.clearLskf(FAKE_OTA_PACKAGE_NAME);
     }
 
     @Test
     public void clearLskf_requestedThenCleared() throws Exception {
         IntentSender intentSender = mock(IntentSender.class);
-        assertThat(mRecoverySystemService.requestLskf("test", intentSender), is(true));
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, intentSender),
+                is(true));
         mRecoverySystemService.onPreparedForReboot(true);
         verify(intentSender).sendIntent(any(), anyInt(), any(), any(), any());
 
-        assertThat(mRecoverySystemService.clearLskf(), is(true));
+        assertThat(mRecoverySystemService.clearLskf(FAKE_OTA_PACKAGE_NAME), is(true));
+        verify(mLockSettingsInternal).clearRebootEscrow();
+    }
+
+    @Test
+    public void clearLskf_callerNotRequested_Success() throws Exception {
+        IntentSender intentSender = mock(IntentSender.class);
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, intentSender),
+                is(true));
+        assertThat(mRecoverySystemService.clearLskf(FAKE_OTHER_PACKAGE_NAME), is(true));
+        verify(mLockSettingsInternal, never()).clearRebootEscrow();
+    }
+
+    @Test
+    public void clearLskf_multiClient_BothClientsClear() throws Exception {
+        IntentSender intentSender = mock(IntentSender.class);
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, intentSender),
+                is(true));
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTHER_PACKAGE_NAME, intentSender),
+                is(true));
+
+        assertThat(mRecoverySystemService.clearLskf(FAKE_OTA_PACKAGE_NAME), is(true));
+        verify(mLockSettingsInternal, never()).clearRebootEscrow();
+        assertThat(mRecoverySystemService.clearLskf(FAKE_OTHER_PACKAGE_NAME), is(true));
         verify(mLockSettingsInternal).clearRebootEscrow();
     }
 
@@ -279,27 +318,84 @@ public class RecoverySystemServiceTest {
 
     @Test(expected = SecurityException.class)
     public void rebootWithLskf_protected() {
-        doThrow(SecurityException.class).when(mContext).enforceCallingOrSelfPermission(
-                eq(android.Manifest.permission.RECOVERY), any());
-        mRecoverySystemService.rebootWithLskf("test1", null);
+        when(mContext.checkCallingOrSelfPermission(anyString())).thenReturn(
+                PackageManager.PERMISSION_DENIED);
+        mRecoverySystemService.rebootWithLskf(FAKE_OTA_PACKAGE_NAME, null, true);
     }
 
     @Test
     public void rebootWithLskf_Success() throws Exception {
-        assertThat(mRecoverySystemService.requestLskf("test", null), is(true));
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, null), is(true));
         mRecoverySystemService.onPreparedForReboot(true);
-        assertThat(mRecoverySystemService.rebootWithLskf("test", "ab-update"), is(true));
+        assertThat(mRecoverySystemService.rebootWithLskf(FAKE_OTA_PACKAGE_NAME, "ab-update", true),
+                is(true));
         verify(mIPowerManager).reboot(anyBoolean(), eq("ab-update"), anyBoolean());
     }
 
     @Test
     public void rebootWithLskf_withoutPrepare_Failure() throws Exception {
-        assertThat(mRecoverySystemService.rebootWithLskf("test1", null), is(false));
+        assertThat(mRecoverySystemService.rebootWithLskf(FAKE_OTA_PACKAGE_NAME, null, true),
+                is(false));
     }
 
     @Test
-    public void rebootWithLskf_withNullUpdateToken_Failure() throws Exception {
-        assertThat(mRecoverySystemService.rebootWithLskf(null, null), is(false));
+    public void rebootWithLskf_withNullCallerId_Failure() throws Exception {
+        assertThat(mRecoverySystemService.rebootWithLskf(null, null, true), is(false));
         verifyNoMoreInteractions(mIPowerManager);
     }
+
+    @Test
+    public void rebootWithLskf_multiClient_ClientASuccess() throws Exception {
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, null), is(true));
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTHER_PACKAGE_NAME, null), is(true));
+        mRecoverySystemService.onPreparedForReboot(true);
+
+        // Client B's clear won't affect client A's preparation.
+        assertThat(mRecoverySystemService.clearLskf(FAKE_OTHER_PACKAGE_NAME), is(true));
+        assertThat(mRecoverySystemService.rebootWithLskf(FAKE_OTA_PACKAGE_NAME, "ab-update", true),
+                is(true));
+        verify(mIPowerManager).reboot(anyBoolean(), eq("ab-update"), anyBoolean());
+    }
+
+
+    @Test
+    public void rebootWithLskf_multiClient_ClientBSuccess() throws Exception {
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, null), is(true));
+        mRecoverySystemService.onPreparedForReboot(true);
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTHER_PACKAGE_NAME, null), is(true));
+
+        assertThat(mRecoverySystemService.clearLskf(FAKE_OTA_PACKAGE_NAME), is(true));
+        assertThat(mRecoverySystemService.rebootWithLskf(FAKE_OTA_PACKAGE_NAME, null, true),
+                is(false));
+        verifyNoMoreInteractions(mIPowerManager);
+
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTHER_PACKAGE_NAME, null), is(true));
+        assertThat(
+                mRecoverySystemService.rebootWithLskf(FAKE_OTHER_PACKAGE_NAME, "ab-update", true),
+                is(true));
+        verify(mIPowerManager).reboot(anyBoolean(), eq("ab-update"), anyBoolean());
+    }
+
+    @Test
+    public void rebootWithLskf_multiClient_BothClientsClear_Failure() throws Exception {
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTA_PACKAGE_NAME, null), is(true));
+        mRecoverySystemService.onPreparedForReboot(true);
+        assertThat(mRecoverySystemService.requestLskf(FAKE_OTHER_PACKAGE_NAME, null), is(true));
+
+        // Client A clears
+        assertThat(mRecoverySystemService.clearLskf(FAKE_OTA_PACKAGE_NAME), is(true));
+        assertThat(mRecoverySystemService.rebootWithLskf(FAKE_OTA_PACKAGE_NAME, null, true),
+                is(false));
+        verifyNoMoreInteractions(mIPowerManager);
+
+        // Client B clears
+        assertThat(mRecoverySystemService.clearLskf(FAKE_OTHER_PACKAGE_NAME), is(true));
+        verify(mLockSettingsInternal).clearRebootEscrow();
+        assertThat(
+                mRecoverySystemService.rebootWithLskf(FAKE_OTHER_PACKAGE_NAME, "ab-update", true),
+                is(false));
+        verifyNoMoreInteractions(mIPowerManager);
+    }
+
+    // TODO(xunchang) add more multi client tests
 }