Snap for 10754184 from 2170302cf62e221782e0aff51200e92246e449c7 to mainline-extservices-releaseaml_ext_341131030

Change-Id: Ie4e88411c875dc9bcdb584ce0b18e07ba650b750

15 files changed, 295 insertions, 126 deletions

diff --git a/app/proguard.flags b/app/proguard.flags
index 6339e99..2e7d09b 100644
--- a/app/proguard.flags
+++ b/app/proguard.flags
@@ -2,3 +2,14 @@
 -keep class com.android.rkpdapp.interfaces.ServiceManagerInterface { *; }
 -keep class com.android.rkpdapp.service.** { *; }
 -keep class com.android.rkpdapp.utils.Settings { *; }
+
+# Required for tests that use Mockito's thenThrow with checked exceptions.
+-keepattributes Exceptions
+
+# Minimal set of keep rules for mocked methods with checked exceptions.
+# This can be relaxed to specific packages if that simplifies testing.
+# See also https://r8.googlesource.com/r8/+/refs/heads/master/compatibility-faq.md#r8-full-mode.
+-keepclassmembers,allowshrinking,allowoptimization,allowobfuscation,allowaccessmodification class android.hardware.security.keymint.IRemotelyProvisionedComponent { public <methods>; }
+-keepclassmembers,allowshrinking,allowoptimization,allowobfuscation,allowaccessmodification class com.android.rkpdapp.IGetRegistrationCallback { public <methods>; }
+-keepclassmembers,allowshrinking,allowoptimization,allowobfuscation,allowaccessmodification class com.android.rkpdapp.interfaces.SystemInterface { public <methods>; }
+-keepclassmembers,allowshrinking,allowoptimization,allowobfuscation,allowaccessmodification class com.android.rkpdapp.provisioner.Provisioner { public <methods>; }
\ No newline at end of file
diff --git a/app/src/com/android/rkpdapp/ThreadPool.java b/app/src/com/android/rkpdapp/ThreadPool.java
index 6d4b9a1..3ed7755 100644
--- a/app/src/com/android/rkpdapp/ThreadPool.java
+++ b/app/src/com/android/rkpdapp/ThreadPool.java
@@ -17,21 +17,29 @@
 package com.android.rkpdapp;
 
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 
 /**
  * This class provides a global thread pool to RKPD app.
  */
 public class ThreadPool {
-    public static final int NUMBER_OF_THREADS = Runtime.getRuntime().availableProcessors();
+    public static final int NUMBER_OF_THREADS = 4;
     /*
-     * This thread pool has a minimum of 0 threads and a maximum of up to the
-     * number of processors. If a thread is idle for more than 30 seconds, it is
-     * terminated. RKPD is idle most of the time. So, this way we can don't keep
-     * unused threads around.
-     *
-     * Each thread has an unbounded queue. This allows RKPD to serve requests
-     * asynchronously.
+     * This thread pool has a minimum of 0 threads and a maximum of up to 4. If
+     * a thread is idle for more than 60 seconds, it is terminated. RKPD is idle
+     * most of the time. So, this way we can don't keep unused threads around.
      */
-    public static final ExecutorService EXECUTOR = Executors.newFixedThreadPool(NUMBER_OF_THREADS);
+    public static final ExecutorService EXECUTOR;
+
+    static {
+        ThreadPoolExecutor executor =
+                new ThreadPoolExecutor(/*corePoolSize=*/ NUMBER_OF_THREADS,
+                    /*maximumPoolSize=*/ NUMBER_OF_THREADS,
+                    /*keepAliveTime=*/ 60L, /*unit=*/ TimeUnit.SECONDS,
+                    /*workQueue=*/ new LinkedBlockingQueue<Runnable>());
+        executor.allowCoreThreadTimeOut(true);
+        EXECUTOR = executor;
+    }
 }
diff --git a/app/src/com/android/rkpdapp/interfaces/ServerInterface.java b/app/src/com/android/rkpdapp/interfaces/ServerInterface.java
index a989cf6..173a2a8 100644
--- a/app/src/com/android/rkpdapp/interfaces/ServerInterface.java
+++ b/app/src/com/android/rkpdapp/interfaces/ServerInterface.java
@@ -18,8 +18,9 @@ package com.android.rkpdapp.interfaces;
 
 import android.content.Context;
 import android.net.ConnectivityManager;
-import android.net.NetworkInfo;
+import android.net.NetworkCapabilities;
 import android.net.TrafficStats;
+import android.net.Uri;
 import android.util.Base64;
 import android.util.Log;
 
@@ -37,6 +38,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.net.HttpURLConnection;
+import java.net.MalformedURLException;
 import java.net.SocketTimeoutException;
 import java.net.URL;
 import java.nio.charset.Charset;
@@ -55,17 +57,28 @@ public class ServerInterface {
 
     private static final int TIMEOUT_MS = 20000;
     private static final int BACKOFF_TIME_MS = 100;
+    private static final int SHORT_RETRY_COUNT = 2;
 
     private static final String TAG = "RkpdServerInterface";
     private static final String GEEK_URL = ":fetchEekChain";
-    private static final String CERTIFICATE_SIGNING_URL = ":signCertificates?";
-    private static final String CHALLENGE_PARAMETER = "challenge=";
-    private static final String REQUEST_ID_PARAMETER = "request_id=";
+    private static final String CERTIFICATE_SIGNING_URL = ":signCertificates";
+    private static final String CHALLENGE_PARAMETER = "challenge";
+    private static final String REQUEST_ID_PARAMETER = "request_id";
     private final Context mContext;
 
     private enum Operation {
-        FETCH_GEEK,
-        SIGN_CERTS;
+        FETCH_GEEK(1),
+        SIGN_CERTS(2);
+
+        private final int mTrafficTag;
+
+        Operation(int trafficTag) {
+            mTrafficTag = trafficTag;
+        }
+
+        public int getTrafficTag() {
+            return mTrafficTag;
+        }
 
         public ProvisioningAttempt.Status getHttpErrorStatus() {
             if (Objects.equals(name(), FETCH_GEEK.name())) {
@@ -114,11 +127,9 @@ public class ServerInterface {
      */
     public List<byte[]> requestSignedCertificates(byte[] csr, byte[] challenge,
             ProvisioningAttempt metrics) throws RkpdException, InterruptedException {
-        final String challengeParam = CHALLENGE_PARAMETER + Base64.encodeToString(challenge,
-                Base64.URL_SAFE | Base64.NO_WRAP);
-        final String fullUrl = CERTIFICATE_SIGNING_URL + String.join("&", challengeParam,
-                REQUEST_ID_PARAMETER + generateAndLogRequestId());
-        final byte[] cborBytes = connectAndGetData(metrics, fullUrl, csr, Operation.SIGN_CERTS);
+        final byte[] cborBytes =
+                connectAndGetData(metrics, generateSignCertsUrl(challenge),
+                                  csr, Operation.SIGN_CERTS);
         List<byte[]> certChains = CborUtils.parseSignedCertificates(cborBytes);
         if (certChains == null) {
             metrics.setStatus(ProvisioningAttempt.Status.INTERNAL_ERROR);
@@ -143,6 +154,22 @@ public class ServerInterface {
         return certChains;
     }
 
+    private URL generateSignCertsUrl(byte[] challenge) throws RkpdException {
+        try {
+            return new URL(Uri.parse(Settings.getUrl(mContext)).buildUpon()
+                    .appendEncodedPath(CERTIFICATE_SIGNING_URL)
+                    .appendQueryParameter(CHALLENGE_PARAMETER,
+                            Base64.encodeToString(challenge, Base64.URL_SAFE | Base64.NO_WRAP))
+                    .appendQueryParameter(REQUEST_ID_PARAMETER, generateAndLogRequestId())
+                    .build()
+                    .toString()
+                    // Needed due to the `:` in the URL endpoint.
+                    .replaceFirst("%3A", ":"));
+        } catch (MalformedURLException e) {
+            throw new RkpdException(RkpdException.ErrorCode.HTTP_CLIENT_ERROR, "Bad URL", e);
+        }
+    }
+
     private String generateAndLogRequestId() {
         String reqId = UUID.randomUUID().toString();
         Log.i(TAG, "request_id: " + reqId);
@@ -163,7 +190,8 @@ public class ServerInterface {
     public GeekResponse fetchGeek(ProvisioningAttempt metrics)
             throws RkpdException, InterruptedException {
         byte[] input = CborUtils.buildProvisioningInfo(mContext);
-        byte[] cborBytes = connectAndGetData(metrics, GEEK_URL, input, Operation.FETCH_GEEK);
+        byte[] cborBytes =
+                connectAndGetData(metrics, generateFetchGeekUrl(), input, Operation.FETCH_GEEK);
         GeekResponse resp = CborUtils.parseGeekResponse(cborBytes);
         if (resp == null) {
             metrics.setStatus(ProvisioningAttempt.Status.FETCH_GEEK_HTTP_ERROR);
@@ -174,6 +202,19 @@ public class ServerInterface {
         return resp;
     }
 
+    private URL generateFetchGeekUrl() throws RkpdException {
+        try {
+            return new URL(Uri.parse(Settings.getUrl(mContext)).buildUpon()
+                            .appendPath(GEEK_URL)
+                            .build()
+                            .toString()
+                            // Needed due to the `:` in the URL endpoint.
+                            .replaceFirst("%3A", ":"));
+        } catch (MalformedURLException e) {
+            throw new RkpdException(RkpdException.ErrorCode.INTERNAL_ERROR, "Bad URL", e);
+        }
+    }
+
     private void checkDataBudget(ProvisioningAttempt metrics)
             throws RkpdException {
         if (!Settings.hasErrDataBudget(mContext, null /* curTime */)) {
@@ -186,9 +227,7 @@ public class ServerInterface {
 
     private RkpdException makeNetworkError(String message,
             ProvisioningAttempt metrics) {
-        ConnectivityManager cm = mContext.getSystemService(ConnectivityManager.class);
-        NetworkInfo networkInfo = cm.getActiveNetworkInfo();
-        if (networkInfo != null && networkInfo.isConnected()) {
+        if (isNetworkConnected(mContext)) {
             return new RkpdException(
                     RkpdException.ErrorCode.NETWORK_COMMUNICATION_ERROR, message);
         }
@@ -198,6 +237,18 @@ public class ServerInterface {
     }
 
     /**
+     * Checks whether network is connected.
+     * @return true if connected else false.
+     */
+    public static boolean isNetworkConnected(Context context) {
+        ConnectivityManager cm = context.getSystemService(ConnectivityManager.class);
+        NetworkCapabilities capabilities = cm.getNetworkCapabilities(cm.getActiveNetwork());
+        return capabilities != null
+                && capabilities.hasCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                && capabilities.hasCapability(NetworkCapabilities.NET_CAPABILITY_VALIDATED);
+    }
+
+    /**
      * Fetch a GEEK from the server and update SettingsManager appropriately with the return
      * values. This will also delete all keys in the attestation key pool if the server has
      * indicated that RKP should be turned off.
@@ -273,45 +324,52 @@ public class ServerInterface {
         }
     }
 
-    private byte[] connectAndGetData(ProvisioningAttempt metrics, String endpoint, byte[] input,
+    private byte[] connectAndGetData(ProvisioningAttempt metrics, URL url, byte[] input,
             Operation operation) throws RkpdException, InterruptedException {
-        TrafficStats.setThreadStatsTag(0);
+        final int oldTrafficTag = TrafficStats.getAndSetThreadStatsTag(operation.getTrafficTag());
         int backoff_time = BACKOFF_TIME_MS;
         int attempt = 1;
+        RkpdException lastSeenRkpdException = null;
         try (StopWatch retryTimer = new StopWatch(TAG)) {
             retryTimer.start();
+            // Retry logic.
+            // Provide longer retries (up to 10s) for RkpdExceptions
+            // Provide shorter retries (once) for everything else.
             while (true) {
                 checkDataBudget(metrics);
                 try {
                     Log.v(TAG, "Requesting data from server. Attempt " + attempt);
-                    return requestData(metrics, new URL(Settings.getUrl(mContext) + endpoint),
-                            input);
+                    return requestData(metrics, url, input);
                 } catch (SocketTimeoutException e) {
                     metrics.setStatus(operation.getTimedOutStatus());
                     Log.e(TAG, "Server timed out. " + e.getMessage());
                 } catch (IOException e) {
                     metrics.setStatus(operation.getIoExceptionStatus());
-                    Log.e(TAG, "Failed to complete request from server." + e.getMessage());
+                    Log.e(TAG, "Failed to complete request from server. " + e.getMessage());
                 } catch (RkpdException e) {
+                    lastSeenRkpdException = e;
                     if (e.getErrorCode() == RkpdException.ErrorCode.DEVICE_NOT_REGISTERED) {
                         metrics.setStatus(
                                 ProvisioningAttempt.Status.SIGN_CERTS_DEVICE_NOT_REGISTERED);
                         throw e;
                     } else {
                         metrics.setStatus(operation.getHttpErrorStatus());
-                        if (e.getErrorCode() == RkpdException.ErrorCode.HTTP_CLIENT_ERROR) {
-                            throw e;
-                        }
                     }
                 }
-                if (retryTimer.getElapsedMillis() > Settings.getMaxRequestTime(mContext)) {
+                // Only RkpdExceptions should get longer retries.
+                if (retryTimer.getElapsedMillis() > Settings.getMaxRequestTime(mContext)
+                        || (attempt >= SHORT_RETRY_COUNT && lastSeenRkpdException == null)) {
                     break;
-                } else {
-                    Thread.sleep(backoff_time);
-                    backoff_time *= 2;
-                    attempt += 1;
                 }
+                Thread.sleep(backoff_time);
+                backoff_time *= 2;
+                attempt += 1;
             }
+        } finally {
+            TrafficStats.setThreadStatsTag(oldTrafficTag);
+        }
+        if (lastSeenRkpdException != null) {
+            throw lastSeenRkpdException;
         }
         Settings.incrementFailureCounter(mContext);
         throw makeNetworkError("Error getting data from server.", metrics);
diff --git a/app/src/com/android/rkpdapp/interfaces/SystemInterface.java b/app/src/com/android/rkpdapp/interfaces/SystemInterface.java
index 87cba0a..9bca6e3 100644
--- a/app/src/com/android/rkpdapp/interfaces/SystemInterface.java
+++ b/app/src/com/android/rkpdapp/interfaces/SystemInterface.java
@@ -182,13 +182,13 @@ public class SystemInterface {
 
         int batchSize = mBinder.getHardwareInfo().supportedNumKeysInCsr;
 
-        if (batchSize <= RpcHardwareInfo.MIN_SUPPORTED_NUM_KEYS_IN_CSR) {
+        if (batchSize < RpcHardwareInfo.MIN_SUPPORTED_NUM_KEYS_IN_CSR) {
             Log.w(TAG, "HAL returned a batch size that's too small (" + batchSize
                     + "), defaulting to " + RpcHardwareInfo.MIN_SUPPORTED_NUM_KEYS_IN_CSR);
             return RpcHardwareInfo.MIN_SUPPORTED_NUM_KEYS_IN_CSR;
         }
 
-        if (batchSize >= maxBatchSize) {
+        if (batchSize > maxBatchSize) {
             Log.w(TAG, "HAL returned a batch size that's too large (" + batchSize
                     + "), defaulting to " + maxBatchSize);
             return maxBatchSize;
diff --git a/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java b/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java
index 2bd2296..8c7179f 100644
--- a/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java
+++ b/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java
@@ -36,6 +36,8 @@ import com.android.rkpdapp.metrics.RkpdStatsLog;
 import com.android.rkpdapp.utils.Settings;
 
 import java.time.Instant;
+import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import co.nstant.in.cbor.CborException;
 
@@ -103,8 +105,8 @@ public class PeriodicProvisioner extends Worker {
 
             Log.i(TAG, "Total services found implementing IRPC: " + irpcs.length);
             Provisioner provisioner = new Provisioner(mContext, mKeyDao);
-            Result result = Result.success();
-            for (SystemInterface irpc : irpcs) {
+            final AtomicBoolean result = new AtomicBoolean(true);
+            Arrays.stream(irpcs).parallel().forEach(irpc -> {
                 Log.i(TAG, "Starting provisioning for " + irpc);
                 try {
                     provisioner.provisionKeys(metrics, irpc, response);
@@ -112,13 +114,13 @@ public class PeriodicProvisioner extends Worker {
                     Log.i(TAG, "Successfully provisioned " + irpc);
                 } catch (CborException e) {
                     Log.e(TAG, "Error parsing CBOR for " + irpc, e);
-                    result = Result.failure();
+                    result.set(false);
                 } catch (InterruptedException | RkpdException e) {
                     Log.e(TAG, "Error provisioning keys for " + irpc, e);
-                    result = Result.failure();
+                    result.set(false);
                 }
-            }
-            return result;
+            });
+            return result.get() ? Result.success() : Result.failure();
         }
     }
 
diff --git a/app/src/com/android/rkpdapp/provisioner/Provisioner.java b/app/src/com/android/rkpdapp/provisioner/Provisioner.java
index eb63b50..b8f6057 100644
--- a/app/src/com/android/rkpdapp/provisioner/Provisioner.java
+++ b/app/src/com/android/rkpdapp/provisioner/Provisioner.java
@@ -124,7 +124,7 @@ public class Provisioner {
             throws RkpdException, CborException, InterruptedException {
         int provisionedSoFar = 0;
         List<byte[]> certChains = new ArrayList<>(keysGenerated.size());
-        int maxBatchSize = 0;
+        int maxBatchSize;
         try {
             maxBatchSize = systemInterface.getBatchSize();
         } catch (RemoteException e) {
@@ -162,9 +162,11 @@ public class Provisioner {
             List<RkpKey> keysGenerated) throws RkpdException {
         List<ProvisionedKey> provisionedKeys = new ArrayList<>();
         for (byte[] chain : certChains) {
-            X509Certificate cert = X509Utils.formatX509Certs(chain)[0];
-            long expirationDate = cert.getNotAfter().getTime();
-            byte[] rawPublicKey = X509Utils.getAndFormatRawPublicKey(cert);
+            X509Certificate[] certChain = X509Utils.formatX509Certs(chain);
+            X509Certificate leafCertificate = certChain[0];
+            long expirationDate = X509Utils.getExpirationTimeForCertificateChain(certChain)
+                    .getTime();
+            byte[] rawPublicKey = X509Utils.getAndFormatRawPublicKey(leafCertificate);
             if (rawPublicKey == null) {
                 Log.e(TAG, "Skipping malformed public key.");
                 continue;
diff --git a/app/src/com/android/rkpdapp/utils/X509Utils.java b/app/src/com/android/rkpdapp/utils/X509Utils.java
index 4e12872..61b55c2 100644
--- a/app/src/com/android/rkpdapp/utils/X509Utils.java
+++ b/app/src/com/android/rkpdapp/utils/X509Utils.java
@@ -40,6 +40,7 @@ import java.security.cert.X509Certificate;
 import java.security.interfaces.ECPublicKey;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Date;
 import java.util.Set;
 
 /**
@@ -158,4 +159,16 @@ public class X509Utils {
             return false;
         }
     }
+
+    /**
+     * Gets the date when the entire cert chain expires. This can be calculated by
+     * checking the individual certificate times and returning the minimum of all.
+     * @return Expiration time for the certificate chain
+     */
+    public static Date getExpirationTimeForCertificateChain(X509Certificate[] certChain) {
+        return Arrays.stream(certChain)
+                .map(X509Certificate::getNotAfter)
+                .min(Date::compareTo)
+                .orElse(null);
+    }
 }
diff --git a/app/tests/e2e/src/com/android/rkpdapp/e2etest/KeystoreIntegrationTest.java b/app/tests/e2e/src/com/android/rkpdapp/e2etest/KeystoreIntegrationTest.java
index 5a8e728..46bd879 100644
--- a/app/tests/e2e/src/com/android/rkpdapp/e2etest/KeystoreIntegrationTest.java
+++ b/app/tests/e2e/src/com/android/rkpdapp/e2etest/KeystoreIntegrationTest.java
@@ -41,11 +41,11 @@ import androidx.work.testing.TestWorkerBuilder;
 import com.android.rkpdapp.database.ProvisionedKey;
 import com.android.rkpdapp.database.ProvisionedKeyDao;
 import com.android.rkpdapp.database.RkpdDatabase;
+import com.android.rkpdapp.interfaces.ServerInterface;
 import com.android.rkpdapp.interfaces.ServiceManagerInterface;
 import com.android.rkpdapp.interfaces.SystemInterface;
 import com.android.rkpdapp.provisioner.PeriodicProvisioner;
 import com.android.rkpdapp.testutil.FakeRkpServer;
-import com.android.rkpdapp.testutil.NetworkUtils;
 import com.android.rkpdapp.testutil.SystemPropertySetter;
 import com.android.rkpdapp.utils.Settings;
 import com.android.rkpdapp.utils.X509Utils;
@@ -108,15 +108,20 @@ public class KeystoreIntegrationTest {
     @BeforeClass
     public static void init() {
         sContext = ApplicationProvider.getApplicationContext();
+    }
 
+    @Before
+    public void setUp() throws Exception {
         assume()
                 .withMessage("The RKP server hostname is not configured -- assume RKP disabled.")
                 .that(SystemProperties.get("remote_provisioning.hostname"))
                 .isNotEmpty();
-    }
 
-    @Before
-    public void setUp() throws Exception {
+        assume()
+                .withMessage("RKP Integration tests rely on network availability.")
+                .that(ServerInterface.isNetworkConnected(sContext))
+                .isTrue();
+
         Settings.clearPreferences(sContext);
 
         mKeyDao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
@@ -132,8 +137,10 @@ public class KeystoreIntegrationTest {
     public void tearDown() throws Exception {
         Settings.clearPreferences(sContext);
 
-        mKeyStore.deleteEntry(getTestKeyAlias());
-        mKeyDao.deleteAllKeys();
+        if (mKeyDao != null) {
+            mKeyStore.deleteEntry(getTestKeyAlias());
+            mKeyDao.deleteAllKeys();
+        }
 
         ServiceManagerInterface.setInstances(null);
     }
@@ -229,13 +236,8 @@ public class KeystoreIntegrationTest {
             assertWithMessage("Should have gotten a KeyStoreException").fail();
         } catch (ProviderException e) {
             assertThat(e.getCause()).isInstanceOf(KeyStoreException.class);
-            if (NetworkUtils.isNetworkConnected(sContext)) {
-                assertThat(((KeyStoreException) e.getCause()).getErrorCode())
-                        .isEqualTo(ResponseCode.OUT_OF_KEYS_TRANSIENT_ERROR);
-            } else {
-                assertThat(((KeyStoreException) e.getCause()).getErrorCode())
-                        .isEqualTo(ResponseCode.OUT_OF_KEYS_PENDING_INTERNET_CONNECTIVITY);
-            }
+            assertThat(((KeyStoreException) e.getCause()).getErrorCode())
+                    .isEqualTo(ResponseCode.OUT_OF_KEYS_TRANSIENT_ERROR);
         }
     }
 
diff --git a/app/tests/e2e/src/com/android/rkpdapp/e2etest/RkpdHostTestHelperTests.java b/app/tests/e2e/src/com/android/rkpdapp/e2etest/RkpdHostTestHelperTests.java
index 51b5e5b..0eb76be 100644
--- a/app/tests/e2e/src/com/android/rkpdapp/e2etest/RkpdHostTestHelperTests.java
+++ b/app/tests/e2e/src/com/android/rkpdapp/e2etest/RkpdHostTestHelperTests.java
@@ -38,6 +38,7 @@ import androidx.work.testing.TestWorkerBuilder;
 import com.android.rkpdapp.database.ProvisionedKey;
 import com.android.rkpdapp.database.ProvisionedKeyDao;
 import com.android.rkpdapp.database.RkpdDatabase;
+import com.android.rkpdapp.interfaces.ServerInterface;
 import com.android.rkpdapp.interfaces.ServiceManagerInterface;
 import com.android.rkpdapp.interfaces.SystemInterface;
 import com.android.rkpdapp.provisioner.PeriodicProvisioner;
@@ -58,7 +59,6 @@ import org.junit.runners.Parameterized;
 import java.security.KeyPairGenerator;
 import java.security.KeyStore;
 import java.security.spec.ECGenParameterSpec;
-import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
 import java.util.concurrent.Executors;
@@ -89,15 +89,20 @@ public class RkpdHostTestHelperTests {
     @BeforeClass
     public static void init() {
         sContext = ApplicationProvider.getApplicationContext();
+    }
 
+    @Before
+    public void setUp() throws Exception {
         assume()
                 .withMessage("The RKP server hostname is not configured -- assume RKP disabled.")
                 .that(SystemProperties.get("remote_provisioning.hostname"))
                 .isNotEmpty();
-    }
 
-    @Before
-    public void setUp() throws Exception {
+        assume()
+                .withMessage("RKP Integration tests rely on network availability.")
+                .that(ServerInterface.isNetworkConnected(sContext))
+                .isTrue();
+
         Settings.clearPreferences(sContext);
         mRealDao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
         mRealDao.deleteAllKeys();
@@ -114,7 +119,10 @@ public class RkpdHostTestHelperTests {
     @After
     public void tearDown() throws Exception {
         Settings.clearPreferences(sContext);
-        mRealDao.deleteAllKeys();
+
+        if (mRealDao != null) {
+            mRealDao.deleteAllKeys();
+        }
 
         KeyStore keyStore = KeyStore.getInstance("AndroidKeyStore");
         keyStore.load(null);
@@ -144,8 +152,6 @@ public class RkpdHostTestHelperTests {
     public void provisionThenExpireThenProvisionAgain() throws Exception {
         assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.success());
 
-        final Instant expiry = Instant.now().plus(Duration.ofHours(1));
-
         List<ProvisionedKey> keys = mTestDao.getAllKeys();
 
         // Expire a key
@@ -177,6 +183,8 @@ public class RkpdHostTestHelperTests {
         StatsProcessor.PoolStats updatedPool = StatsProcessor.processPool(mRealDao, mServiceName,
                 Settings.getExtraSignedKeysAvailable(sContext),
                 Settings.getExpirationTime(sContext));
-        assertThat(updatedPool.toString()).isEqualTo(pool.toString());
+
+        assertThat(updatedPool.keysInUse + updatedPool.keysUnassigned)
+                .isEqualTo(pool.keysInUse + pool.keysUnassigned);
     }
 }
diff --git a/app/tests/unit/src/com/android/rkpdapp/unittest/PeriodicProvisionerTests.java b/app/tests/unit/src/com/android/rkpdapp/unittest/PeriodicProvisionerTests.java
index 9a51bec..a278f50 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/PeriodicProvisionerTests.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/PeriodicProvisionerTests.java
@@ -24,7 +24,6 @@ import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import android.content.Context;
 
@@ -46,6 +45,7 @@ import com.android.rkpdapp.database.RkpdDatabase;
 import com.android.rkpdapp.interfaces.ServiceManagerInterface;
 import com.android.rkpdapp.interfaces.SystemInterface;
 import com.android.rkpdapp.provisioner.PeriodicProvisioner;
+import com.android.rkpdapp.service.RegistrationBinder;
 import com.android.rkpdapp.testutil.FakeRkpServer;
 import com.android.rkpdapp.testutil.SystemPropertySetter;
 import com.android.rkpdapp.utils.Settings;
@@ -176,8 +176,8 @@ public class PeriodicProvisionerTests {
             ServiceManagerInterface.setInstances(new SystemInterface[]{mockHal});
             assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.failure());
 
-            // we should have failed before making any local HAl calls
-            verifyNoMoreInteractions(mockHal);
+            // we should have failed before trying to generate any keys
+            verify(mockHal, never()).generateKey(any());
         }
     }
 
@@ -197,8 +197,8 @@ public class PeriodicProvisionerTests {
             ServiceManagerInterface.setInstances(new SystemInterface[]{mockHal});
             assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.success());
 
-            // since RKP is disabled, there should be no interactions with the HAL
-            verifyNoMoreInteractions(mockHal);
+            // since RKP is disabled, there should be no keys generated
+            verify(mockHal, never()).generateKey(any());
         }
 
         // when RKP is detected as disabled, the provisioner is supposed to delete all keys
@@ -209,9 +209,12 @@ public class PeriodicProvisionerTests {
     public void provisioningExpiresOldKeys() throws Exception {
         ProvisionedKeyDao dao = RkpdDatabase.getDatabase(mContext).provisionedKeyDao();
         ProvisionedKey oldKey = new ProvisionedKey(new byte[1], "fake-irpc", new byte[2],
-                new byte[3], Instant.now().minusSeconds(120));
+                new byte[3],
+                Instant.now().minus(RegistrationBinder.MIN_KEY_LIFETIME.multipliedBy(2)));
+        // Add 2 hours so that this key does not get deleted in case getKeyWorker comes alive.
         ProvisionedKey freshKey = new ProvisionedKey(new byte[11], "fake-irpc", new byte[12],
-                new byte[13], Instant.now().plusSeconds(120));
+                new byte[13],
+                Instant.now().plus(RegistrationBinder.MIN_KEY_LIFETIME.multipliedBy(2)));
         dao.insertKeys(List.of(oldKey, freshKey));
         assertThat(dao.getTotalKeysForIrpc("fake-irpc")).isEqualTo(2);
 
diff --git a/app/tests/unit/src/com/android/rkpdapp/unittest/ServerInterfaceTest.java b/app/tests/unit/src/com/android/rkpdapp/unittest/ServerInterfaceTest.java
index d7971df..725b1a6 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/ServerInterfaceTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/ServerInterfaceTest.java
@@ -21,7 +21,7 @@ import static com.google.common.truth.Truth.assertWithMessage;
 
 import android.content.Context;
 import android.net.ConnectivityManager;
-import android.net.NetworkInfo;
+import android.net.NetworkCapabilities;
 import android.util.Base64;
 
 import androidx.test.core.app.ApplicationProvider;
@@ -81,7 +81,8 @@ public class ServerInterfaceTest {
                     ProvisioningAttempt.createScheduledAttemptMetrics(sContext));
             assertWithMessage("Expected RkpdException.").fail();
         } catch (RkpdException e) {
-            // should throw this
+            assertThat(e.getErrorCode()).isEqualTo(RkpdException.ErrorCode.HTTP_SERVER_ERROR);
+            assertThat(e).hasMessageThat().contains("HTTP error status encountered");
         }
     }
 
@@ -183,6 +184,7 @@ public class ServerInterfaceTest {
                 FakeRkpServer.Response.SIGN_CERTS_USER_UNAUTHORIZED)) {
             Settings.setDeviceConfig(sContext, 2 /* extraKeys */,
                     TIME_TO_REFRESH_HOURS /* expiringBy */, server.getUrl());
+            Settings.setMaxRequestTime(sContext, 100);
             ProvisioningAttempt metrics = ProvisioningAttempt.createScheduledAttemptMetrics(
                     sContext);
             mServerInterface.requestSignedCertificates(new byte[0], new byte[0], metrics);
@@ -227,11 +229,18 @@ public class ServerInterfaceTest {
 
     @Test
     public void testDataBudgetEmptyFetchGeekNetworkConnected() throws Exception {
-        // Check the data budget in order to initialize a rolling window.
-        assertThat(Settings.hasErrDataBudget(sContext, null /* curTime */)).isTrue();
-        Settings.consumeErrDataBudget(sContext, Settings.FAILURE_DATA_USAGE_MAX);
-        ProvisioningAttempt metrics = ProvisioningAttempt.createScheduledAttemptMetrics(sContext);
-        try {
+        try (FakeRkpServer server = new FakeRkpServer(
+                FakeRkpServer.Response.FETCH_EEK_OK,
+                FakeRkpServer.Response.SIGN_CERTS_OK_VALID_CBOR)) {
+            Settings.setDeviceConfig(sContext, 2 /* extraKeys */,
+                    TIME_TO_REFRESH_HOURS /* expiringBy */, server.getUrl());
+
+            // Check the data budget in order to initialize a rolling window.
+            assertThat(Settings.hasErrDataBudget(sContext, null /* curTime */)).isTrue();
+            Settings.consumeErrDataBudget(sContext, Settings.FAILURE_DATA_USAGE_MAX);
+            ProvisioningAttempt metrics = ProvisioningAttempt.createScheduledAttemptMetrics(
+                    sContext);
+
             // We are okay in mocking connectivity failure since err data budget is the first thing
             // to be checked.
             mockConnectivityFailure(ConnectivityState.CONNECTED);
@@ -246,15 +255,21 @@ public class ServerInterfaceTest {
 
     @Test
     public void testDataBudgetEmptyFetchGeekNetworkDisconnected() throws Exception {
-        // Check the data budget in order to initialize a rolling window.
-        try {
-            // We are okay in mocking connectivity failure since err data budget is the first thing
-            // to be checked.
-            mockConnectivityFailure(ConnectivityState.DISCONNECTED);
+        try (FakeRkpServer server = new FakeRkpServer(
+                FakeRkpServer.Response.FETCH_EEK_OK,
+                FakeRkpServer.Response.SIGN_CERTS_OK_VALID_CBOR)) {
+            Settings.setDeviceConfig(sContext, 2 /* extraKeys */,
+                    TIME_TO_REFRESH_HOURS /* expiringBy */, server.getUrl());
+
+            // Check the data budget in order to initialize a rolling window.
             assertThat(Settings.hasErrDataBudget(sContext, null /* curTime */)).isTrue();
             Settings.consumeErrDataBudget(sContext, Settings.FAILURE_DATA_USAGE_MAX);
             ProvisioningAttempt metrics = ProvisioningAttempt.createScheduledAttemptMetrics(
                     sContext);
+
+            // We are okay in mocking connectivity failure since err data budget is the first thing
+            // to be checked.
+            mockConnectivityFailure(ConnectivityState.DISCONNECTED);
             mServerInterface.fetchGeek(metrics);
             assertWithMessage("Network transaction should not have proceeded.").fail();
         } catch (RkpdException e) {
@@ -376,12 +391,16 @@ public class ServerInterfaceTest {
 
     private void mockConnectivityFailure(ConnectivityState state) {
         ConnectivityManager mockedConnectivityManager = Mockito.mock(ConnectivityManager.class);
-        NetworkInfo mockedNetwork = Mockito.mock(NetworkInfo.class);
 
         Mockito.when(sContext.getSystemService(ConnectivityManager.class))
                 .thenReturn(mockedConnectivityManager);
-        Mockito.when(mockedConnectivityManager.getActiveNetworkInfo()).thenReturn(mockedNetwork);
-        Mockito.when(mockedNetwork.isConnected()).thenReturn(state == ConnectivityState.CONNECTED);
+        NetworkCapabilities.Builder builder = new NetworkCapabilities.Builder();
+        builder.addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET);
+        if (state == ConnectivityState.CONNECTED) {
+            builder.addCapability(NetworkCapabilities.NET_CAPABILITY_VALIDATED);
+        }
+        Mockito.when(mockedConnectivityManager.getNetworkCapabilities(Mockito.any()))
+                .thenReturn(builder.build());
     }
 
     private enum ConnectivityState {
diff --git a/app/tests/unit/src/com/android/rkpdapp/unittest/Utils.java b/app/tests/unit/src/com/android/rkpdapp/unittest/Utils.java
index d478db7..bbce2d9 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/Utils.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/Utils.java
@@ -144,6 +144,12 @@ public class Utils {
 
     public static X509Certificate signPublicKey(KeyPair issuerKeyPair, PublicKey publicKeyToSign)
             throws Exception {
+        return signPublicKey(issuerKeyPair, publicKeyToSign,
+                Instant.now().plus(Duration.ofDays(1)));
+    }
+
+    public static X509Certificate signPublicKey(KeyPair issuerKeyPair, PublicKey publicKeyToSign,
+            Instant expirationInstant) throws Exception {
         X500Principal issuer = new X500Principal("CN=TEE");
         BigInteger serial = BigInteger.ONE;
         X500Principal subject = new X500Principal("CN=TEE");
@@ -153,7 +159,7 @@ public class Utils {
         certificateBuilder.setIssuerDN(issuer);
         certificateBuilder.setSerialNumber(serial);
         certificateBuilder.setNotBefore(Date.from(now));
-        certificateBuilder.setNotAfter(Date.from(now.plus(Duration.ofDays(1))));
+        certificateBuilder.setNotAfter(Date.from(expirationInstant));
         certificateBuilder.setSignatureAlgorithm("SHA256WITHECDSA");
         certificateBuilder.setSubjectDN(subject);
         certificateBuilder.setPublicKey(publicKeyToSign);
diff --git a/app/tests/unit/src/com/android/rkpdapp/unittest/X509UtilsTest.java b/app/tests/unit/src/com/android/rkpdapp/unittest/X509UtilsTest.java
index d4f9d0a..83378e4 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/X509UtilsTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/X509UtilsTest.java
@@ -39,6 +39,9 @@ import java.security.KeyPair;
 import java.security.cert.CertificateException;
 import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Date;
 
 @RunWith(AndroidJUnit4.class)
 public class X509UtilsTest {
@@ -133,6 +136,36 @@ public class X509UtilsTest {
         assertThat(X509Utils.isCertChainValid(certChain)).isFalse();
     }
 
+    @Test
+    public void testCertChainExpirationTimeWhenRootExpiresLater() throws Exception {
+        KeyPair root = generateEcdsaKeyPair();
+        KeyPair leaf = generateEcdsaKeyPair();
+        X509Certificate[] certs = new X509Certificate[2];
+
+        // root cert expires later.
+        certs[0] = signPublicKey(root, leaf.getPublic(), Instant.now().plus(Duration.ofDays(2)));
+        certs[1] = signPublicKey(root, root.getPublic(), Instant.now().plus(Duration.ofDays(1)));
+
+        Date expirationTime = X509Utils.getExpirationTimeForCertificateChain(certs);
+        assertThat(expirationTime).isEqualTo(certs[1].getNotAfter());
+        assertThat(expirationTime).isNotEqualTo(certs[0].getNotAfter());
+    }
+
+    @Test
+    public void testCertChainExpirationTimeWhenLeafExpiresLater() throws Exception {
+        KeyPair root = generateEcdsaKeyPair();
+        KeyPair leaf = generateEcdsaKeyPair();
+        X509Certificate[] certs = new X509Certificate[2];
+
+        // leaf cert expires later.
+        certs[0] = signPublicKey(root, leaf.getPublic(), Instant.now().plus(Duration.ofDays(1)));
+        certs[1] = signPublicKey(root, root.getPublic(), Instant.now().plus(Duration.ofDays(2)));
+
+        Date expirationTime = X509Utils.getExpirationTimeForCertificateChain(certs);
+        assertThat(expirationTime).isEqualTo(certs[0].getNotAfter());
+        assertThat(expirationTime).isNotEqualTo(certs[1].getNotAfter());
+    }
+
     private X509Certificate generateCertificateFromEncodedBytes(String encodedCert)
             throws CertificateException {
         CertificateFactory certFactory = CertificateFactory.getInstance("X.509");
diff --git a/app/tests/util/src/com/android/rkpdapp/testutil/NetworkUtils.java b/app/tests/util/src/com/android/rkpdapp/testutil/NetworkUtils.java
deleted file mode 100644
index 77a5d96..0000000
--- a/app/tests/util/src/com/android/rkpdapp/testutil/NetworkUtils.java
+++ /dev/null
@@ -1,32 +0,0 @@
-/*
- * Copyright (C) 2023 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.rkpdapp.testutil;
-
-import android.content.Context;
-import android.net.ConnectivityManager;
-import android.net.NetworkInfo;
-
-public class NetworkUtils {
-    public static boolean isNetworkConnected(Context context) {
-        ConnectivityManager cm = context.getSystemService(ConnectivityManager.class);
-        NetworkInfo networkInfo = cm.getActiveNetworkInfo();
-        if (networkInfo != null && networkInfo.isConnected()) {
-            return true;
-        }
-        return false;
-    }
-}
diff --git a/system-server/tests/unit/AndroidTest.xml b/system-server/tests/unit/AndroidTest.xml
new file mode 100644
index 0000000..5c83c54
--- /dev/null
+++ b/system-server/tests/unit/AndroidTest.xml
@@ -0,0 +1,36 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- 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.
+-->
+<configuration description="Unit tests for the rkpd system server.">
+    <option name="test-suite-tag" value="apct" />
+    <option name="test-suite-tag" value="apct-instrumentation" />
+
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="cleanup-apks" value="true" />
+        <option name="test-file-name" value="service-rkp-unittest.apk" />
+    </target_preparer>
+
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="android.security.rkp.service.test" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+    </test>
+
+    <!-- Only run if RKPD mainline module is installed -->
+    <object type="module_controller"
+            class="com.android.tradefed.testtype.suite.module.MainlineTestModuleController">
+        <option name="enable" value="true" />
+        <option name="mainline-module-package-name" value="com.android.rkpd" />
+    </object>
+</configuration>