Propagate more details about RKP failures am: 7497fdaee1 am: 658a1a076b am: 12b9a62d84

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

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

7 files changed, 350 insertions, 81 deletions

diff --git a/src/com/android/remoteprovisioner/PeriodicProvisioner.java b/src/com/android/remoteprovisioner/PeriodicProvisioner.java
index b050e1b..d9fcc79 100644
--- a/src/com/android/remoteprovisioner/PeriodicProvisioner.java
+++ b/src/com/android/remoteprovisioner/PeriodicProvisioner.java
@@ -28,7 +28,6 @@ import android.security.remoteprovisioning.ImplInfo;
 import android.util.Log;
 
 import androidx.annotation.NonNull;
-import androidx.work.ListenableWorker.Result;
 import androidx.work.Worker;
 import androidx.work.WorkerParameters;
 
@@ -102,9 +101,6 @@ public class PeriodicProvisioner extends Worker {
                     // So long as the connection is unmetered, go ahead and grab an updated
                     // device configuration file.
                     resp = ServerInterface.fetchGeek(mContext);
-                    if (!checkGeekResp(resp)) {
-                        return Result.failure();
-                    }
                     SettingsManager.setDeviceConfig(mContext,
                             resp.numExtraAttestationKeys,
                             resp.timeToRefresh,
@@ -116,9 +112,6 @@ public class PeriodicProvisioner extends Worker {
                 return Result.success();
             }
             resp = ServerInterface.fetchGeek(mContext);
-            if (!checkGeekResp(resp)) {
-                return Result.failure();
-            }
             SettingsManager.setDeviceConfig(mContext,
                         resp.numExtraAttestationKeys,
                         resp.timeToRefresh,
@@ -145,13 +138,20 @@ public class PeriodicProvisioner extends Worker {
         } catch (InterruptedException e) {
             Log.e(TAG, "Provisioner thread interrupted.", e);
             return Result.failure();
+        } catch (RemoteProvisioningException e) {
+            Log.e(TAG, "Encountered RemoteProvisioningException", e);
+            if (SettingsManager.getFailureCounter(mContext) > FAILURE_MAXIMUM) {
+                Log.e(TAG, "Too many failures, resetting defaults.");
+                SettingsManager.clearPreferences(mContext);
+            }
+            return Result.failure();
         }
     }
 
     public static void batchProvision(IRemoteProvisioning binder, Context context,
                                int keysToProvision, int secLevel,
                                byte[] geekChain, byte[] challenge)
-                               throws RemoteException {
+            throws RemoteException, RemoteProvisioningException {
         while (keysToProvision != 0) {
             int batchSize = min(keysToProvision, SAFE_CSR_BATCH_SIZE);
             Log.i(TAG, "Requesting " + batchSize + " keys to be provisioned.");
@@ -165,18 +165,6 @@ public class PeriodicProvisioner extends Worker {
         }
     }
 
-    private boolean checkGeekResp(GeekResponse resp) {
-        if (resp == null) {
-            Log.e(TAG, "Failed to get a response from the server.");
-            if (SettingsManager.getFailureCounter(mContext) > FAILURE_MAXIMUM) {
-                Log.e(TAG, "Too many failures, resetting defaults.");
-                SettingsManager.clearPreferences(mContext);
-            }
-            return false;
-        }
-        return true;
-    }
-
     private boolean isProvisioningNeeded(
             IRemoteProvisioning binder, long expiringBy, ImplInfo[] implInfos,
             int[] keysNeededForSecLevel)
diff --git a/src/com/android/remoteprovisioner/Provisioner.java b/src/com/android/remoteprovisioner/Provisioner.java
index 8093192..293ab00 100644
--- a/src/com/android/remoteprovisioner/Provisioner.java
+++ b/src/com/android/remoteprovisioner/Provisioner.java
@@ -20,6 +20,7 @@ import android.annotation.NonNull;
 import android.content.Context;
 import android.hardware.security.keymint.DeviceInfo;
 import android.hardware.security.keymint.ProtectedData;
+import android.security.IGenerateRkpKeyService;
 import android.security.remoteprovisioning.IRemoteProvisioning;
 import android.util.Log;
 
@@ -58,14 +59,15 @@ public class Provisioner {
      *               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}.
+     * @return The number of certificates provisioned. Ideally, this should equal {@code numKeys}.
      */
     public static int provisionCerts(int numKeys, int secLevel, byte[] geekChain, byte[] challenge,
-            @NonNull IRemoteProvisioning binder, Context context) {
+            @NonNull IRemoteProvisioning binder, Context context) throws
+            RemoteProvisioningException {
         Log.i(TAG, "Request for " + numKeys + " keys to be provisioned.");
         if (numKeys < 1) {
-            Log.e(TAG, "Request at least 1 key to be signed. Num requested: " + numKeys);
-            return 0;
+            throw new RemoteProvisioningException(IGenerateRkpKeyService.Status.INTERNAL_ERROR,
+                    "Request at least 1 key to be signed. Num requested: " + numKeys);
         }
         DeviceInfo deviceInfo = new DeviceInfo();
         ProtectedData protectedData = new ProtectedData();
@@ -73,8 +75,8 @@ public class Provisioner {
                 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 0;
+            throw new RemoteProvisioningException(IGenerateRkpKeyService.Status.INTERNAL_ERROR,
+                    "Keystore failed to generate a payload");
         }
         byte[] certificateRequest =
                 CborUtils.buildCertificateRequest(deviceInfo.deviceInfo,
@@ -82,14 +84,16 @@ public class Provisioner {
                                                   protectedData.protectedData,
                                                   macedKeysToSign);
         if (certificateRequest == null) {
-            Log.e(TAG, "Failed to serialize the payload generated by keystore.");
-            return 0;
+            throw new RemoteProvisioningException(IGenerateRkpKeyService.Status.INTERNAL_ERROR,
+                    "Failed to serialize the payload generated by keystore.");
         }
         List<byte[]> certChains = ServerInterface.requestSignedCertificates(context,
                         certificateRequest, challenge);
         if (certChains == null) {
-            Log.e(TAG, "Server response failed on provisioning attempt.");
-            return 0;
+            // This is marked as an internal error, because ServerInterface should never return
+            // null, and if it does it's indicative of a bug.
+            throw new RemoteProvisioningException(IGenerateRkpKeyService.Status.INTERNAL_ERROR,
+                    "Server response failed on provisioning attempt.");
         }
         Log.i(TAG, "Received " + certChains.size() + " certificate chains from the server.");
         int provisioned = 0;
@@ -100,8 +104,8 @@ public class Provisioner {
             try {
                 cert = X509Utils.formatX509Certs(certChain)[0];
             } catch (CertificateException e) {
-                Log.e(TAG, "Failed to interpret DER encoded certificate chain", e);
-                return 0;
+                throw new RemoteProvisioningException(IGenerateRkpKeyService.Status.INTERNAL_ERROR,
+                        "Failed to interpret DER encoded certificate chain", e);
             }
             // getTime returns the time in *milliseconds* since the epoch.
             long expirationDate = cert.getNotAfter().getTime();
@@ -116,10 +120,19 @@ public class Provisioner {
                     provisioned++;
                 }
             } catch (CertificateEncodingException e) {
-                Log.e(TAG, "Somehow can't re-encode the decoded batch cert...", e);
-                return provisioned;
+                throw new RemoteProvisioningException(IGenerateRkpKeyService.Status.INTERNAL_ERROR,
+                        "Error re-encoding the decoded batch cert", e);
             }
         }
+
+        // In lieu of metrics, do our best to spot when things are "going really badly" by looking
+        // for edge cases where fewer than half of the key requests succeed. In the future, we
+        // should add metrics to get a picture of how healthy this process is.
+        if (provisioned < (numKeys / 2)) {
+            throw new RemoteProvisioningException(IGenerateRkpKeyService.Status.INTERNAL_ERROR,
+                    "Requested " + numKeys + " keys, provisioned " + provisioned);
+        }
+
         Log.i(TAG, "In provisionCerts: Requested " + numKeys + " keys. "
                    + provisioned + " were provisioned.");
         return provisioned;
diff --git a/src/com/android/remoteprovisioner/RemoteProvisioningException.java b/src/com/android/remoteprovisioner/RemoteProvisioningException.java
new file mode 100644
index 0000000..5652556
--- /dev/null
+++ b/src/com/android/remoteprovisioner/RemoteProvisioningException.java
@@ -0,0 +1,94 @@
+/**
+ * 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.annotation.IntDef;
+import android.security.IGenerateRkpKeyService.Status;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+/**
+ * Represents an error that occurred while contacting the remote key provisioning server.
+ */
+public final class RemoteProvisioningException extends Exception {
+    /** @hide */
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef(flag = true, value = {
+            Status.NO_NETWORK_CONNECTIVITY,
+            Status.NETWORK_COMMUNICATION_ERROR,
+            Status.DEVICE_NOT_REGISTERED,
+            Status.HTTP_CLIENT_ERROR,
+            Status.HTTP_SERVER_ERROR,
+            Status.HTTP_UNKNOWN_ERROR,
+            Status.INTERNAL_ERROR,
+    })
+    public @interface ErrorCode {
+    }
+
+    private static final int HTTP_STATUS_DEVICE_NOT_REGISTERED = 444;
+    private static final int HTTP_CLIENT_ERROR_HUNDREDS_DIGIT = 4;
+    private static final int HTTP_SERVER_ERROR_HUNDREDS_DIGIT = 5;
+
+    @ErrorCode
+    private final int mErrorCode;
+
+    /**
+     * @param errorCode the underlying ServerInterface error
+     * @param message describes the exception
+     */
+    public RemoteProvisioningException(@ErrorCode int errorCode, String message) {
+        super(message);
+        mErrorCode = errorCode;
+    }
+
+    /**
+     * @param errorCode the underlying ServerInterface error
+     * @param message describes the exception
+     * @param cause the underlying error that led this exception
+     */
+    public RemoteProvisioningException(@ErrorCode int errorCode, String message, Throwable cause) {
+        super(message, cause);
+        mErrorCode = errorCode;
+    }
+
+    /**
+     * @param httpStatus the HTTP status that lead to the error
+     * @return a newly created RemoteProvisioningException that indicates an HTTP error occurred
+     */
+    public static RemoteProvisioningException createFromHttpError(@ErrorCode int httpStatus) {
+        String message = "HTTP error status encountered: " + httpStatus;
+        if (httpStatus == HTTP_STATUS_DEVICE_NOT_REGISTERED) {
+            return new RemoteProvisioningException(Status.DEVICE_NOT_REGISTERED, message);
+        }
+        if ((httpStatus / 100) == HTTP_CLIENT_ERROR_HUNDREDS_DIGIT) {
+            return new RemoteProvisioningException(Status.HTTP_CLIENT_ERROR, message);
+        }
+        if ((httpStatus / 100) == HTTP_SERVER_ERROR_HUNDREDS_DIGIT) {
+            return new RemoteProvisioningException(Status.HTTP_SERVER_ERROR, message);
+        }
+        return new RemoteProvisioningException(Status.HTTP_UNKNOWN_ERROR, message);
+    }
+
+    /**
+     * @return the underlying error that caused the failure
+     */
+    @ErrorCode
+    public int getErrorCode() {
+        return mErrorCode;
+    }
+}
diff --git a/src/com/android/remoteprovisioner/ServerInterface.java b/src/com/android/remoteprovisioner/ServerInterface.java
index 0ce70ef..697d8dd 100644
--- a/src/com/android/remoteprovisioner/ServerInterface.java
+++ b/src/com/android/remoteprovisioner/ServerInterface.java
@@ -17,6 +17,9 @@
 package com.android.remoteprovisioner;
 
 import android.content.Context;
+import android.net.ConnectivityManager;
+import android.net.NetworkInfo;
+import android.security.IGenerateRkpKeyService;
 import android.util.Base64;
 import android.util.Log;
 
@@ -34,7 +37,7 @@ import java.util.List;
  */
 public class ServerInterface {
 
-    private static final int TIMEOUT_MS = 5000;
+    private static final int TIMEOUT_MS = 20000;
 
     private static final String TAG = "ServerInterface";
     private static final String GEEK_URL = ":fetchEekChain";
@@ -55,7 +58,8 @@ public class ServerInterface {
      *                    chain for one attestation key pair.
      */
     public static List<byte[]> requestSignedCertificates(Context context, byte[] csr,
-                                                         byte[] challenge) {
+                                                         byte[] challenge) throws
+            RemoteProvisioningException {
         try {
             URL url = new URL(SettingsManager.getUrl(context) + CERTIFICATE_SIGNING_URL
                               + Base64.encodeToString(challenge, Base64.URL_SAFE));
@@ -63,6 +67,7 @@ public class ServerInterface {
             con.setRequestMethod("POST");
             con.setDoOutput(true);
             con.setConnectTimeout(TIMEOUT_MS);
+            con.setReadTimeout(TIMEOUT_MS);
 
             // May not be able to use try-with-resources here if the connection gets closed due to
             // the output stream being automatically closed.
@@ -74,7 +79,7 @@ public class ServerInterface {
                 int failures = SettingsManager.incrementFailureCounter(context);
                 Log.e(TAG, "Server connection for signing failed, response code: "
                         + con.getResponseCode() + "\nRepeated failure count: " + failures);
-                return null;
+                throw RemoteProvisioningException.createFromHttpError(con.getResponseCode());
             }
             SettingsManager.clearFailureCounter(context);
             BufferedInputStream inputStream = new BufferedInputStream(con.getInputStream());
@@ -88,11 +93,11 @@ public class ServerInterface {
         } catch (SocketTimeoutException e) {
             SettingsManager.incrementFailureCounter(context);
             Log.e(TAG, "Server timed out", e);
-            return null;
+            throw makeNetworkError(context, "Server timed out");
         } catch (IOException e) {
             SettingsManager.incrementFailureCounter(context);
             Log.e(TAG, "Failed to request signed certificates from the server", e);
-            return null;
+            throw makeNetworkError(context, e.getMessage());
         }
     }
 
@@ -108,12 +113,13 @@ public class ServerInterface {
      * @param context The application context which is required to use SettingsManager.
      * @return A GeekResponse object which optionally contains configuration data.
      */
-    public static GeekResponse fetchGeek(Context context) {
+    public static GeekResponse fetchGeek(Context context) throws RemoteProvisioningException {
         try {
             URL url = new URL(SettingsManager.getUrl(context) + GEEK_URL);
             HttpURLConnection con = (HttpURLConnection) url.openConnection();
             con.setRequestMethod("POST");
             con.setConnectTimeout(TIMEOUT_MS);
+            con.setReadTimeout(TIMEOUT_MS);
             con.setDoOutput(true);
 
             byte[] config = CborUtils.buildProvisioningInfo(context);
@@ -125,7 +131,7 @@ public class ServerInterface {
                 int failures = SettingsManager.incrementFailureCounter(context);
                 Log.e(TAG, "Server connection for GEEK failed, response code: "
                         + con.getResponseCode() + "\nRepeated failure count: " + failures);
-                return null;
+                throw RemoteProvisioningException.createFromHttpError(con.getResponseCode());
             }
             SettingsManager.clearFailureCounter(context);
 
@@ -146,6 +152,17 @@ public class ServerInterface {
             SettingsManager.incrementFailureCounter(context);
             Log.e(TAG, "Failed to fetch GEEK from the servers.", e);
         }
-        return null;
+        throw makeNetworkError(context, "Error fetching GEEK");
+    }
+
+    private static RemoteProvisioningException makeNetworkError(Context context, String message) {
+        ConnectivityManager cm = context.getSystemService(ConnectivityManager.class);
+        NetworkInfo networkInfo = cm.getActiveNetworkInfo();
+        if (networkInfo != null && networkInfo.isConnected()) {
+            return new RemoteProvisioningException(
+                    IGenerateRkpKeyService.Status.NETWORK_COMMUNICATION_ERROR, message);
+        }
+        return new RemoteProvisioningException(
+                IGenerateRkpKeyService.Status.NO_NETWORK_CONNECTIVITY, message);
     }
 }
diff --git a/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java b/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java
index 268a8b5..6d8f48a 100644
--- a/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java
+++ b/src/com/android/remoteprovisioner/service/GenerateRkpKeyService.java
@@ -24,19 +24,20 @@ 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.security.remoteprovisioning.ImplInfo;
 import android.util.Log;
 
 import com.android.remoteprovisioner.GeekResponse;
 import com.android.remoteprovisioner.PeriodicProvisioner;
+import com.android.remoteprovisioner.RemoteProvisioningException;
 import com.android.remoteprovisioner.ServerInterface;
 
 import java.time.Duration;
 import java.util.concurrent.locks.ReentrantLock;
 
 /**
- * Provides the implementation for IGenerateKeyService.aidl
+ * Provides the implementation for IGenerateRkpKeyService.aidl
  */
 public class GenerateRkpKeyService extends Service {
     private static final int KEY_GENERATION_PAUSE_MS = 1000;
@@ -47,6 +48,11 @@ public class GenerateRkpKeyService extends Service {
 
     private static final ReentrantLock sLock = new ReentrantLock();
 
+    private enum Concurrency {
+        BLOCKING,
+        NON_BLOCKING
+    };
+
     @Override
     public void onCreate() {
         super.onCreate();
@@ -59,35 +65,30 @@ public class GenerateRkpKeyService extends Service {
 
     private final IGenerateRkpKeyService.Stub mBinder = new IGenerateRkpKeyService.Stub() {
         @Override
-        public void generateKey(int securityLevel) {
-            try {
-                Log.i(TAG, "generateKey ping for secLevel: " + securityLevel);
-                IRemoteProvisioning binder =
-                        IRemoteProvisioning.Stub.asInterface(ServiceManager.getService(SERVICE));
-                checkAndFillPool(binder, securityLevel);
-            } catch (RemoteException e) {
-                Log.e(TAG, "Remote Exception: ", e);
-            }
+        public int generateKey(int securityLevel) {
+            Log.i(TAG, "generateKey ping for secLevel: " + securityLevel);
+            IRemoteProvisioning binder =
+                    IRemoteProvisioning.Stub.asInterface(ServiceManager.getService(SERVICE));
+            return checkAndFillPool(binder, securityLevel, Concurrency.BLOCKING);
         }
 
         @Override
         public void notifyKeyGenerated(int securityLevel) {
-            try {
-                Log.i(TAG, "Notify key generated ping for secLevel: " + securityLevel);
-                IRemoteProvisioning binder =
-                        IRemoteProvisioning.Stub.asInterface(ServiceManager.getService(SERVICE));
-                checkAndFillPool(binder, securityLevel);
-            } catch (RemoteException e) {
-                Log.e(TAG, "Remote Exception: ", e);
-            }
+            Log.i(TAG, "Notify key generated ping for secLevel: " + securityLevel);
+            IRemoteProvisioning binder =
+                    IRemoteProvisioning.Stub.asInterface(ServiceManager.getService(SERVICE));
+            checkAndFillPool(binder, securityLevel, Concurrency.NON_BLOCKING);
         }
 
-        private void checkAndFillPool(IRemoteProvisioning binder, int secLevel)
-                throws RemoteException {
+        private int checkAndFillPool(IRemoteProvisioning binder, int secLevel,
+                Concurrency concurrency) {
             // No need to hammer the pool check with a ton of redundant requests.
-            if (!sLock.tryLock()) {
+            if (concurrency == Concurrency.BLOCKING) {
+                Log.i(TAG, "Waiting on lock to check pool status.");
+                sLock.lock();
+            } else if (!sLock.tryLock()) {
                 Log.i(TAG, "Exiting check; another process already started the check.");
-                return;
+                return Status.OK;
             }
             try {
                 AttestationPoolStatus pool =
@@ -104,27 +105,26 @@ public class GenerateRkpKeyService extends Service {
                 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) {
+                                LOOKAHEAD_TIME.toMillis(),
+                                secLevel);
+                if (keysToProvision != 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());
+                            resp.getGeekChain(curve), resp.getChallenge());
                 }
             } catch (InterruptedException e) {
                 Log.e(TAG, "Provisioner thread interrupted.", e);
+            } catch (RemoteProvisioningException e) {
+                Log.e(TAG, "RemoteProvisioningException: ", e);
+                return e.getErrorCode();
+            } catch (RemoteException e) {
+                Log.e(TAG, "Remote Exception: ", e);
+                return Status.INTERNAL_ERROR;
             } finally {
                 sLock.unlock();
             }
+            return Status.OK;
         }
     };
 }
diff --git a/tests/unittests/src/com/android/remoteprovisioner/unittest/RemoteProvisioningExceptionTest.java b/tests/unittests/src/com/android/remoteprovisioner/unittest/RemoteProvisioningExceptionTest.java
new file mode 100644
index 0000000..a018627
--- /dev/null
+++ b/tests/unittests/src/com/android/remoteprovisioner/unittest/RemoteProvisioningExceptionTest.java
@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2021 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 android.platform.test.annotations.Presubmit;
+import android.security.IGenerateRkpKeyService.Status;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.remoteprovisioner.RemoteProvisioningException;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+public class RemoteProvisioningExceptionTest {
+    @Presubmit
+    @Test
+    public void handlesArbitraryErrors() {
+        for (int i : new int[]{1, 11, 123, 0x8675309}) {
+            RemoteProvisioningException ex = new RemoteProvisioningException(i, "error: " + i);
+            Assert.assertEquals(i, ex.getErrorCode());
+            Assert.assertEquals("error: " + i, ex.getMessage());
+        }
+    }
+
+    @Presubmit
+    @Test
+    public void handlesUnknownHttpStatus() {
+        RemoteProvisioningException ex = RemoteProvisioningException.createFromHttpError(123);
+        Assert.assertNotNull(ex);
+        Assert.assertEquals(ex.getErrorCode(), Status.HTTP_UNKNOWN_ERROR);
+    }
+
+    @Presubmit
+    @Test
+    public void handlesServerErrors() {
+        for (int httpStatus = 500; httpStatus < 600; ++httpStatus) {
+            RemoteProvisioningException ex = RemoteProvisioningException.createFromHttpError(
+                    httpStatus);
+            Assert.assertNotNull(ex);
+            Assert.assertEquals(httpStatus + " should have been a server error", ex.getErrorCode(),
+                    Status.HTTP_SERVER_ERROR);
+            Assert.assertTrue(ex.getMessage().contains("HTTP"));
+        }
+    }
+
+    @Presubmit
+    @Test
+    public void handlesClientErrors() {
+        for (int httpStatus = 400; httpStatus < 500; ++httpStatus) {
+            RemoteProvisioningException ex = RemoteProvisioningException.createFromHttpError(
+                    httpStatus);
+            Assert.assertNotNull(ex);
+            if (httpStatus == 444) {
+                Assert.assertEquals(ex.getErrorCode(), Status.DEVICE_NOT_REGISTERED);
+            } else {
+                Assert.assertEquals(httpStatus + " should have been a client error",
+                        ex.getErrorCode(), Status.HTTP_CLIENT_ERROR);
+            }
+            Assert.assertTrue(ex.getMessage().contains("HTTP"));
+        }
+    }
+}
diff --git a/tests/unittests/src/com/android/remoteprovisioner/unittest/ServerToSystemTest.java b/tests/unittests/src/com/android/remoteprovisioner/unittest/ServerToSystemTest.java
index f4fdf06..924caff 100644
--- a/tests/unittests/src/com/android/remoteprovisioner/unittest/ServerToSystemTest.java
+++ b/tests/unittests/src/com/android/remoteprovisioner/unittest/ServerToSystemTest.java
@@ -27,13 +27,18 @@ import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
 import android.content.Context;
+import android.os.ParcelFileDescriptor;
 import android.os.ServiceManager;
+import android.os.SystemProperties;
+import android.security.KeyStoreException;
 import android.security.keystore.KeyGenParameterSpec;
 import android.security.remoteprovisioning.AttestationPoolStatus;
-import android.security.remoteprovisioning.ImplInfo;
 import android.security.remoteprovisioning.IRemoteProvisioning;
+import android.security.remoteprovisioning.ImplInfo;
+import android.system.keystore2.ResponseCode;
 
 import androidx.test.core.app.ApplicationProvider;
+import androidx.test.platform.app.InstrumentationRegistry;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.remoteprovisioner.GeekResponse;
@@ -42,13 +47,17 @@ import com.android.remoteprovisioner.ServerInterface;
 import com.android.remoteprovisioner.SettingsManager;
 
 import org.junit.After;
+import org.junit.Assert;
+import org.junit.Assume;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.IOException;
 import java.security.KeyPairGenerator;
 import java.security.KeyStore;
+import java.security.ProviderException;
 import java.security.cert.Certificate;
 import java.time.Duration;
 import java.util.Arrays;
@@ -58,6 +67,7 @@ public class ServerToSystemTest {
 
     private static final boolean IS_TEST_MODE = false;
     private static final String SERVICE = "android.security.remoteprovisioning";
+    private static final String RKP_ONLY_PROP = "remote_provisioning.tee.rkp_only";
 
     private static Context sContext;
     private static IRemoteProvisioning sBinder;
@@ -65,6 +75,26 @@ public class ServerToSystemTest {
 
     private Duration mDuration;
 
+    // Helper class that sets rkp_only to true if it's not already set, then restores the state on
+    // close. Intended to be used in a try expression: try (RkpOnlyContext c = new RkpOnlyContext())
+    private static class ForceRkpOnlyContext implements AutoCloseable {
+        private final boolean mOriginalPropertyValue;
+
+        ForceRkpOnlyContext() {
+            mOriginalPropertyValue = SystemProperties.getBoolean(RKP_ONLY_PROP, false);
+            if (!mOriginalPropertyValue) {
+                SystemProperties.set(RKP_ONLY_PROP, "true");
+            }
+        }
+
+        @Override
+        public void close() {
+            if (!mOriginalPropertyValue) {
+                SystemProperties.set(RKP_ONLY_PROP, "false");
+            }
+        }
+    }
+
     private void assertPoolStatus(int total, int attested,
                                   int unassigned, int expiring, Duration time) throws Exception {
         AttestationPoolStatus pool = sBinder.getPoolStatus(time.toMillis(), TRUSTED_ENVIRONMENT);
@@ -77,6 +107,9 @@ public class ServerToSystemTest {
     private static Certificate[] generateKeyStoreKey(String alias) throws Exception {
         KeyStore keyStore = KeyStore.getInstance("AndroidKeyStore");
         keyStore.load(null);
+        if (keyStore.containsAlias(alias)) {
+            keyStore.deleteEntry(alias);
+        }
         KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM_EC,
                 "AndroidKeyStore");
         KeyGenParameterSpec spec = new KeyGenParameterSpec.Builder(alias, PURPOSE_SIGN)
@@ -139,6 +172,10 @@ public class ServerToSystemTest {
 
     @Test
     public void testFallback() throws Exception {
+        Assume.assumeFalse(
+                "Skipping test as this system does not support fallback from RKP keys",
+                SystemProperties.getBoolean(RKP_ONLY_PROP, false));
+
         // Feed a fake URL into the device config to ensure that remote provisioning fails.
         SettingsManager.setDeviceConfig(sContext, 1 /* extraKeys */, mDuration /* expiringBy */,
                                         "Not even a URL" /* url */);
@@ -165,9 +202,6 @@ public class ServerToSystemTest {
                                         "Not even a URL" /* url */);
         // Even if there is an unsigned key hanging around, fallback should still occur.
         Certificate[] fallbackKeyCerts2 = generateKeyStoreKey("test3");
-        // Due to there being no attested keys in the pool, the provisioning service should not
-        // have even attempted to provision more certificates.
-        assertEquals(0, SettingsManager.getFailureCounter(sContext));
         assertTrue(fallbackKeyCerts1.length == fallbackKeyCerts2.length);
         for (int i = 1; i < fallbackKeyCerts1.length; i++) {
             assertArrayEquals("Cert: " + i, fallbackKeyCerts1[i].getEncoded(),
@@ -179,4 +213,48 @@ public class ServerToSystemTest {
                     Arrays.equals(fallbackKeyCerts1[fallbackKeyCerts1.length - 1].getEncoded(),
                               provisionedKeyCerts[provisionedKeyCerts.length - 1].getEncoded()));
     }
+
+    @Test
+    public void testRetryableRkpError() throws Exception {
+        try (ForceRkpOnlyContext c = new ForceRkpOnlyContext()) {
+            SettingsManager.setDeviceConfig(sContext, 1 /* extraKeys */, mDuration /* expiringBy */,
+                    "Not even a URL" /* url */);
+            generateKeyStoreKey("should-never-succeed");
+            Assert.fail("Expected a keystore exception");
+        } catch (ProviderException e) {
+            Assert.assertTrue(e.getCause() instanceof KeyStoreException);
+            KeyStoreException keyStoreException = (KeyStoreException) e.getCause();
+            Assert.assertEquals(ResponseCode.OUT_OF_KEYS, keyStoreException.getErrorCode());
+            Assert.assertTrue(keyStoreException.isTransientFailure());
+            Assert.assertEquals(KeyStoreException.RETRY_WITH_EXPONENTIAL_BACKOFF,
+                    keyStoreException.getRetryPolicy());
+        }
+    }
+
+    private void setAirplaneMode(boolean isEnabled) throws IOException {
+        String command = "cmd connectivity airplane-mode " + (isEnabled ? "enable" : "disable");
+        ParcelFileDescriptor fd =
+                InstrumentationRegistry.getInstrumentation().getUiAutomation().executeShellCommand(
+                        command);
+        fd.close();
+    }
+
+    @Test
+    public void testRetryWithoutNetworkTee() throws Exception {
+        setAirplaneMode(true);
+        try (ForceRkpOnlyContext c = new ForceRkpOnlyContext()) {
+            assertPoolStatus(0, 0, 0, 0, mDuration);
+            generateKeyStoreKey("should-never-succeed");
+            Assert.fail("Expected a keystore exception");
+        } catch (ProviderException e) {
+            Assert.assertTrue(e.getCause() instanceof KeyStoreException);
+            KeyStoreException keyStoreException = (KeyStoreException) e.getCause();
+            Assert.assertEquals(ResponseCode.OUT_OF_KEYS, keyStoreException.getErrorCode());
+            Assert.assertTrue(keyStoreException.isTransientFailure());
+            Assert.assertEquals(KeyStoreException.RETRY_WHEN_CONNECTIVITY_AVAILABLE,
+                    keyStoreException.getRetryPolicy());
+        } finally {
+            setAirplaneMode(false);
+        }
+    }
 }