Refactoring app for server changes, P256 support

This change adds support for sending some device information up to the
server and receiving a CBOR encoded configuration package in response
during the fetching of an EEK certificate chain. This configuration
information is stored and accessed in the SettingsManager class, which
acts as a wrapper around the SharedPreferences API.

Since this device configuration API change is not yet ready, the app
will support both the current request/response and the
soon-to-be-accepted request/response on the fetchEek API in order to
facilitate a seamless change to the next version.

In addition to the device configuration, the server will also send down
an array of EekChains instead of just one. This is due to the fact that
P256 support has been added for the exchange between the server and the
underlying KM implementation to support StrongBox KMs that don't
currently have support for curve25519.

Finally, this change also adds logic to better support provisioning the
number of keys that are actually needed based on which keys are
expiring, how many unsigned keypairs are already generated and currently
available, and how many keys are actually assigned and in active use.
These numbers can start to diverge from what is expected due to the
ability for keystore to nudge the RemoteProvisioner app to provision
keys on demand in case there are more apps on device using attestation
than what the default number of signed keys is in the pool.

Bug: 189018262
Test: atest RemoteProvisionerUnitTests
Change-Id: Id861dcd66ce6d2773aa1d5b5e429bdba2a085a67
diff --git a/src/com/android/remoteprovisioner/BootReceiver.java b/src/com/android/remoteprovisioner/BootReceiver.java
index f39d155..aabbff3 100644
--- a/src/com/android/remoteprovisioner/BootReceiver.java
+++ b/src/com/android/remoteprovisioner/BootReceiver.java
@@ -34,6 +34,7 @@
     @Override
     public void onReceive(Context context, Intent intent) {
         Log.d(TAG, "Caught boot intent, waking up.");
+        SettingsManager.generateAndSetId(context);
         JobInfo info = new JobInfo
                 .Builder(1, new ComponentName(context, PeriodicProvisioner.class))
                 .setRequiredNetworkType(JobInfo.NETWORK_TYPE_ANY)
diff --git a/src/com/android/remoteprovisioner/CborUtils.java b/src/com/android/remoteprovisioner/CborUtils.java
index d30f8f5..d3fc3d7 100644
--- a/src/com/android/remoteprovisioner/CborUtils.java
+++ b/src/com/android/remoteprovisioner/CborUtils.java
@@ -16,12 +16,14 @@
 
 package com.android.remoteprovisioner;
 
+import android.content.Context;
 import android.os.Build;
 import android.util.Log;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -35,8 +37,16 @@
 import co.nstant.in.cbor.model.MajorType;
 import co.nstant.in.cbor.model.Map;
 import co.nstant.in.cbor.model.UnicodeString;
+import co.nstant.in.cbor.model.UnsignedInteger;
 
 public class CborUtils {
+    public static final int EC_CURVE_P256 = 1;
+    public static final int EC_CURVE_25519 = 2;
+
+    public static final String EXTRA_KEYS = "num_extra_attestation_keys";
+    public static final String TIME_TO_REFRESH = "time_to_refresh_hours";
+    public static final String PROVISIONING_URL = "provisioning_url";
+
     private static final int RESPONSE_CERT_ARRAY_INDEX = 0;
     private static final int RESPONSE_ARRAY_SIZE = 1;
 
@@ -44,11 +54,18 @@
     private static final int UNIQUE_CERTIFICATES_INDEX = 1;
     private static final int CERT_ARRAY_ENTRIES = 2;
 
-    private static final int EEK_INDEX = 0;
+    private static final int EEK_AND_CURVE_INDEX = 0;
     private static final int CHALLENGE_INDEX = 1;
-    private static final int EEK_ARRAY_ENTRIES = 2;
+    private static final int CONFIG_INDEX = 2;
 
+    private static final int CURVE_AND_EEK_CHAIN_LENGTH = 2;
+    private static final int CURVE_INDEX = 0;
+    private static final int EEK_CERT_CHAIN_INDEX = 1;
+
+    private static final int EEK_ARRAY_ENTRIES_NO_CONFIG = 2;
+    private static final int EEK_ARRAY_ENTRIES_WITH_CONFIG = 3;
     private static final String TAG = "RemoteProvisioningService";
+    private static final byte[] EMPTY_MAP = new byte[] {(byte) 0xA0};
 
     /**
      * Parses the signed certificate chains returned by the server. In order to reduce data use over
@@ -67,10 +84,10 @@
             ByteArrayInputStream bais = new ByteArrayInputStream(serverResp);
             List<DataItem> dataItems = new CborDecoder(bais).decode();
             if (dataItems.size() != RESPONSE_ARRAY_SIZE
-                    || dataItems.get(RESPONSE_CERT_ARRAY_INDEX).getMajorType() != MajorType.ARRAY) {
+                    || !checkType(dataItems.get(RESPONSE_CERT_ARRAY_INDEX),
+                                  MajorType.ARRAY, "CborResponse")) {
                 Log.e(TAG, "Improper formatting of CBOR response. Expected size 1. Actual: "
-                            + dataItems.size() + "\nExpected major type: Array. Actual: "
-                            + dataItems.get(0).getMajorType().name());
+                            + dataItems.size());
                 return null;
             }
             dataItems = ((Array) dataItems.get(RESPONSE_CERT_ARRAY_INDEX)).getDataItems();
@@ -79,12 +96,10 @@
                             + dataItems.size());
                 return null;
             }
-            if (dataItems.get(SHARED_CERTIFICATES_INDEX).getMajorType() != MajorType.BYTE_STRING
-                    || dataItems.get(UNIQUE_CERTIFICATES_INDEX).getMajorType() != MajorType.ARRAY) {
-                Log.e(TAG, "Incorrect CBOR types. Expected 'Byte String' and 'Array'. Got: "
-                            + dataItems.get(SHARED_CERTIFICATES_INDEX).getMajorType().name()
-                            + " and "
-                            + dataItems.get(UNIQUE_CERTIFICATES_INDEX).getMajorType().name());
+            if (!checkType(dataItems.get(SHARED_CERTIFICATES_INDEX),
+                           MajorType.BYTE_STRING, "SharedCertificates")
+                    || !checkType(dataItems.get(UNIQUE_CERTIFICATES_INDEX),
+                                  MajorType.ARRAY, "UniqueCertificates")) {
                 return null;
             }
             byte[] sharedCertificates =
@@ -92,9 +107,7 @@
             Array uniqueCertificates = (Array) dataItems.get(UNIQUE_CERTIFICATES_INDEX);
             List<byte[]> uniqueCertificateChains = new ArrayList<byte[]>();
             for (DataItem entry : uniqueCertificates.getDataItems()) {
-                if (entry.getMajorType() != MajorType.BYTE_STRING) {
-                    Log.e(TAG, "Incorrect CBOR type. Expected: 'Byte String'. Actual:"
-                                + entry.getMajorType().name());
+                if (!checkType(entry, MajorType.BYTE_STRING, "UniqueCertificate")) {
                     return null;
                 }
                 ByteArrayOutputStream concat = new ByteArrayOutputStream();
@@ -112,39 +125,109 @@
         return null;
     }
 
+    private static boolean checkType(DataItem item, MajorType majorType, String field) {
+        if (item.getMajorType() != majorType) {
+            Log.e(TAG, "Incorrect CBOR type for field: " + field + ". Expected " + majorType.name()
+                        + ". Actual: " + item.getMajorType().name());
+            return false;
+        }
+        return true;
+    }
+
+    private static boolean parseDeviceConfig(GeekResponse resp, DataItem deviceConfig) {
+        if (!checkType(deviceConfig, MajorType.MAP, "DeviceConfig")) {
+            return false;
+        }
+        Map deviceConfiguration = (Map) deviceConfig;
+        DataItem extraKeys =
+                deviceConfiguration.get(new UnicodeString(EXTRA_KEYS));
+        DataItem timeToRefreshHours =
+                deviceConfiguration.get(new UnicodeString(TIME_TO_REFRESH));
+        DataItem newUrl =
+                deviceConfiguration.get(new UnicodeString(PROVISIONING_URL));
+        if (extraKeys != null) {
+            if (!checkType(extraKeys, MajorType.UNSIGNED_INTEGER, "ExtraKeys")) {
+                return false;
+            }
+            resp.numExtraAttestationKeys = ((UnsignedInteger) extraKeys).getValue().intValue();
+        }
+        if (timeToRefreshHours != null) {
+            if (!checkType(timeToRefreshHours, MajorType.UNSIGNED_INTEGER, "TimeToRefresh")) {
+                return false;
+            }
+            resp.timeToRefresh =
+                    Duration.ofHours(((UnsignedInteger) timeToRefreshHours).getValue().intValue());
+        }
+        if (newUrl != null) {
+            if (!checkType(newUrl, MajorType.UNICODE_STRING, "ProvisioningURL")) {
+                return false;
+            }
+            resp.provisioningUrl = ((UnicodeString) newUrl).getString();
+        }
+        return true;
+    }
+
     /**
      * Parses the Google Endpoint Encryption Key response provided by the server which contains a
      * Google signed EEK and a challenge for use by the underlying IRemotelyProvisionedComponent HAL
      */
     public static GeekResponse parseGeekResponse(byte[] serverResp) {
         try {
+            GeekResponse resp = new GeekResponse();
             ByteArrayInputStream bais = new ByteArrayInputStream(serverResp);
             List<DataItem> dataItems = new CborDecoder(bais).decode();
             if (dataItems.size() != RESPONSE_ARRAY_SIZE
-                    || dataItems.get(RESPONSE_CERT_ARRAY_INDEX).getMajorType() != MajorType.ARRAY) {
+                    || !checkType(dataItems.get(RESPONSE_CERT_ARRAY_INDEX),
+                                  MajorType.ARRAY, "CborResponse")) {
                 Log.e(TAG, "Improper formatting of CBOR response. Expected size 1. Actual: "
-                            + dataItems.size() + "\nExpected major type: Array. Actual: "
-                            + dataItems.get(0).getMajorType().name());
-                return null;
-            }
-            dataItems = ((Array) dataItems.get(RESPONSE_CERT_ARRAY_INDEX)).getDataItems();
-            if (dataItems.size() != EEK_ARRAY_ENTRIES) {
-                Log.e(TAG, "Incorrect number of certificate array entries. Expected: 2. Actual: "
                             + dataItems.size());
                 return null;
             }
-            if (dataItems.get(EEK_INDEX).getMajorType() != MajorType.ARRAY
-                    || dataItems.get(CHALLENGE_INDEX).getMajorType() != MajorType.BYTE_STRING) {
-                Log.e(TAG, "Incorrect CBOR types. Expected 'Array' and 'Byte String'. Got: "
-                            + dataItems.get(EEK_INDEX).getMajorType().name()
-                            + " and "
-                            + dataItems.get(CHALLENGE_INDEX).getMajorType().name());
+            List<DataItem> respItems =
+                    ((Array) dataItems.get(RESPONSE_CERT_ARRAY_INDEX)).getDataItems();
+            if (respItems.size() != EEK_ARRAY_ENTRIES_NO_CONFIG
+                    && respItems.size() != EEK_ARRAY_ENTRIES_WITH_CONFIG) {
+                Log.e(TAG, "Incorrect number of certificate array entries. Expected: "
+                            + EEK_ARRAY_ENTRIES_NO_CONFIG + " or " + EEK_ARRAY_ENTRIES_WITH_CONFIG
+                            + ". Actual: " + respItems.size());
                 return null;
             }
-            ByteArrayOutputStream baos = new ByteArrayOutputStream();
-            new CborEncoder(baos).encode(dataItems.get(EEK_INDEX));
-            return new GeekResponse(baos.toByteArray(),
-                                    ((ByteString) dataItems.get(CHALLENGE_INDEX)).getBytes());
+            if (!checkType(respItems.get(EEK_AND_CURVE_INDEX), MajorType.ARRAY, "EekAndCurveArr")) {
+                return null;
+            }
+            List<DataItem> curveAndEekChains =
+                    ((Array) respItems.get(EEK_AND_CURVE_INDEX)).getDataItems();
+            for (int i = 0; i < curveAndEekChains.size(); i++) {
+                if (!checkType(curveAndEekChains.get(i), MajorType.ARRAY, "EekAndCurve")) {
+                    return null;
+                }
+                List<DataItem> curveAndEekChain =
+                        ((Array) curveAndEekChains.get(i)).getDataItems();
+                if (curveAndEekChain.size() != CURVE_AND_EEK_CHAIN_LENGTH) {
+                    Log.e(TAG, "Wrong size. Expected: " + CURVE_AND_EEK_CHAIN_LENGTH + ". Actual: "
+                               + curveAndEekChain.size());
+                    return null;
+                }
+                if (!checkType(curveAndEekChain.get(CURVE_INDEX),
+                               MajorType.UNSIGNED_INTEGER, "Curve")
+                        || !checkType(curveAndEekChain.get(EEK_CERT_CHAIN_INDEX),
+                                                           MajorType.ARRAY, "EekCertChain")) {
+                    return null;
+                }
+                ByteArrayOutputStream baos = new ByteArrayOutputStream();
+                new CborEncoder(baos).encode(curveAndEekChain.get(EEK_CERT_CHAIN_INDEX));
+                UnsignedInteger curve = (UnsignedInteger) curveAndEekChain.get(CURVE_INDEX);
+                resp.addGeek(curve.getValue().intValue(), baos.toByteArray());
+            }
+            if (!checkType(respItems.get(CHALLENGE_INDEX), MajorType.BYTE_STRING, "Challenge")) {
+                return null;
+            }
+            resp.setChallenge(((ByteString) respItems.get(CHALLENGE_INDEX)).getBytes());
+            if (respItems.size() == EEK_ARRAY_ENTRIES_WITH_CONFIG
+                    && !parseDeviceConfig(resp, respItems.get(CONFIG_INDEX))) {
+                return null;
+            }
+            return resp;
         } catch (CborException e) {
             Log.e(TAG, "CBOR parsing/serializing failed.", e);
             return null;
@@ -152,22 +235,46 @@
     }
 
     /**
+     * Creates the bundle of data that the server needs in order to make a decision over what
+     * device configuration values to return. In general, this boils down to if remote provisioning
+     * is turned on at all or not.
+     *
+     * @return the CBOR encoded provisioning information relevant to the server.
+     */
+    public static byte[] buildProvisioningInfo(Context context) {
+        try {
+            ByteArrayOutputStream baos = new ByteArrayOutputStream();
+            new CborEncoder(baos).encode(new CborBuilder()
+                    .addMap()
+                        .put("fingerprint", Build.FINGERPRINT)
+                        .put(new UnicodeString("id"),
+                             new UnsignedInteger(SettingsManager.getId(context)))
+                        .end()
+                    .build());
+            return baos.toByteArray();
+        } catch (CborException e) {
+            Log.e(TAG, "CBOR serialization failed.", e);
+            return EMPTY_MAP;
+        }
+    }
+
+    /**
      * Takes the various fields fetched from the server and the remote provisioning service and
      * formats them in the CBOR blob the server is expecting as defined by the
      * IRemotelyProvisionedComponent HAL AIDL files.
      */
     public static byte[] buildCertificateRequest(byte[] deviceInfo, byte[] challenge,
                                                  byte[] protectedData, byte[] macedKeysToSign) {
-            // This CBOR library doesn't support adding already serialized CBOR structures into a
-            // CBOR builder. Because of this, we have to first deserialize the provided parameters
-            // back into the library's CBOR object types, and then reserialize them into the
-            // desired structure.
+        // This CBOR library doesn't support adding already serialized CBOR structures into a
+        // CBOR builder. Because of this, we have to first deserialize the provided parameters
+        // back into the library's CBOR object types, and then reserialize them into the
+        // desired structure.
         try {
             // Deserialize the protectedData blob
             ByteArrayInputStream bais = new ByteArrayInputStream(protectedData);
             List<DataItem> dataItems = new CborDecoder(bais).decode();
-            if (dataItems.size() != 1 || dataItems.get(0).getMajorType() != MajorType.ARRAY) {
-                Log.e(TAG, "protectedData is carrying unexpected data.");
+            if (dataItems.size() != 1
+                    || !checkType(dataItems.get(0), MajorType.ARRAY, "ProtectedData")) {
                 return null;
             }
             Array protectedDataArray = (Array) dataItems.get(0);
@@ -175,8 +282,8 @@
             // Deserialize macedKeysToSign
             bais = new ByteArrayInputStream(macedKeysToSign);
             dataItems = new CborDecoder(bais).decode();
-            if (dataItems.size() != 1 || dataItems.get(0).getMajorType() != MajorType.ARRAY) {
-                Log.e(TAG, "macedKeysToSign is carrying unexpected data.");
+            if (dataItems.size() != 1
+                    || !checkType(dataItems.get(0), MajorType.ARRAY, "MacedKeysToSign")) {
                 return null;
             }
             Array macedKeysToSignArray = (Array) dataItems.get(0);
@@ -184,8 +291,8 @@
             // Deserialize deviceInfo
             bais = new ByteArrayInputStream(deviceInfo);
             dataItems = new CborDecoder(bais).decode();
-            if (dataItems.size() != 1 || dataItems.get(0).getMajorType() != MajorType.MAP) {
-                Log.e(TAG, "macedKeysToSign is carrying unexpected data.");
+            if (dataItems.size() != 1
+                    || !checkType(dataItems.get(0), MajorType.MAP, "DeviceInfo")) {
                 return null;
             }
             Map verifiedDeviceInfoMap = (Map) dataItems.get(0);
diff --git a/src/com/android/remoteprovisioner/GeekResponse.java b/src/com/android/remoteprovisioner/GeekResponse.java
index 0a1c798..dcfe040 100644
--- a/src/com/android/remoteprovisioner/GeekResponse.java
+++ b/src/com/android/remoteprovisioner/GeekResponse.java
@@ -16,24 +16,83 @@
 
 package com.android.remoteprovisioner;
 
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
+
 /**
  * Convenience class for packaging up the values returned by the server when initially requesting
  * an Endpoint Encryption Key for remote provisioning. Those values are described by the following
  * CDDL Schema:
  *    GeekResponse = [
- *        EekChain,
+ *        [+CurveAndEek],
  *        challenge : bstr,
+ *        ? Config,
  *    ]
+ *    CurveAndEek = [
+ *        curve: uint,
+ *        EekChain
+ *    ]
+ *    Config = {
+ *        ? "num_extra_attestation_keys": uint,
+ *        ? "time_to_refresh_hours" : uint,
+ *        ? "provisioning_url": tstr,
+ *    }
  *
  * The CDDL that defines EekChain is defined in the RemoteProvisioning HAL, but this app does not
  * require any semantic understanding of the format to perform its function.
  */
 public class GeekResponse {
-    public byte[] challenge;
-    public byte[] geek;
+    public static final int NO_EXTRA_KEY_UPDATE = -1;
+    private byte[] mChallenge;
+    private Map<Integer, byte[]> mCurveToGeek;
+    public int numExtraAttestationKeys;
+    public Duration timeToRefresh;
+    public String provisioningUrl;
 
-    public GeekResponse(byte[] geek, byte[] challenge) {
-        this.geek = geek;
-        this.challenge = challenge;
+    /**
+     * Default initializer.
+     */
+    public GeekResponse() {
+        mCurveToGeek = new HashMap();
+        numExtraAttestationKeys = NO_EXTRA_KEY_UPDATE;
+    }
+
+    /**
+     * Add a CBOR encoded array containing a GEEK and the corresponding certificate chain, keyed
+     * on the EC {@code curve}.
+     *
+     * @param curve an integer which represents an EC curve.
+     * @param geekChain the encoded CBOR array containing an ECDH key and corresponding certificate
+     *                  chain.
+     */
+    public void addGeek(int curve, byte[] geekChain) {
+        mCurveToGeek.put(curve, geekChain);
+    }
+
+    /**
+     * Returns the encoded CBOR array with an ECDH key corresponding to the provided {@code curve}.
+     *
+     * @param curve an integer which represents an EC curve.
+     * @return the corresponding encoded CBOR array.
+     */
+    public byte[] getGeekChain(int curve) {
+        return mCurveToGeek.get(curve);
+    }
+
+    /**
+     * Sets the {@code challenge}.
+     */
+    public void setChallenge(byte[] challenge) {
+        mChallenge = challenge;
+    }
+
+    /**
+     * Returns the {@code challenge}.
+     *
+     * @return the challenge that will be embedded in the CSR sent to the server.
+     */
+    public byte[] getChallenge() {
+        return mChallenge;
     }
 }
diff --git a/src/com/android/remoteprovisioner/PeriodicProvisioner.java b/src/com/android/remoteprovisioner/PeriodicProvisioner.java
index 397ea9a..52ab4bf 100644
--- a/src/com/android/remoteprovisioner/PeriodicProvisioner.java
+++ b/src/com/android/remoteprovisioner/PeriodicProvisioner.java
@@ -18,22 +18,27 @@
 
 import android.app.job.JobParameters;
 import android.app.job.JobService;
+import android.content.Context;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.security.remoteprovisioning.AttestationPoolStatus;
+import android.security.remoteprovisioning.ImplInfo;
 import android.security.remoteprovisioning.IRemoteProvisioning;
 import android.util.Log;
 
+import java.time.Duration;
+
 /**
  * A class that extends JobService in order to be scheduled to check the status of the attestation
  * key pool at regular intervals. If the job determines that more keys need to be generated and
  * signed, it drives that process.
  */
 public class PeriodicProvisioner extends JobService {
-    //TODO(b/176249146): Replace default values here with values fetched from the server.
-    private static int sTotalSignedKeys = 0;
-    // Check for expiring certs in the next 3 days
-    private static int sExpiringBy = 1000 * 60 * 60 * 24 * 3;
+
+    private static final int FAILURE_MAXIMUM = 5;
+
+    // How long to wait in between key pair generations to avoid flooding keystore with requests.
+    private static final Duration KEY_GENERATION_PAUSE = Duration.ofMillis(1000);
 
     private static final String SERVICE = "android.security.remoteprovisioning";
     private static final String TAG = "RemoteProvisioningService";
@@ -45,7 +50,7 @@
      */
     public boolean onStartJob(JobParameters params) {
         Log.d(TAG, "Starting provisioning job");
-        mProvisionerThread = new ProvisionerThread(params);
+        mProvisionerThread = new ProvisionerThread(params, this);
         mProvisionerThread.start();
         return true;
     }
@@ -59,14 +64,25 @@
     }
 
     private class ProvisionerThread extends Thread {
+        private Context mContext;
         private JobParameters mParams;
 
-        ProvisionerThread(JobParameters params) {
+        ProvisionerThread(JobParameters params, Context context) {
             mParams = params;
+            mContext = context;
         }
 
         public void run() {
             try {
+                if (SettingsManager.getExtraSignedKeysAvailable(mContext) == 0) {
+                    // Provisioning is disabled. Check with the server if it's time to turn it back
+                    // on. If not, quit.
+                    GeekResponse check = ServerInterface.fetchGeek(mContext);
+                    if (check.numExtraAttestationKeys == 0) {
+                        jobFinished(mParams, false /* wantsReschedule */);
+                        return;
+                    }
+                }
                 IRemoteProvisioning binder =
                         IRemoteProvisioning.Stub.asInterface(ServiceManager.getService(SERVICE));
                 if (binder == null) {
@@ -74,18 +90,54 @@
                     jobFinished(mParams, true /* wantsReschedule */);
                     return;
                 }
-                int[] securityLevels = binder.getSecurityLevels();
-                if (securityLevels == null) {
+                ImplInfo[] implInfos = binder.getImplementationInfo();
+                if (implInfos == null) {
                     Log.e(TAG, "No instances of IRemotelyProvisionedComponent registered in "
                                + SERVICE);
                     jobFinished(mParams, true /* wantsReschedule */);
                     return;
                 }
-                for (int i = 0; i < securityLevels.length; i++) {
-                    // TODO(b/176249146): Replace expiration date parameter with value fetched from
-                    //                    server
-                    checkAndProvision(binder, System.currentTimeMillis() + sExpiringBy,
-                            securityLevels[i]);
+                int[] keysNeededForSecLevel = new int[implInfos.length];
+                boolean provisioningNeeded = false;
+                for (int i = 0; i < implInfos.length; i++) {
+                    keysNeededForSecLevel[i] =
+                            generateNumKeysNeeded(binder,
+                                       SettingsManager.getExpiringBy(mContext)
+                                                      .plusMillis(System.currentTimeMillis())
+                                                      .toMillis(),
+                                       implInfos[i].secLevel);
+                    if (keysNeededForSecLevel[i] > 0) {
+                        provisioningNeeded = true;
+                    }
+                }
+                if (provisioningNeeded) {
+                    GeekResponse resp = ServerInterface.fetchGeek(mContext);
+                    if (resp == null) {
+                        if (SettingsManager.getFailureCounter(mContext) > FAILURE_MAXIMUM) {
+                            SettingsManager.clearPreferences(mContext);
+                        }
+                        jobFinished(mParams, true /* wantsReschedule */);
+                        return;
+                    }
+                    // Updates to configuration will take effect on the next check.
+                    SettingsManager.setDeviceConfig(mContext,
+                                                    resp.numExtraAttestationKeys,
+                                                    resp.timeToRefresh,
+                                                    resp.provisioningUrl);
+                    if (resp.numExtraAttestationKeys == 0) {
+                        // If the server has sent this, deactivate RKP.
+                        binder.deleteAllKeys();
+                        jobFinished(mParams, false /* wantsReschedule */);
+                        return;
+                    }
+                    for (int i = 0; i < implInfos.length; i++) {
+                        Provisioner.provisionCerts(keysNeededForSecLevel[i],
+                                                   implInfos[i].secLevel,
+                                                   resp.getGeekChain(implInfos[i].supportedCurve),
+                                                   resp.getChallenge(),
+                                                   binder,
+                                                   mContext);
+                    }
                 }
                 jobFinished(mParams, false /* wantsReschedule */);
             } catch (RemoteException e) {
@@ -97,19 +149,38 @@
             }
         }
 
-        private void checkAndProvision(IRemoteProvisioning binder, long expiringBy, int secLevel)
+        /**
+         * 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 = binder.getPoolStatus(expiringBy, secLevel);
-            int generated = 0;
-            while (generated + pool.total - pool.expiring < sTotalSignedKeys) {
-                generated++;
+            int unattestedKeys = pool.total - pool.attested;
+            int validKeys = pool.attested - pool.expiring;
+            int keysInUse = pool.attested - pool.unassigned;
+            int totalSignedKeys = keysInUse + SettingsManager.getExtraSignedKeysAvailable(mContext);
+            int generated;
+            for (generated = 0;
+                    generated + unattestedKeys + validKeys < totalSignedKeys; generated++) {
                 binder.generateKeyPair(false /* isTestMode */, secLevel);
-                Thread.sleep(5000);
+                // 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 (generated > 0) {
-                Log.d(TAG, "Keys generated, moving to provisioning process.");
-                Provisioner.provisionCerts(generated, secLevel, binder);
+            if (totalSignedKeys - validKeys > 0) {
+                return generated + unattestedKeys;
             }
+            return 0;
         }
     }
 }
diff --git a/src/com/android/remoteprovisioner/Provisioner.java b/src/com/android/remoteprovisioner/Provisioner.java
index eabe7d3..91341fc 100644
--- a/src/com/android/remoteprovisioner/Provisioner.java
+++ b/src/com/android/remoteprovisioner/Provisioner.java
@@ -17,6 +17,7 @@
 package com.android.remoteprovisioner;
 
 import android.annotation.NonNull;
+import android.content.Context;
 import android.hardware.security.keymint.DeviceInfo;
 import android.hardware.security.keymint.ProtectedData;
 import android.security.remoteprovisioning.IRemoteProvisioning;
@@ -25,7 +26,7 @@
 import java.security.cert.CertificateEncodingException;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
-import java.util.ArrayList;
+import java.util.List;
 
 /**
  * Provides an easy package to run the provisioning process from start to finish, interfacing
@@ -36,57 +37,57 @@
     private static final String TAG = "RemoteProvisioningService";
 
     /**
-     * Drives the process of provisioning certs. The method first contacts the provided backend
-     * server to retrieve an Endpoing Encryption Key with an accompanying certificate chain and a
-     * challenge. It passes this data and the requested number of keys to the remote provisioning
-     * system backend, which then works with KeyMint in order to get a CSR bundle generated, along
-     * with an encrypted package containing metadata that the server needs in order to make
-     * decisions about provisioning.
+     * Drives the process of provisioning certs. The method passes the data fetched from the
+     * provisioning server along with the requested number of keys to the remote provisioning
+     * system backend. The backend will talk to the underlying IRemotelyProvisionedComponent
+     * interface in order to get a CSR bundle generated, along with an encrypted package containing
+     * metadata that the server needs in order to make decisions about provisioning.
      *
      * This method then passes that bundle back out to the server backend, waits for the response,
      * and, if successful, passes the certificate chains back to the remote provisioning service to
      * be stored and later assigned to apps requesting a key attestation.
      *
      * @param numKeys The number of keys to be signed. The service will do a best-effort to
-     *                     provision the number requested, but if the number requested is larger
-     *                     than the number of unsigned attestation key pairs available, it will
-     *                     only sign the number that is available at time of calling.
-     *
+     *                provision the number requested, but if the number requested is larger
+     *                than the number of unsigned attestation key pairs available, it will
+     *                only sign the number that is available at time of calling.
      * @param secLevel Which KM instance should be used to provision certs.
+     * @param geekChain The certificate chain that signs the endpoint encryption key.
+     * @param challenge A server provided challenge to ensure freshness of the response.
      * @param binder The IRemoteProvisioning binder interface needed by the method to handle talking
-     *                     to the remote provisioning system component.
-     *
-     * @return True if certificates were successfully provisioned for the signing keys.
+     *               to the remote provisioning system component.
+     * @param context The application context object which enables this method to make use of
+     *                SettingsManager.
+     * @return The number of certificates provisoned. Ideally, this should equal {@code numKeys}.
      */
-    public static boolean provisionCerts(int numKeys, int secLevel,
-            @NonNull IRemoteProvisioning binder) {
+    public static int provisionCerts(int numKeys, int secLevel, byte[] geekChain, byte[] challenge,
+            @NonNull IRemoteProvisioning binder, Context context) {
         if (numKeys < 1) {
             Log.e(TAG, "Request at least 1 key to be signed. Num requested: " + numKeys);
-            return false;
-        }
-        GeekResponse geek = ServerInterface.fetchGeek();
-        if (geek == null) {
-            Log.e(TAG, "The geek is null");
-            return false;
+            return 0;
         }
         DeviceInfo deviceInfo = new DeviceInfo();
         ProtectedData protectedData = new ProtectedData();
         byte[] macedKeysToSign =
-                SystemInterface.generateCsr(false /* testMode */, numKeys, secLevel, geek,
-                                            protectedData, deviceInfo, binder);
+                SystemInterface.generateCsr(false /* testMode */, numKeys, secLevel, geekChain,
+                                            challenge, protectedData, deviceInfo, binder);
         if (macedKeysToSign == null || protectedData.protectedData == null
                 || deviceInfo.deviceInfo == null) {
             Log.e(TAG, "Keystore failed to generate a payload");
-            return false;
+            return 0;
         }
         byte[] certificateRequest =
                 CborUtils.buildCertificateRequest(deviceInfo.deviceInfo,
-                                                  geek.challenge,
+                                                  challenge,
                                                   protectedData.protectedData,
                                                   macedKeysToSign);
-        ArrayList<byte[]> certChains =
-                new ArrayList<byte[]>(ServerInterface.requestSignedCertificates(
-                        certificateRequest, geek.challenge));
+        List<byte[]> certChains = ServerInterface.requestSignedCertificates(context,
+                        certificateRequest, challenge);
+        if (certChains == null) {
+            Log.e(TAG, "Server response failed on provisioning attempt.");
+            return 0;
+        }
+        int provisioned = 0;
         for (byte[] certChain : certChains) {
             // DER encoding specifies leaf to root ordering. Pull the public key and expiration
             // date from the leaf.
@@ -95,21 +96,21 @@
                 cert = X509Utils.formatX509Certs(certChain)[0];
             } catch (CertificateException e) {
                 Log.e(TAG, "Failed to interpret DER encoded certificate chain", e);
-                return false;
+                return 0;
             }
             // getTime returns the time in *milliseconds* since the epoch.
             long expirationDate = cert.getNotAfter().getTime();
             byte[] rawPublicKey = X509Utils.getAndFormatRawPublicKey(cert);
             try {
-                return SystemInterface.provisionCertChain(rawPublicKey, cert.getEncoded(),
-                                                          certChain, expirationDate, secLevel,
-                                                          binder);
+                if (SystemInterface.provisionCertChain(rawPublicKey, cert.getEncoded(), certChain,
+                                                       expirationDate, secLevel, binder)) {
+                    provisioned++;
+                }
             } catch (CertificateEncodingException e) {
                 Log.e(TAG, "Somehow can't re-encode the decoded batch cert...", e);
-                return false;
+                return provisioned;
             }
         }
-        Log.d(TAG, "Reaching this return statement implies the server returned 0 signed certs.");
-        return false;
+        return provisioned;
     }
 }
diff --git a/src/com/android/remoteprovisioner/ServerInterface.java b/src/com/android/remoteprovisioner/ServerInterface.java
index 576e913..0ce70ef 100644
--- a/src/com/android/remoteprovisioner/ServerInterface.java
+++ b/src/com/android/remoteprovisioner/ServerInterface.java
@@ -16,6 +16,7 @@
 
 package com.android.remoteprovisioner;
 
+import android.content.Context;
 import android.util.Base64;
 import android.util.Log;
 
@@ -36,30 +37,28 @@
     private static final int TIMEOUT_MS = 5000;
 
     private static final String TAG = "ServerInterface";
-    private static final String PROVISIONING_URL = "https://remoteprovisioning.googleapis.com";
-    private static final String GEEK_URL = PROVISIONING_URL + "/v1alpha1:fetchEekChain";
-    private static final String CERTIFICATE_SIGNING_URL =
-            PROVISIONING_URL + "/v1alpha1:signCertificates?challenge=";
+    private static final String GEEK_URL = ":fetchEekChain";
+    private static final String CERTIFICATE_SIGNING_URL = ":signCertificates?challenge=";
 
     /**
      * Ferries the CBOR blobs returned by KeyMint to the provisioning server. The data sent to the
      * provisioning server contains the MAC'ed CSRs and encrypted bundle containing the MAC key and
      * the hardware unique public key.
      *
+     * @param context The application context which is required to use SettingsManager.
      * @param csr The CBOR encoded data containing the relevant pieces needed for the server to
      *                    sign the CSRs. The data encoded within comes from Keystore / KeyMint.
-     *
      * @param challenge The challenge that was sent from the server. It is included here even though
      *                    it is also included in `cborBlob` in order to allow the server to more
      *                    easily reject bad requests.
-     *
      * @return A List of byte arrays, where each array contains an entire DER-encoded certificate
      *                    chain for one attestation key pair.
      */
-    public static List<byte[]> requestSignedCertificates(byte[] csr, byte[] challenge) {
+    public static List<byte[]> requestSignedCertificates(Context context, byte[] csr,
+                                                         byte[] challenge) {
         try {
-            URL url = new URL(CERTIFICATE_SIGNING_URL
-                        + Base64.encodeToString(challenge, Base64.URL_SAFE));
+            URL url = new URL(SettingsManager.getUrl(context) + CERTIFICATE_SIGNING_URL
+                              + Base64.encodeToString(challenge, Base64.URL_SAFE));
             HttpURLConnection con = (HttpURLConnection) url.openConnection();
             con.setRequestMethod("POST");
             con.setDoOutput(true);
@@ -69,15 +68,15 @@
             // the output stream being automatically closed.
             try (OutputStream os = con.getOutputStream()) {
                 os.write(csr, 0, csr.length);
-            } catch (Exception e) {
-                return null;
             }
 
             if (con.getResponseCode() != HttpURLConnection.HTTP_OK) {
+                int failures = SettingsManager.incrementFailureCounter(context);
                 Log.e(TAG, "Server connection for signing failed, response code: "
-                        + con.getResponseCode());
+                        + con.getResponseCode() + "\nRepeated failure count: " + failures);
                 return null;
             }
+            SettingsManager.clearFailureCounter(context);
             BufferedInputStream inputStream = new BufferedInputStream(con.getInputStream());
             ByteArrayOutputStream cborBytes = new ByteArrayOutputStream();
             byte[] buffer = new byte[1024];
@@ -87,9 +86,11 @@
             }
             return CborUtils.parseSignedCertificates(cborBytes.toByteArray());
         } catch (SocketTimeoutException e) {
+            SettingsManager.incrementFailureCounter(context);
             Log.e(TAG, "Server timed out", e);
             return null;
         } catch (IOException e) {
+            SettingsManager.incrementFailureCounter(context);
             Log.e(TAG, "Failed to request signed certificates from the server", e);
             return null;
         }
@@ -103,19 +104,30 @@
      *
      * A challenge is also returned from the server so that it can check freshness of the follow-up
      * request to get keys signed.
+     *
+     * @param context The application context which is required to use SettingsManager.
+     * @return A GeekResponse object which optionally contains configuration data.
      */
-    public static GeekResponse fetchGeek() {
+    public static GeekResponse fetchGeek(Context context) {
         try {
-            URL url = new URL(GEEK_URL);
+            URL url = new URL(SettingsManager.getUrl(context) + GEEK_URL);
             HttpURLConnection con = (HttpURLConnection) url.openConnection();
-            con.setRequestMethod("GET");
+            con.setRequestMethod("POST");
             con.setConnectTimeout(TIMEOUT_MS);
+            con.setDoOutput(true);
+
+            byte[] config = CborUtils.buildProvisioningInfo(context);
+            try (OutputStream os = con.getOutputStream()) {
+                os.write(config, 0, config.length);
+            }
 
             if (con.getResponseCode() != HttpURLConnection.HTTP_OK) {
+                int failures = SettingsManager.incrementFailureCounter(context);
                 Log.e(TAG, "Server connection for GEEK failed, response code: "
-                        + con.getResponseCode());
+                        + con.getResponseCode() + "\nRepeated failure count: " + failures);
                 return null;
             }
+            SettingsManager.clearFailureCounter(context);
 
             BufferedInputStream inputStream = new BufferedInputStream(con.getInputStream());
             ByteArrayOutputStream cborBytes = new ByteArrayOutputStream();
@@ -127,8 +139,11 @@
             inputStream.close();
             return CborUtils.parseGeekResponse(cborBytes.toByteArray());
         } catch (SocketTimeoutException e) {
+            SettingsManager.incrementFailureCounter(context);
             Log.e(TAG, "Server timed out", e);
         } catch (IOException e) {
+            // This exception will trigger on a completely malformed URL.
+            SettingsManager.incrementFailureCounter(context);
             Log.e(TAG, "Failed to fetch GEEK from the servers.", e);
         }
         return null;
diff --git a/src/com/android/remoteprovisioner/SystemInterface.java b/src/com/android/remoteprovisioner/SystemInterface.java
index 160f185..79f4cea 100644
--- a/src/com/android/remoteprovisioner/SystemInterface.java
+++ b/src/com/android/remoteprovisioner/SystemInterface.java
@@ -68,15 +68,15 @@
      * Sends a generateCsr request over the binder interface. `dataBlob` is an out parameter that
      * will be populated by the underlying binder service.
      */
-    public static byte[] generateCsr(boolean testMode, int numKeys, int secLevel, GeekResponse geek,
-            ProtectedData protectedData, DeviceInfo deviceInfo,
+    public static byte[] generateCsr(boolean testMode, int numKeys, int secLevel,
+            byte[] geekChain, byte[] challenge, ProtectedData protectedData, DeviceInfo deviceInfo,
             @NonNull IRemoteProvisioning binder) {
         try {
             ProtectedData dataBundle = new ProtectedData();
             byte[] macedPublicKeys = binder.generateCsr(testMode,
                                                         numKeys,
-                                                        geek.geek,
-                                                        geek.challenge,
+                                                        geekChain,
+                                                        challenge,
                                                         secLevel,
                                                         protectedData,
                                                         deviceInfo);
diff --git a/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java b/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java
index f505d28..b3292b5 100644
--- a/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java
+++ b/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java
@@ -17,21 +17,27 @@
 package com.android.remoteprovisioner.service;
 
 import android.app.Service;
+import android.content.Context;
 import android.content.Intent;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.security.IGenerateRkpKeyService;
 import android.security.remoteprovisioning.AttestationPoolStatus;
+import android.security.remoteprovisioning.ImplInfo;
 import android.security.remoteprovisioning.IRemoteProvisioning;
 import android.util.Log;
 
+import com.android.remoteprovisioner.GeekResponse;
 import com.android.remoteprovisioner.Provisioner;
+import com.android.remoteprovisioner.ServerInterface;
+import com.android.remoteprovisioner.SettingsManager;
 
 /**
  * Provides the implementation for IGenerateKeyService.aidl
  */
 public class GenerateRkpKeyService extends Service {
+    private static final int KEY_GENERATION_PAUSE_MS = 1000;
     private static final String SERVICE = "android.security.remoteprovisioning";
     private static final String TAG = "RemoteProvisioningService";
 
@@ -51,7 +57,6 @@
             try {
                 IRemoteProvisioning binder =
                         IRemoteProvisioning.Stub.asInterface(ServiceManager.getService(SERVICE));
-                // Iterate through each security level backend
                 checkAndFillPool(binder, securityLevel);
             } catch (RemoteException e) {
                 Log.e(TAG, "Remote Exception: ", e);
@@ -63,7 +68,6 @@
             try {
                 IRemoteProvisioning binder =
                         IRemoteProvisioning.Stub.asInterface(ServiceManager.getService(SERVICE));
-                // Iterate through each security level backend
                 checkAndFillPool(binder, securityLevel);
             } catch (RemoteException e) {
                 Log.e(TAG, "Remote Exception: ", e);
@@ -74,13 +78,38 @@
                 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;
+                }
+            }
             // If there are no unassigned keys, go ahead and provision some. If there are no keys
             // at all on 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.total != 0) {
-                Log.d(TAG, "All signed keys are currently in use, provisioning more.");
-                binder.generateKeyPair(false /* isTestMode */, secLevel);
-                Provisioner.provisionCerts(1 /* numCsr */, secLevel, binder);
+                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);
+                    }
+                } catch (InterruptedException e) {
+                    Log.i(TAG, "Thread interrupted", e);
+                }
+                GeekResponse resp = ServerInterface.fetchGeek(context);
+                if (resp == null) {
+                    Log.e(TAG, "Server unavailable");
+                    return;
+                }
+                Provisioner.provisionCerts(keysToProvision, secLevel, resp.getGeekChain(curve),
+                                           resp.getChallenge(), binder, context);
             }
         }
     };
diff --git a/tests/unittests/Android.bp b/tests/unittests/Android.bp
index c2ed374..2aafdc7 100644
--- a/tests/unittests/Android.bp
+++ b/tests/unittests/Android.bp
@@ -19,6 +19,7 @@
     name: "RemoteProvisionerUnitTests",
     srcs: ["src/**/*.java"],
     static_libs: [
+        "androidx.test.core",
         "androidx.test.rules",
         "android.security.remoteprovisioning-java",
         "platform-test-annotations",
diff --git a/tests/unittests/src/com/android/remoteprovisioner/unittest/CborUtilsTest.java b/tests/unittests/src/com/android/remoteprovisioner/unittest/CborUtilsTest.java
index 4c30a6e..cc561a8 100644
--- a/tests/unittests/src/com/android/remoteprovisioner/unittest/CborUtilsTest.java
+++ b/tests/unittests/src/com/android/remoteprovisioner/unittest/CborUtilsTest.java
@@ -31,8 +31,12 @@
 import co.nstant.in.cbor.CborDecoder;
 import co.nstant.in.cbor.CborEncoder;
 import co.nstant.in.cbor.model.Array;
+import co.nstant.in.cbor.model.ByteString;
 import co.nstant.in.cbor.model.DataItem;
 import co.nstant.in.cbor.model.MajorType;
+import co.nstant.in.cbor.model.Map;
+import co.nstant.in.cbor.model.UnicodeString;
+import co.nstant.in.cbor.model.UnsignedInteger;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -47,10 +51,48 @@
 public class CborUtilsTest {
 
     private ByteArrayOutputStream mBaos;
+    private Array mGeekChain1;
+    private byte[] mEncodedmGeekChain1;
+    private Array mGeekChain2;
+    private byte[] mEncodedGeekChain2;
+    private Map mDeviceConfig;
+    private static final byte[] CHALLENGE = new byte[]{0x0a, 0x0b, 0x0c};
+    private static final int TEST_EXTRA_KEYS = 18;
+    private static final int TEST_TIME_TO_REFRESH_HOURS = 42;
+    private static final String TEST_URL = "https://www.wonderifthisisvalid.combutjustincase";
+
+    private byte[] encodeDataItem(DataItem toEncode) throws Exception {
+        new CborEncoder(mBaos).encode(new CborBuilder()
+                        .add(toEncode)
+                        .build());
+        byte[] encoded = mBaos.toByteArray();
+        mBaos.reset();
+        return encoded;
+    }
 
     @Before
     public void setUp() throws Exception {
         mBaos = new ByteArrayOutputStream();
+
+        mGeekChain1 = new Array();
+        mGeekChain1.add(new ByteString(new byte[] {0x01, 0x02, 0x03}))
+                   .add(new ByteString(new byte[] {0x04, 0x05, 0x06}))
+                   .add(new ByteString(new byte[] {0x07, 0x08, 0x09}));
+        mEncodedmGeekChain1 = encodeDataItem(mGeekChain1);
+
+        mGeekChain2 = new Array();
+        mGeekChain2.add(new ByteString(new byte[] {0x09, 0x08, 0x07}))
+                   .add(new ByteString(new byte[] {0x06, 0x05, 0x04}))
+                   .add(new ByteString(new byte[] {0x03, 0x02, 0x01}));
+        mEncodedGeekChain2 = encodeDataItem(mGeekChain2);
+
+        mDeviceConfig = new Map();
+        mDeviceConfig.put(new UnicodeString(CborUtils.EXTRA_KEYS),
+                          new UnsignedInteger(TEST_EXTRA_KEYS))
+                     .put(new UnicodeString(CborUtils.TIME_TO_REFRESH),
+                          new UnsignedInteger(TEST_TIME_TO_REFRESH_HOURS))
+                     .put(new UnicodeString(CborUtils.PROVISIONING_URL),
+                          new UnicodeString(TEST_URL));
     }
 
     @Presubmit
@@ -116,39 +158,182 @@
     public void testParseGeekResponseFakeData() throws Exception {
         new CborEncoder(mBaos).encode(new CborBuilder()
                 .addArray()
-                    .addArray()                                       // GEEK Chain
-                        .add(new byte[] {0x01, 0x02, 0x03})
-                        .add(new byte[] {0x04, 0x05, 0x06})
-                        .add(new byte[] {0x07, 0x08, 0x09})
+                    .addArray()                                       // GEEK Curve to Chains
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_25519))
+                            .add(mGeekChain1)
+                            .end()
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_P256))
+                            .add(mGeekChain2)
+                            .end()
                         .end()
-                    .add(new byte[] {0x0a, 0x0b, 0x0c})               // Challenge
+                    .add(CHALLENGE)
+                    .add(mDeviceConfig)
                     .end()
                 .build());
         GeekResponse resp = CborUtils.parseGeekResponse(mBaos.toByteArray());
         mBaos.reset();
+        assertArrayEquals(mEncodedmGeekChain1, resp.getGeekChain(CborUtils.EC_CURVE_25519));
+        assertArrayEquals(mEncodedGeekChain2, resp.getGeekChain(CborUtils.EC_CURVE_P256));
+        assertArrayEquals(CHALLENGE, resp.getChallenge());
+        assertEquals(TEST_EXTRA_KEYS, resp.numExtraAttestationKeys);
+        assertEquals(TEST_TIME_TO_REFRESH_HOURS, resp.timeToRefresh.toHours());
+        assertEquals(TEST_URL, resp.provisioningUrl);
+    }
+
+    @Test
+    public void testExtraDeviceConfigEntriesDontFail() throws Exception {
         new CborEncoder(mBaos).encode(new CborBuilder()
                 .addArray()
-                    .add(new byte[] {0x01, 0x02, 0x03})
-                    .add(new byte[] {0x04, 0x05, 0x06})
-                    .add(new byte[] {0x07, 0x08, 0x09})
+                    .addArray()                                       // GEEK Curve to Chains
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_25519))
+                            .add(mGeekChain1)
+                            .end()
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_P256))
+                            .add(mGeekChain2)
+                            .end()
+                        .end()
+                    .add(CHALLENGE)
+                    .add(mDeviceConfig.put(new UnicodeString("new_field"),
+                                          new UnsignedInteger(84)))
                     .end()
                 .build());
-        byte[] expectedGeek = mBaos.toByteArray();
-        assertArrayEquals(expectedGeek, resp.geek);
-        assertArrayEquals(new byte[] {0x0a, 0x0b, 0x0c}, resp.challenge);
+        GeekResponse resp = CborUtils.parseGeekResponse(mBaos.toByteArray());
+        mBaos.reset();
+        assertArrayEquals(mEncodedmGeekChain1, resp.getGeekChain(CborUtils.EC_CURVE_25519));
+        assertArrayEquals(mEncodedGeekChain2, resp.getGeekChain(CborUtils.EC_CURVE_P256));
+        assertArrayEquals(CHALLENGE, resp.getChallenge());
+        assertEquals(TEST_EXTRA_KEYS, resp.numExtraAttestationKeys);
+        assertEquals(TEST_TIME_TO_REFRESH_HOURS, resp.timeToRefresh.toHours());
+        assertEquals(TEST_URL, resp.provisioningUrl);
+    }
+
+    @Test
+    public void testMissingDeviceConfigDoesntFail() throws Exception {
+        new CborEncoder(mBaos).encode(new CborBuilder()
+                .addArray()
+                    .addArray()                                       // GEEK Curve to Chains
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_25519))
+                            .add(mGeekChain1)
+                            .end()
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_P256))
+                            .add(mGeekChain2)
+                            .end()
+                        .end()
+                    .add(CHALLENGE)
+                    .end()
+                .build());
+        GeekResponse resp = CborUtils.parseGeekResponse(mBaos.toByteArray());
+        mBaos.reset();
+        assertArrayEquals(mEncodedmGeekChain1, resp.getGeekChain(CborUtils.EC_CURVE_25519));
+        assertArrayEquals(mEncodedGeekChain2, resp.getGeekChain(CborUtils.EC_CURVE_P256));
+        assertArrayEquals(CHALLENGE, resp.getChallenge());
+        assertEquals(GeekResponse.NO_EXTRA_KEY_UPDATE, resp.numExtraAttestationKeys);
+        assertEquals(null, resp.timeToRefresh);
+        assertEquals(null, resp.provisioningUrl);
+    }
+
+    @Test
+    public void testMissingDeviceConfigEntriesDoesntFail() throws Exception {
+        mDeviceConfig.remove(new UnicodeString(CborUtils.TIME_TO_REFRESH));
+        new CborEncoder(mBaos).encode(new CborBuilder()
+                .addArray()
+                    .addArray()                                       // GEEK Curve to Chains
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_25519))
+                            .add(mGeekChain1)
+                            .end()
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_P256))
+                            .add(mGeekChain2)
+                            .end()
+                        .end()
+                    .add(CHALLENGE)
+                    .add(mDeviceConfig)
+                    .end()
+                .build());
+        GeekResponse resp = CborUtils.parseGeekResponse(mBaos.toByteArray());
+        mBaos.reset();
+        assertArrayEquals(mEncodedmGeekChain1, resp.getGeekChain(CborUtils.EC_CURVE_25519));
+        assertArrayEquals(mEncodedGeekChain2, resp.getGeekChain(CborUtils.EC_CURVE_P256));
+        assertArrayEquals(CHALLENGE, resp.getChallenge());
+        assertEquals(TEST_EXTRA_KEYS, resp.numExtraAttestationKeys);
+        assertEquals(null, resp.timeToRefresh);
+        assertEquals(TEST_URL, resp.provisioningUrl);
+    }
+
+    @Test
+    public void testParseGeekResponseFailsOnWrongType() throws Exception {
+        new CborEncoder(mBaos).encode(new CborBuilder()
+                .addArray()
+                    .addArray()
+                        .addArray()
+                            .add("String instead of curve enum")
+                            .add(mGeekChain1)
+                            .end()
+                        .end()
+                    .add(CHALLENGE)
+                    .add(mDeviceConfig)
+                    .end()
+                .build());
+        assertNull(CborUtils.parseGeekResponse(mBaos.toByteArray()));
+        mBaos.reset();
+        new CborEncoder(mBaos).encode(new CborBuilder()
+                .addArray()
+                    .addArray()
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_25519))
+                            .add(new ByteString(CHALLENGE)) // Must be an array of bstrs
+                            .end()
+                        .end()
+                    .add(CHALLENGE)
+                    .add(mDeviceConfig)
+                    .end()
+                .build());
+        assertNull(CborUtils.parseGeekResponse(mBaos.toByteArray()));
+        mBaos.reset();
+        new CborEncoder(mBaos).encode(new CborBuilder()
+                .addArray()
+                    .addArray()
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_25519))
+                            .add(mGeekChain1)
+                            .end()
+                        .end()
+                    .add(new UnicodeString("tstr instead of bstr"))
+                    .add(mDeviceConfig)
+                    .end()
+                .build());
+        assertNull(CborUtils.parseGeekResponse(mBaos.toByteArray()));
+        mBaos.reset();
+        new CborEncoder(mBaos).encode(new CborBuilder()
+                .addArray()
+                    .addArray()
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_25519))
+                            .add(mGeekChain1)
+                            .end()
+                        .end()
+                    .add(CHALLENGE)
+                    .add(CHALLENGE)
+                    .end()
+                .build());
+        assertNull(CborUtils.parseGeekResponse(mBaos.toByteArray()));
     }
 
     @Test
     public void testParseGeekResponseWrongSize() throws Exception {
         new CborEncoder(mBaos).encode(new CborBuilder()
                 .addArray()
-                    .addArray()
-                        .add(new byte[] {0x01, 0x02, 0x03})
-                        .add(new byte[] {0x04, 0x05, 0x06})
-                        .add(new byte[] {0x07, 0x08, 0x09})
-                        .end()
-                    .add(new byte[] {0x0a, 0x0b, 0x0c})
-                    .add("One more entry than there should be")
+                    .add("one entry")
+                    .add("two entries")
+                    .add("three entries")
+                    .add("whoops")
                     .end()
                 .build());
         assertNull(CborUtils.parseGeekResponse(mBaos.toByteArray()));
diff --git a/tests/unittests/src/com/android/remoteprovisioner/unittest/SystemInterfaceTest.java b/tests/unittests/src/com/android/remoteprovisioner/unittest/SystemInterfaceTest.java
index 2a2adc4..44be452 100644
--- a/tests/unittests/src/com/android/remoteprovisioner/unittest/SystemInterfaceTest.java
+++ b/tests/unittests/src/com/android/remoteprovisioner/unittest/SystemInterfaceTest.java
@@ -38,7 +38,7 @@
 
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.remoteprovisioner.GeekResponse;
+import com.android.remoteprovisioner.CborUtils;
 import com.android.remoteprovisioner.SystemInterface;
 import com.android.remoteprovisioner.X509Utils;
 
@@ -119,11 +119,11 @@
         ProtectedData encryptedBundle = new ProtectedData();
         byte[] eek = new byte[32];
         new Random().nextBytes(eek);
-        GeekResponse geek = new GeekResponse(generateEekChain(eek), new byte[] {0x02});
         byte[] bundle =
             SystemInterface.generateCsr(true /* testMode */, 0 /* numKeys */,
                                         SecurityLevel.TRUSTED_ENVIRONMENT,
-                                        geek, encryptedBundle, deviceInfo, mBinder);
+                                        generateEekChain(eek),
+                                        new byte[] {0x02}, encryptedBundle, deviceInfo, mBinder);
         // encryptedBundle should contain a COSE_Encrypt message
         ByteArrayInputStream bais = new ByteArrayInputStream(encryptedBundle.protectedData);
         List<DataItem> dataItems = new CborDecoder(bais).decode();
@@ -156,14 +156,14 @@
         int numKeys = 10;
         byte[] eek = new byte[32];
         new Random().nextBytes(eek);
-        GeekResponse geek = new GeekResponse(generateEekChain(eek), new byte[] {0x02});
         for (int i = 0; i < numKeys; i++) {
             mBinder.generateKeyPair(true /* testMode */, SecurityLevel.TRUSTED_ENVIRONMENT);
         }
         byte[] bundle =
             SystemInterface.generateCsr(true /* testMode */, numKeys,
                                         SecurityLevel.TRUSTED_ENVIRONMENT,
-                                        geek, encryptedBundle, deviceInfo, mBinder);
+                                        generateEekChain(eek),
+                                        new byte[] {0x02}, encryptedBundle, deviceInfo, mBinder);
         assertNotNull(bundle);
         // The return value of generateCsr should be a COSE_Mac0 message
         ByteArrayInputStream bais = new ByteArrayInputStream(bundle);
@@ -280,12 +280,12 @@
         int numKeys = 1;
         byte[] eekPriv = X25519.generatePrivateKey();
         byte[] eekPub = X25519.publicFromPrivate(eekPriv);
-        GeekResponse geek = new GeekResponse(generateEekChain(eekPub), new byte[] {0x02});
         mBinder.generateKeyPair(true /* testMode */, SecurityLevel.TRUSTED_ENVIRONMENT);
         byte[] bundle =
             SystemInterface.generateCsr(true /* testMode */, numKeys,
                                         SecurityLevel.TRUSTED_ENVIRONMENT,
-                                        geek, encryptedBundle, deviceInfo, mBinder);
+                                        generateEekChain(eekPub),
+                                        new byte[] {0x02}, encryptedBundle, deviceInfo, mBinder);
         ByteArrayInputStream bais = new ByteArrayInputStream(encryptedBundle.protectedData);
         List<DataItem> dataItems = new CborDecoder(bais).decode();
         // Parse encMsg into components: protected and unprotected headers, payload, and recipient