Snap for 8589293 from 0fd786ab16e464b3b450a50aa6f510430e3be94d to sc-v2-platform-release am: 469539ffc0android12L-gsi

Original change: https://googleplex-android-review.googlesource.com/c/platform/packages/apps/RemoteProvisioner/+/18439732

Change-Id: Icfc7bb3fd1853889f3e1540c4e82799194b43a85
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

6 files changed, 528 insertions, 101 deletions

diff --git a/src/com/android/remoteprovisioner/PeriodicProvisioner.java b/src/com/android/remoteprovisioner/PeriodicProvisioner.java
index 1c2bc88..026b399 100644
--- a/src/com/android/remoteprovisioner/PeriodicProvisioner.java
+++ b/src/com/android/remoteprovisioner/PeriodicProvisioner.java
@@ -154,19 +154,10 @@ public class PeriodicProvisioner extends JobService {
                 for (int i = 0; i < implInfos.length; i++) {
                     // Break very large CSR requests into chunks, so as not to overwhelm the
                     // backend.
-                    int keysToCertify = keysNeededForSecLevel[i];
-                    while (keysToCertify != 0) {
-                        int batchSize = min(keysToCertify, SAFE_CSR_BATCH_SIZE);
-                        Log.i(TAG, "Requesting " + batchSize + " keys to be provisioned.");
-                        Provisioner.provisionCerts(batchSize,
-                                                   implInfos[i].secLevel,
-                                                   resp.getGeekChain(implInfos[i].supportedCurve),
-                                                   resp.getChallenge(),
-                                                   binder,
-                                                   mContext);
-                        keysToCertify -= batchSize;
-                    }
-                }
+                    int keysToProvision = keysNeededForSecLevel[i];
+                    batchProvision(binder, mContext, keysToProvision, implInfos[i].secLevel,
+                            resp.getGeekChain(implInfos[i].supportedCurve), resp.getChallenge());
+            }
                 jobFinished(mParams, false /* wantsReschedule */);
             } catch (RemoteException e) {
                 jobFinished(mParams, false /* wantsReschedule */);
@@ -203,6 +194,7 @@ public class PeriodicProvisioner extends JobService {
             for (int i = 0; i < implInfos.length; i++) {
                 keysNeededForSecLevel[i] =
                         generateNumKeysNeeded(binder,
+                                   mContext,
                                    expiringBy,
                                    implInfos[i].secLevel);
                 if (keysNeededForSecLevel[i] > 0) {
@@ -211,60 +203,67 @@ public class PeriodicProvisioner extends JobService {
             }
             return provisioningNeeded;
         }
+    }
 
-        /**
-         * This method will generate and bundle up keys for signing to make sure that there will be
-         * enough keys available for use by the system when current keys expire.
-         *
-         * Enough keys is defined by checking how many keys are currently assigned to apps and
-         * generating enough keys to cover any expiring certificates plus a bit of buffer room
-         * defined by {@code sExtraSignedKeysAvailable}.
-         *
-         * This allows devices to dynamically resize their key pools as the user downloads and
-         * removes apps that may also use attestation.
-         */
-        private int generateNumKeysNeeded(IRemoteProvisioning binder, long expiringBy, int secLevel)
-                throws InterruptedException, RemoteException {
-            AttestationPoolStatus pool =
-                    SystemInterface.getPoolStatus(expiringBy, secLevel, binder);
-            if (pool == null) {
-                Log.e(TAG, "Failed to fetch pool status.");
-                return 0;
-            }
-            Log.i(TAG, "Pool status.\nTotal: " + pool.total
-                       + "\nAttested: " + pool.attested
-                       + "\nUnassigned: " + pool.unassigned
-                       + "\nExpiring: " + pool.expiring);
-            int unattestedKeys = pool.total - pool.attested;
-            int keysInUse = pool.attested - pool.unassigned;
-            int totalSignedKeys = keysInUse + SettingsManager.getExtraSignedKeysAvailable(mContext);
-            int generated;
-            // If nothing is expiring, and the amount of available unassigned keys is sufficient,
-            // then do nothing. Otherwise, generate the complete amount of totalSignedKeys. It will
-            // reduce network usage if the app just provisions an entire new batch in one go, rather
-            // than consistently grabbing just a few at a time as the expiration dates become
-            // misaligned.
-            if (pool.expiring < pool.unassigned && pool.attested >= totalSignedKeys) {
-                Log.i(TAG,
-                        "No keys expiring and the expected number of attested keys are available");
-                return 0;
-            }
-            for (generated = 0;
-                    generated + unattestedKeys < totalSignedKeys; generated++) {
-                SystemInterface.generateKeyPair(false /* isTestMode */, secLevel, binder);
-                // Prioritize provisioning if there are no keys available. No keys being available
-                // indicates that this is the first time a device is being brought online.
-                if (pool.total != 0) {
-                    Thread.sleep(KEY_GENERATION_PAUSE.toMillis());
-                }
-            }
-            if (totalSignedKeys > 0) {
-                Log.i(TAG, "Generated " + generated + " keys. "
-                        + (generated + unattestedKeys) + " keys are now available for signing.");
-                return generated + unattestedKeys;
-            }
-            Log.i(TAG, "No keys generated.");
+    public static void batchProvision(IRemoteProvisioning binder, Context context,
+                        int keysToProvision, int secLevel,
+                        byte[] geekChain, byte[] challenge)
+                        throws RemoteException {
+        while (keysToProvision != 0) {
+            int batchSize = min(keysToProvision, SAFE_CSR_BATCH_SIZE);
+            Log.i(TAG, "Requesting " + batchSize + " keys to be provisioned.");
+            Provisioner.provisionCerts(batchSize,
+                                    secLevel,
+                                    geekChain,
+                                    challenge,
+                                    binder,
+                                    context);
+            keysToProvision -= batchSize;
+        }
+    }
+
+    /**
+     * This method will generate and bundle up keys for signing to make sure that there will be
+     * enough keys available for use by the system when current keys expire.
+     *
+     * Enough keys is defined by checking how many keys are currently assigned to apps and
+     * generating enough keys to cover any expiring certificates plus a bit of buffer room
+     * defined by {@code sExtraSignedKeysAvailable}.
+     *
+     * This allows devices to dynamically resize their key pools as the user downloads and
+     * removes apps that may also use attestation.
+     */
+    public static int generateNumKeysNeeded(IRemoteProvisioning binder, Context context,
+                                            long expiringBy, int secLevel)
+            throws InterruptedException, RemoteException {
+        AttestationPoolStatus pool =
+                SystemInterface.getPoolStatus(expiringBy, secLevel, binder);
+        if (pool == null) {
+            Log.e(TAG, "Failed to fetch pool status.");
             return 0;
         }
+        Log.i(TAG, "Pool status.\nTotal: " + pool.total
+                   + "\nAttested: " + pool.attested
+                   + "\nUnassigned: " + pool.unassigned
+                   + "\nExpiring: " + pool.expiring);
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(
+                    pool, SettingsManager.getExtraSignedKeysAvailable(context));
+        if (!stats.provisioningNeeded) {
+            Log.i(TAG, "No provisioning needed.");
+            return 0;
+        }
+        Log.i(TAG, "Need to generate " + stats.keysToGenerate + " keys.");
+        int generated;
+        for (generated = 0; generated < stats.keysToGenerate; generated++) {
+            SystemInterface.generateKeyPair(false /* isTestMode */, secLevel, binder);
+            // Prioritize provisioning if there are no keys available. No keys being available
+            // indicates that this is the first time a device is being brought online.
+            if (pool.total != 0) {
+                Thread.sleep(KEY_GENERATION_PAUSE.toMillis());
+            }
+        }
+        Log.i(TAG, "Generated " + generated + " keys. " + stats.unattestedKeys
+                    + " keys were also available for signing previous to generation.");
+        return stats.idealTotalSignedKeys;
     }
 }
diff --git a/src/com/android/remoteprovisioner/StatsProcessor.java b/src/com/android/remoteprovisioner/StatsProcessor.java
new file mode 100644
index 0000000..85582ab
--- /dev/null
+++ b/src/com/android/remoteprovisioner/StatsProcessor.java
@@ -0,0 +1,73 @@
+/**
+ * Copyright (C) 2022 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.remoteprovisioner;
+
+import android.security.remoteprovisioning.AttestationPoolStatus;
+import android.util.Log;
+
+public class StatsProcessor {
+    public static final double LIMIT_SCALER = .4;
+
+    private static final String TAG = "RemoteProvisioningService.KeyPoolStats";
+
+    private StatsProcessor() {}
+
+    public static int calcMinUnassignedToTriggerProvisioning(int extraSignedKeysAvailable) {
+        return (int) Math.ceil(LIMIT_SCALER * extraSignedKeysAvailable);
+    }
+
+    /**
+     * Creates a PoolStats. Takes an {@Code AttestationPoolStatus} and calculates different
+     * pieces of status to inform the caller if any action needs to be taken to reprovision the pool
+     * and what action is needed in terms of keys to generate.
+     *
+     * @parameter pool the current status of the keypool in Keystore2
+     * @parameter extraSignedKeysAvailable how many extra attested keys should ideally be available
+     *                                     for assignment.
+     * @return the PoolStats object describing higher level info about the state of the key pool.
+     */
+    public static PoolStats processPool(
+            AttestationPoolStatus pool, int extraSignedKeysAvailable) {
+        PoolStats stats = new PoolStats();
+        stats.unattestedKeys = pool.total - pool.attested;
+        stats.keysInUse = pool.attested - pool.unassigned;
+        stats.idealTotalSignedKeys = stats.keysInUse + extraSignedKeysAvailable;
+        // If nothing is expiring, and the amount of available unassigned keys is sufficient,
+        // then do nothing. Otherwise, generate the complete amount of idealTotalSignedKeys. It will
+        // reduce network usage if the app just provisions an entire new batch in one go, rather
+        // than consistently grabbing just a few at a time as the expiration dates become
+        // misaligned.
+        stats.provisioningNeeded =
+                pool.unassigned - pool.expiring
+                <= calcMinUnassignedToTriggerProvisioning(extraSignedKeysAvailable);
+        if (!stats.provisioningNeeded) {
+            Log.i(TAG, "Sufficient keys are available, no CSR needed.");
+            stats.keysToGenerate = 0;
+        } else {
+            stats.keysToGenerate = Math.max(0, stats.idealTotalSignedKeys - stats.unattestedKeys);
+        }
+        return stats;
+    }
+
+    public static class PoolStats {
+        public int unattestedKeys;
+        public int keysInUse;
+        public int idealTotalSignedKeys;
+        public int keysToGenerate;
+        public boolean provisioningNeeded = true;
+    }
+}
diff --git a/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java b/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java
index ab2ec8d..268a8b5 100644
--- a/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java
+++ b/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java
@@ -29,18 +29,24 @@ import android.security.remoteprovisioning.IRemoteProvisioning;
 import android.util.Log;
 
 import com.android.remoteprovisioner.GeekResponse;
-import com.android.remoteprovisioner.Provisioner;
+import com.android.remoteprovisioner.PeriodicProvisioner;
 import com.android.remoteprovisioner.ServerInterface;
-import com.android.remoteprovisioner.SettingsManager;
+
+import java.time.Duration;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * Provides the implementation for IGenerateKeyService.aidl
  */
 public class GenerateRkpKeyService extends Service {
     private static final int KEY_GENERATION_PAUSE_MS = 1000;
+    private static final Duration LOOKAHEAD_TIME = Duration.ofDays(1);
+
     private static final String SERVICE = "android.security.remoteprovisioning";
     private static final String TAG = "RemoteProvisioningService";
 
+    private static final ReentrantLock sLock = new ReentrantLock();
+
     @Override
     public void onCreate() {
         super.onCreate();
@@ -78,41 +84,46 @@ public class GenerateRkpKeyService extends Service {
 
         private void checkAndFillPool(IRemoteProvisioning binder, int secLevel)
                 throws RemoteException {
-            AttestationPoolStatus pool =
-                    binder.getPoolStatus(System.currentTimeMillis(), secLevel);
-            ImplInfo[] implInfos = binder.getImplementationInfo();
-            int curve = 0;
-            for (int i = 0; i < implInfos.length; i++) {
-                if (implInfos[i].secLevel == secLevel) {
-                    curve = implInfos[i].supportedCurve;
-                    break;
-                }
+            // No need to hammer the pool check with a ton of redundant requests.
+            if (!sLock.tryLock()) {
+                Log.i(TAG, "Exiting check; another process already started the check.");
+                return;
             }
-            // If there are no unassigned keys, go ahead and provision some. If there are no
-            // attested keys at all on the system, this implies that it is a hybrid
-            // rkp/factory-provisioned system that has turned off RKP. In that case, do
-            // not provision.
-            if (pool.unassigned == 0 && pool.attested != 0) {
-                Log.i(TAG, "All signed keys are currently in use, provisioning more.");
-                Context context = getApplicationContext();
-                int keysToProvision = SettingsManager.getExtraSignedKeysAvailable(context);
-                int existingUnsignedKeys = pool.total - pool.attested;
-                int keysToGenerate = keysToProvision - existingUnsignedKeys;
-                try {
-                    for (int i = 0; i < keysToGenerate; i++) {
-                        binder.generateKeyPair(false /* isTestMode */, secLevel);
-                        Thread.sleep(KEY_GENERATION_PAUSE_MS);
+            try {
+                AttestationPoolStatus pool =
+                        binder.getPoolStatus(System.currentTimeMillis(), secLevel);
+                ImplInfo[] implInfos = binder.getImplementationInfo();
+                int curve = 0;
+                for (int i = 0; i < implInfos.length; i++) {
+                    if (implInfos[i].secLevel == secLevel) {
+                        curve = implInfos[i].supportedCurve;
+                        break;
                     }
-                } catch (InterruptedException e) {
-                    Log.i(TAG, "Thread interrupted", e);
                 }
-                GeekResponse resp = ServerInterface.fetchGeek(context);
-                if (resp == null) {
-                    Log.e(TAG, "Server unavailable");
-                    return;
+
+                Context context = getApplicationContext();
+                int keysToProvision =
+                        PeriodicProvisioner.generateNumKeysNeeded(binder, context,
+                                                                  LOOKAHEAD_TIME.toMillis(),
+                                                                  secLevel);
+                // If there are no unassigned keys, go ahead and provision some. If there are no
+                // attested keys at all on the system, this implies that it is a hybrid
+                // rkp/factory-provisioned system that has turned off RKP. In that case, do
+                // not provision.
+                if (keysToProvision != 0 && pool.attested != 0) {
+                    Log.i(TAG, "All signed keys are currently in use, provisioning more.");
+                    GeekResponse resp = ServerInterface.fetchGeek(context);
+                    if (resp == null) {
+                        Log.e(TAG, "Server unavailable");
+                        return;
+                    }
+                    PeriodicProvisioner.batchProvision(binder, context, keysToProvision, secLevel,
+                                                     resp.getGeekChain(curve), resp.getChallenge());
                 }
-                Provisioner.provisionCerts(keysToProvision, secLevel, resp.getGeekChain(curve),
-                                           resp.getChallenge(), binder, context);
+            } catch (InterruptedException e) {
+                Log.e(TAG, "Provisioner thread interrupted.", e);
+            } finally {
+                sLock.unlock();
             }
         }
     };
diff --git a/tests/unittests/src/com/android/remoteprovisioner/unittest/ServerToSystemTest.java b/tests/unittests/src/com/android/remoteprovisioner/unittest/ServerToSystemTest.java
index 94206af..f4fdf06 100644
--- a/tests/unittests/src/com/android/remoteprovisioner/unittest/ServerToSystemTest.java
+++ b/tests/unittests/src/com/android/remoteprovisioner/unittest/ServerToSystemTest.java
@@ -140,14 +140,16 @@ public class ServerToSystemTest {
     @Test
     public void testFallback() throws Exception {
         // Feed a fake URL into the device config to ensure that remote provisioning fails.
-        SettingsManager.setDeviceConfig(sContext, 2 /* extraKeys */, mDuration /* expiringBy */,
+        SettingsManager.setDeviceConfig(sContext, 1 /* extraKeys */, mDuration /* expiringBy */,
                                         "Not even a URL" /* url */);
         int numTestKeys = 1;
         assertPoolStatus(0, 0, 0, 0, mDuration);
+        // Note that due to the GenerateRkpKeyService, this call to generate an attested key will
+        // still cause the service to generate keys up the number specified as `extraKeys` in the
+        // `setDeviceConfig`. This will provide us 1 key for the followup call to provisionCerts.
         Certificate[] fallbackKeyCerts1 = generateKeyStoreKey("test1");
 
         SettingsManager.clearPreferences(sContext);
-        sBinder.generateKeyPair(IS_TEST_MODE, TRUSTED_ENVIRONMENT);
         GeekResponse geek = ServerInterface.fetchGeek(sContext);
         int numProvisioned =
                 Provisioner.provisionCerts(numTestKeys, TRUSTED_ENVIRONMENT,
diff --git a/tests/unittests/src/com/android/remoteprovisioner/unittest/StatsProcessorTest.java b/tests/unittests/src/com/android/remoteprovisioner/unittest/StatsProcessorTest.java
new file mode 100644
index 0000000..3f30ebb
--- /dev/null
+++ b/tests/unittests/src/com/android/remoteprovisioner/unittest/StatsProcessorTest.java
@@ -0,0 +1,342 @@
+/*
+ * Copyright (C) 2022 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.remoteprovisioner.unittest;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import android.platform.test.annotations.Presubmit;
+import android.security.remoteprovisioning.AttestationPoolStatus;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.remoteprovisioner.StatsProcessor;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+public class StatsProcessorTest {
+
+    private AttestationPoolStatus mPool = new AttestationPoolStatus();
+    private int mNumExtraKeys;
+
+    @Before
+    public void setUp() {
+        mNumExtraKeys = 0;
+        mPool.expiring = 0;
+        mPool.unassigned = 0;
+        mPool.attested = 0;
+        mPool.total = 0;
+    }
+
+    @Presubmit
+    @Test
+    public void testMinUnassignedBoundary() throws Exception {
+        mNumExtraKeys = 5;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = StatsProcessor.calcMinUnassignedToTriggerProvisioning(mNumExtraKeys);
+        assertEquals(2, mPool.unassigned);
+        // Add an unassigned key to avoid the threshold for triggering reprovisioning.
+        mPool.unassigned += 1;
+        mPool.expiring = 0;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertFalse(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(7, stats.keysInUse);
+        assertEquals(12, stats.idealTotalSignedKeys);
+        assertEquals(0, stats.keysToGenerate);
+        // Now test provisioning needed boundary
+        mPool.unassigned = StatsProcessor.calcMinUnassignedToTriggerProvisioning(mNumExtraKeys);
+        assertEquals(2, mPool.unassigned);
+        mPool.expiring = 0;
+        stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(8, stats.keysInUse);
+        assertEquals(13, stats.idealTotalSignedKeys);
+        assertEquals(13, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testStatsNoProvisioning() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = StatsProcessor.calcMinUnassignedToTriggerProvisioning(mNumExtraKeys);
+        assertEquals(2, mPool.unassigned);
+        // Add an unassigned key to avoid the threshold for triggering reprovisioning.
+        mPool.unassigned += 1;
+        mPool.expiring = 0;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertFalse(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(7, stats.keysInUse);
+        assertEquals(11, stats.idealTotalSignedKeys);
+        assertEquals(0, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testStatsProvisioning1() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = StatsProcessor.calcMinUnassignedToTriggerProvisioning(mNumExtraKeys);
+        assertEquals(2, mPool.unassigned);
+        mPool.expiring = 0;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(8, stats.keysInUse);
+        assertEquals(12, stats.idealTotalSignedKeys);
+        assertEquals(12, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testStatsProvisioning2() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = StatsProcessor.calcMinUnassignedToTriggerProvisioning(mNumExtraKeys);
+        assertEquals(2, mPool.unassigned);
+        mPool.unassigned -= 1;
+        mPool.expiring = 0;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(9, stats.keysInUse);
+        assertEquals(13, stats.idealTotalSignedKeys);
+        assertEquals(13, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testStatsProvisioningSomePreGenerated() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 15;
+        mPool.attested = 10;
+        mPool.unassigned = StatsProcessor.calcMinUnassignedToTriggerProvisioning(mNumExtraKeys);
+        assertEquals(2, mPool.unassigned);
+        mPool.expiring = 0;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(5, stats.unattestedKeys);
+        assertEquals(8, stats.keysInUse);
+        assertEquals(12, stats.idealTotalSignedKeys);
+        assertEquals(7, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testStatsProvisioningAllPreGenerated() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 22;
+        mPool.attested = 10;
+        mPool.unassigned = StatsProcessor.calcMinUnassignedToTriggerProvisioning(mNumExtraKeys);
+        assertEquals(2, mPool.unassigned);
+        mPool.expiring = 0;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(12, stats.unattestedKeys);
+        assertEquals(8, stats.keysInUse);
+        assertEquals(12, stats.idealTotalSignedKeys);
+        assertEquals(0, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testStatsProvisioningTonsPreGenerated() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 33;
+        mPool.attested = 10;
+        mPool.unassigned = StatsProcessor.calcMinUnassignedToTriggerProvisioning(mNumExtraKeys);
+        assertEquals(2, mPool.unassigned);
+        mPool.expiring = 0;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(23, stats.unattestedKeys);
+        assertEquals(8, stats.keysInUse);
+        assertEquals(12, stats.idealTotalSignedKeys);
+        assertEquals(0, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testExpiringProvisioningNeeded1() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = 5;
+        mPool.expiring = 6;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(5, stats.keysInUse);
+        assertEquals(9, stats.idealTotalSignedKeys);
+        assertEquals(9, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testExpiringProvisioningNeeded2() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = 5;
+        mPool.expiring = 10;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(5, stats.keysInUse);
+        assertEquals(9, stats.idealTotalSignedKeys);
+        assertEquals(9, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testExpiringProvisioningNeeded3() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = 5;
+        mPool.expiring = 5;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(5, stats.keysInUse);
+        assertEquals(9, stats.idealTotalSignedKeys);
+        assertEquals(9, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testExpiringProvisioningNeeded4() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = 10;
+        mPool.expiring = 10;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(0, stats.keysInUse);
+        assertEquals(4, stats.idealTotalSignedKeys);
+        assertEquals(4, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testExpiringProvisioningNeeded5() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = 5;
+        mPool.expiring = 3;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(5, stats.keysInUse);
+        assertEquals(9, stats.idealTotalSignedKeys);
+        assertEquals(9, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testExpiringProvisioningNotNeeded1() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = 5;
+        mPool.expiring = 0;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertFalse(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(5, stats.keysInUse);
+        assertEquals(9, stats.idealTotalSignedKeys);
+        assertEquals(0, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testExpiringProvisioningNotNeeded2() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = 5;
+        mPool.expiring = 2;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertFalse(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(5, stats.keysInUse);
+        assertEquals(9, stats.idealTotalSignedKeys);
+        assertEquals(0, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testExpiringProvisioningNeededSomeKeysPregenerated() throws Exception {
+        mNumExtraKeys = 4;
+        mPool.total = 12;
+        mPool.attested = 10;
+        mPool.unassigned = 5;
+        mPool.expiring = 6;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(2, stats.unattestedKeys);
+        assertEquals(5, stats.keysInUse);
+        assertEquals(9, stats.idealTotalSignedKeys);
+        assertEquals(7, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testBothExpiringAndBelowMinimumExtraKeysAvailable() throws Exception {
+        mNumExtraKeys = 5;
+        mPool.total = 10;
+        mPool.attested = 10;
+        mPool.unassigned = 1;
+        mPool.expiring = 6;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(0, stats.unattestedKeys);
+        assertEquals(9, stats.keysInUse);
+        assertEquals(14, stats.idealTotalSignedKeys);
+        assertEquals(14, stats.keysToGenerate);
+    }
+
+    @Presubmit
+    @Test
+    public void testBothExpiringAndBelowMinimumExtraKeysAvailableWithPreGenKeys() throws Exception {
+        mNumExtraKeys = 5;
+        mPool.total = 14;
+        mPool.attested = 10;
+        mPool.unassigned = 1;
+        mPool.expiring = 6;
+        StatsProcessor.PoolStats stats = StatsProcessor.processPool(mPool, mNumExtraKeys);
+        assertTrue(stats.provisioningNeeded);
+        assertEquals(4, stats.unattestedKeys);
+        assertEquals(9, stats.keysInUse);
+        assertEquals(14, stats.idealTotalSignedKeys);
+        assertEquals(10, stats.keysToGenerate);
+    }
+}
diff --git a/tests/unittests/src/com/android/remoteprovisioner/unittest/SystemInterfaceTest.java b/tests/unittests/src/com/android/remoteprovisioner/unittest/SystemInterfaceTest.java
index 44be452..ec35a0c 100644
--- a/tests/unittests/src/com/android/remoteprovisioner/unittest/SystemInterfaceTest.java
+++ b/tests/unittests/src/com/android/remoteprovisioner/unittest/SystemInterfaceTest.java
@@ -200,7 +200,7 @@ public class SystemInterfaceTest {
             SystemInterface.provisionCertChain(X509Utils.getAndFormatRawPublicKey(certChain[i][0]),
                                                certChain[i][0].getEncoded() /* leafCert */,
                                                os.toByteArray() /* certChain */,
-                                               System.currentTimeMillis() + 2000 /* validity */,
+                                               System.currentTimeMillis() + 25000 /* validity */,
                                                SecurityLevel.TRUSTED_ENVIRONMENT,
                                                mBinder);
         }