Propagate more details about RKP failures

This change reworks the internals of the RemoteProvisioner APK in order
to do a better job of propagating error information from deeper within
the call stack up to the entry points. This is primarily to convey error
information to callers of GenerateRkpKeyService, to allow them to
differentiate whether or not a failure is permanent or transient.

Add unit tests to verify proper error handling and retry detection.

Bug: 227306369
Test: atest RemoteProvisionerUnitTests
Change-Id: If44d9c5b46b2dfce65c31f3cbd372cc754e1dfbe

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);
+        }
+    }
 }