Merge "Add connection_timeout system property." into main
diff --git a/app/src/com/android/rkpdapp/GeekResponse.java b/app/src/com/android/rkpdapp/GeekResponse.java
index 3e40436..be2c13d 100644
--- a/app/src/com/android/rkpdapp/GeekResponse.java
+++ b/app/src/com/android/rkpdapp/GeekResponse.java
@@ -17,6 +17,7 @@
 package com.android.rkpdapp;
 
 import java.time.Duration;
+import java.time.Instant;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -49,6 +50,8 @@
     public int numExtraAttestationKeys;
     public Duration timeToRefresh;
     public String provisioningUrl;
+    public Instant lastBadCertTimeStart;
+    public Instant lastBadCertTimeEnd;
 
     /**
      * Default initializer.
@@ -56,6 +59,8 @@
     public GeekResponse() {
         mCurveToGeek = new HashMap<>();
         numExtraAttestationKeys = NO_EXTRA_KEY_UPDATE;
+        lastBadCertTimeStart = null;
+        lastBadCertTimeEnd = null;
     }
 
     /**
diff --git a/app/src/com/android/rkpdapp/database/ProvisionedKeyDao.java b/app/src/com/android/rkpdapp/database/ProvisionedKeyDao.java
index 947f4eb..af18a25 100644
--- a/app/src/com/android/rkpdapp/database/ProvisionedKeyDao.java
+++ b/app/src/com/android/rkpdapp/database/ProvisionedKeyDao.java
@@ -45,10 +45,16 @@
     public abstract void updateKey(ProvisionedKey key);
 
     /**
-     * Delete all expiring keys provided by given Instant.
+     * Gets all the keys in the database.
      */
-    @Query("DELETE FROM provisioned_keys WHERE expiration_time < :expiryTime")
-    public abstract void deleteExpiringKeys(Instant expiryTime);
+    @Query("SELECT * FROM provisioned_keys")
+    public abstract List<ProvisionedKey> getAllKeys();
+
+    /**
+     * Deletes a specific key from the database.
+     */
+    @Query("DELETE from provisioned_keys WHERE key_blob = :keyBlob")
+    public abstract void deleteKey(byte[] keyBlob);
 
     /**
      * Delete all the provisioned keys.
@@ -57,6 +63,12 @@
     public abstract void deleteAllKeys();
 
     /**
+     * Delete all expiring keys provided by given Instant.
+     */
+    @Query("DELETE FROM provisioned_keys WHERE expiration_time < :expiryTime")
+    public abstract void deleteExpiringKeys(Instant expiryTime);
+
+    /**
      * Get a count of provisioned keys for a specific IRPC that are expiring at a given Instant.
      */
     @Query("SELECT COUNT(*) FROM provisioned_keys"
diff --git a/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java b/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java
index 3e4a7a3..cb1a979 100644
--- a/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java
+++ b/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java
@@ -131,6 +131,8 @@
 
             Log.i(TAG, "Total services found implementing IRPC: " + irpcs.length);
             Provisioner provisioner = new Provisioner(mContext, mKeyDao, IS_ASYNC);
+            provisioner.clearBadAttestationKeys(response);
+
             final AtomicBoolean result = new AtomicBoolean(true);
             Arrays.stream(irpcs).parallel().forEach(irpc -> {
                 Log.i(TAG, "Starting provisioning for " + irpc);
diff --git a/app/src/com/android/rkpdapp/provisioner/Provisioner.java b/app/src/com/android/rkpdapp/provisioner/Provisioner.java
index 19db3a7..2119c0c 100644
--- a/app/src/com/android/rkpdapp/provisioner/Provisioner.java
+++ b/app/src/com/android/rkpdapp/provisioner/Provisioner.java
@@ -35,6 +35,7 @@
 
 import java.security.cert.X509Certificate;
 import java.time.Instant;
+import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -168,7 +169,7 @@
             X509Certificate[] certChain = X509Utils.formatX509Certs(chain);
             X509Certificate leafCertificate = certChain[0];
             long expirationDate = X509Utils.getExpirationTimeForCertificateChain(certChain)
-                    .getTime();
+                    .toInstant().toEpochMilli();
             byte[] rawPublicKey = X509Utils.getAndFormatRawPublicKey(leafCertificate);
             if (rawPublicKey == null) {
                 Log.e(TAG, "Skipping malformed public key.");
@@ -203,4 +204,44 @@
             throw new InterruptedException();
         }
     }
+
+    /**
+     * Clears bad attestation keys on the basis of information provided in the FetchGeek response.
+     */
+    public void clearBadAttestationKeys(GeekResponse resp) {
+        if (resp.lastBadCertTimeStart == null || resp.lastBadCertTimeEnd == null) {
+            // if there is no time sent, no need to do anything.
+            return;
+        }
+        if (resp.lastBadCertTimeStart.equals(Settings.getLastBadCertTimeStart(mContext))
+                && resp.lastBadCertTimeEnd.equals(Settings.getLastBadCertTimeEnd(mContext))) {
+            // if the time is same as already stored version, no need to do anything.
+            return;
+        }
+        // clear the attestation keys on the basis of time.
+        checkAndDeleteBadKeys(resp.lastBadCertTimeStart, resp.lastBadCertTimeEnd);
+
+        // store the time.
+        Settings.setLastBadCertTimeRange(mContext, resp.lastBadCertTimeStart,
+                resp.lastBadCertTimeEnd);
+    }
+
+    private void checkAndDeleteBadKeys(Instant startTime, Instant endTime) {
+        try {
+            List<ProvisionedKey> allKeys = mKeyDao.getAllKeys();
+            for (int i = 0; i < allKeys.size(); i++) {
+                ProvisionedKey key = allKeys.get(i);
+                X509Certificate[] certChain = X509Utils.formatX509Certs(key.certificateChain);
+                X509Certificate leafCertificate = certChain[0];
+                Instant creationTime = leafCertificate.getNotBefore().toInstant()
+                        .truncatedTo(ChronoUnit.MILLIS);
+
+                if (!creationTime.isBefore(startTime) && !creationTime.isAfter(endTime)) {
+                    mKeyDao.deleteKey(key.keyBlob);
+                }
+            }
+        } catch (RkpdException ex) {
+            Log.e(TAG, "Could not convert certificate chain to X509 certificates.", ex);
+        }
+    }
 }
diff --git a/app/src/com/android/rkpdapp/utils/CborUtils.java b/app/src/com/android/rkpdapp/utils/CborUtils.java
index dc9eab5..2eee6d6 100644
--- a/app/src/com/android/rkpdapp/utils/CborUtils.java
+++ b/app/src/com/android/rkpdapp/utils/CborUtils.java
@@ -23,6 +23,7 @@
 
 import com.android.rkpdapp.GeekResponse;
 import com.android.rkpdapp.RkpdException;
+import com.android.rkpdapp.database.InstantConverter;
 import com.android.rkpdapp.database.RkpKey;
 
 import java.io.ByteArrayInputStream;
@@ -52,6 +53,8 @@
     public static final String EXTRA_KEYS = "num_extra_attestation_keys";
     public static final String TIME_TO_REFRESH = "time_to_refresh_hours";
     public static final String PROVISIONING_URL = "provisioning_url";
+    public static final String LAST_BAD_CERT_TIME_START_MILLIS = "bad_cert_start";
+    public static final String LAST_BAD_CERT_TIME_END_MILLIS = "bad_cert_end";
 
     private static final int RESPONSE_CERT_ARRAY_INDEX = 0;
     private static final int RESPONSE_ARRAY_SIZE = 1;
@@ -155,6 +158,10 @@
                 deviceConfiguration.get(new UnicodeString(TIME_TO_REFRESH));
         DataItem newUrl =
                 deviceConfiguration.get(new UnicodeString(PROVISIONING_URL));
+        DataItem lastBadCertTimeStart =
+                deviceConfiguration.get(new UnicodeString(LAST_BAD_CERT_TIME_START_MILLIS));
+        DataItem lastBadCertTimeEnd =
+                deviceConfiguration.get(new UnicodeString(LAST_BAD_CERT_TIME_END_MILLIS));
         if (extraKeys != null) {
             if (!checkType(extraKeys, MajorType.UNSIGNED_INTEGER, "ExtraKeys")) {
                 return false;
@@ -174,6 +181,20 @@
             }
             resp.provisioningUrl = ((UnicodeString) newUrl).getString();
         }
+        if (lastBadCertTimeStart != null) {
+            if (!checkType(lastBadCertTimeStart, MajorType.UNSIGNED_INTEGER, "BadCertTimeStart")) {
+                return false;
+            }
+            resp.lastBadCertTimeStart = InstantConverter.fromTimestamp(
+                    ((UnsignedInteger) lastBadCertTimeStart).getValue().longValue());
+        }
+        if (lastBadCertTimeEnd != null) {
+            if (!checkType(lastBadCertTimeEnd, MajorType.UNSIGNED_INTEGER, "BadCertTimeEnd")) {
+                return false;
+            }
+            resp.lastBadCertTimeEnd = InstantConverter.fromTimestamp(
+                    ((UnsignedInteger) lastBadCertTimeEnd).getValue().longValue());
+        }
         return true;
     }
 
diff --git a/app/src/com/android/rkpdapp/utils/Settings.java b/app/src/com/android/rkpdapp/utils/Settings.java
index c421909..df41306 100644
--- a/app/src/com/android/rkpdapp/utils/Settings.java
+++ b/app/src/com/android/rkpdapp/utils/Settings.java
@@ -22,6 +22,7 @@
 import android.util.Log;
 
 import com.android.rkpdapp.GeekResponse;
+import com.android.rkpdapp.database.InstantConverter;
 
 import java.net.MalformedURLException;
 import java.net.URL;
@@ -45,6 +46,7 @@
     public static final int FAILURE_DATA_USAGE_MAX = 1024 * 1024;
     public static final Duration FAILURE_DATA_USAGE_WINDOW = Duration.ofDays(1);
     public static final int MAX_REQUEST_TIME_MS_DEFAULT = 20000;
+    public static final int NO_TIME_PROVIDED = -1;
 
     private static final String KEY_EXPIRING_BY = "expiring_by";
     private static final String KEY_EXTRA_KEYS = "extra_keys";
@@ -54,6 +56,8 @@
     private static final String KEY_FAILURE_BYTES = "failure_data";
     private static final String KEY_URL = "url";
     private static final String KEY_MAX_REQUEST_TIME = "max_request_time";
+    private static final String KEY_LAST_BAD_CERT_TIME_START = "bad_cert_time_start";
+    private static final String KEY_LAST_BAD_CERT_TIME_END = "bad_cert_time_end";
     private static final String PREFERENCES_NAME = "com.android.rkpdapp.utils.preferences";
     private static final String TAG = "RkpdSettings";
 
@@ -63,7 +67,7 @@
      * resulted in errors, then false will be returned.
      * <p>
      * Additionally, the rolling window of data usage is managed within this call. The used data
-     * budget will be reset if a time greater than @{code FAILURE_DATA_USAGE_WINDOW} has passed.
+     * budget will be reset if a time greater than {@code FAILURE_DATA_USAGE_WINDOW} has passed.
      *
      * @param context The application context
      * @param curTime An instant representing the current time to measure the window against. If
@@ -303,6 +307,58 @@
     }
 
     /**
+     * Gets start time in milliseconds when the bad certificates were provided by server.
+     */
+    public static Instant getLastBadCertTimeStart(Context context) {
+        SharedPreferences sharedPref = getSharedPreferences(context);
+        long lastBadCertTimeStartMillis =
+                sharedPref.getLong(KEY_LAST_BAD_CERT_TIME_START, NO_TIME_PROVIDED);
+        if (lastBadCertTimeStartMillis != -1) {
+            return InstantConverter.fromTimestamp(lastBadCertTimeStartMillis);
+        } else {
+            return null;
+        }
+    }
+
+    /**
+     * Gets end time in milliseconds when the bad certificates were provided by server.
+     */
+    public static Instant getLastBadCertTimeEnd(Context context) {
+        SharedPreferences sharedPref = getSharedPreferences(context);
+        long lastBadCertTimeEndMillis =
+                sharedPref.getLong(KEY_LAST_BAD_CERT_TIME_END, NO_TIME_PROVIDED);
+        if (lastBadCertTimeEndMillis != -1) {
+            return InstantConverter.fromTimestamp(lastBadCertTimeEndMillis);
+        } else {
+            return null;
+        }
+    }
+
+    /**
+     * Sets the time range for the last bad certificates.
+     */
+    public static void setLastBadCertTimeRange(Context context, Instant lastBadCertTimeStart,
+            Instant lastBadCertTimeEnd) {
+        long startMillis = lastBadCertTimeStart.toEpochMilli();
+        long endMillis = lastBadCertTimeEnd.toEpochMilli();
+        SharedPreferences sharedPref = getSharedPreferences(context);
+        SharedPreferences.Editor editor = sharedPref.edit();
+
+        boolean isUpdated = false;
+        if (sharedPref.getLong(KEY_LAST_BAD_CERT_TIME_START, NO_TIME_PROVIDED) != startMillis) {
+            editor.putLong(KEY_LAST_BAD_CERT_TIME_START, startMillis);
+            isUpdated = true;
+        }
+        if (sharedPref.getLong(KEY_LAST_BAD_CERT_TIME_END, NO_TIME_PROVIDED) != endMillis) {
+            editor.putLong(KEY_LAST_BAD_CERT_TIME_END, endMillis);
+            isUpdated = true;
+        }
+        if (isUpdated) {
+            editor.apply();
+        }
+    }
+
+    /**
      * Clears all preferences, thus restoring the defaults.
      */
     public static void clearPreferences(Context context) {
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 21d0c78..4a35efc 100644
--- a/app/tests/e2e/src/com/android/rkpdapp/e2etest/RkpdHostTestHelperTests.java
+++ b/app/tests/e2e/src/com/android/rkpdapp/e2etest/RkpdHostTestHelperTests.java
@@ -19,8 +19,6 @@
 import static android.security.keystore.KeyProperties.KEY_ALGORITHM_EC;
 import static android.security.keystore.KeyProperties.PURPOSE_SIGN;
 
-import static com.android.rkpdapp.database.RkpdDatabase.DB_NAME;
-
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.TruthJUnit.assume;
 
@@ -30,7 +28,6 @@
 import android.os.SystemProperties;
 import android.security.keystore.KeyGenParameterSpec;
 
-import androidx.room.Room;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.work.ListenableWorker;
 import androidx.work.testing.TestWorkerBuilder;
@@ -43,8 +40,6 @@
 import com.android.rkpdapp.interfaces.SystemInterface;
 import com.android.rkpdapp.provisioner.PeriodicProvisioner;
 import com.android.rkpdapp.testutil.SystemInterfaceSelector;
-import com.android.rkpdapp.testutil.TestDatabase;
-import com.android.rkpdapp.testutil.TestProvisionedKeyDao;
 import com.android.rkpdapp.utils.Settings;
 import com.android.rkpdapp.utils.StatsProcessor;
 
@@ -70,8 +65,7 @@
     private static Context sContext;
     private final String mInstanceName;
     private final String mServiceName;
-    private ProvisionedKeyDao mRealDao;
-    private TestProvisionedKeyDao mTestDao;
+    private ProvisionedKeyDao mKeyDao;
     private PeriodicProvisioner mProvisioner;
     private AutoCloseable mPeriodicProvisionerLock;
 
@@ -107,9 +101,8 @@
 
         mPeriodicProvisionerLock = PeriodicProvisioner.lock();
         Settings.clearPreferences(sContext);
-        mRealDao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
-        mRealDao.deleteAllKeys();
-        mTestDao = Room.databaseBuilder(sContext, TestDatabase.class, DB_NAME).build().dao();
+        mKeyDao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
+        mKeyDao.deleteAllKeys();
 
         mProvisioner = TestWorkerBuilder.from(
                 sContext,
@@ -125,8 +118,8 @@
     public void tearDown() throws Exception {
         Settings.clearPreferences(sContext);
 
-        if (mRealDao != null) {
-            mRealDao.deleteAllKeys();
+        if (mKeyDao != null) {
+            mKeyDao.deleteAllKeys();
         }
 
         KeyStore keyStore = KeyStore.getInstance("AndroidKeyStore");
@@ -158,19 +151,19 @@
     }
 
     @Test
-    public void provisionThenExpireThenProvisionAgain() throws Exception {
+    public void provisionThenExpireThenProvisionAgain() {
         assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.success());
 
-        List<ProvisionedKey> keys = mTestDao.getAllKeys();
+        List<ProvisionedKey> keys = mKeyDao.getAllKeys();
 
         // Expire a key
         keys.get(0).expirationTime = Instant.now().minusSeconds(60);
-        mRealDao.updateKey(keys.get(0));
+        mKeyDao.updateKey(keys.get(0));
 
         // Mark two more keys as expiring soon
         for (int i = 1; i < 3; ++i) {
             keys.get(i).expirationTime = Instant.now().plusSeconds(60);
-            mRealDao.updateKey(keys.get(i));
+            mKeyDao.updateKey(keys.get(i));
         }
 
         assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.success());
@@ -182,14 +175,14 @@
         // key pool to ensure that the PeriodicProvisioner just noops.
         // This test is purely to test out proper metrics.
         assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.success());
-        StatsProcessor.PoolStats pool = StatsProcessor.processPool(mRealDao, mServiceName,
+        StatsProcessor.PoolStats pool = StatsProcessor.processPool(mKeyDao, mServiceName,
                 Settings.getExtraSignedKeysAvailable(sContext),
                 Settings.getExpirationTime(sContext));
 
         // The metrics host test will perform additional validation by ensuring correct metrics
         // are recorded.
         assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.success());
-        StatsProcessor.PoolStats updatedPool = StatsProcessor.processPool(mRealDao, mServiceName,
+        StatsProcessor.PoolStats updatedPool = StatsProcessor.processPool(mKeyDao, mServiceName,
                 Settings.getExtraSignedKeysAvailable(sContext),
                 Settings.getExpirationTime(sContext));
 
diff --git a/app/tests/unit/src/com/android/rkpdapp/unittest/CborUtilsTest.java b/app/tests/unit/src/com/android/rkpdapp/unittest/CborUtilsTest.java
index f99522f..ae1ffcc 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/CborUtilsTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/CborUtilsTest.java
@@ -35,6 +35,8 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -58,9 +60,12 @@
     private Array mGeekChain2;
     private byte[] mEncodedGeekChain2;
     private Map mDeviceConfig;
+    private Map mDeviceConfigWithBadCertInfo;
     private static final byte[] CHALLENGE = new byte[]{0x0a, 0x0b, 0x0c};
     private static final int TEST_EXTRA_KEYS = 18;
     private static final int TEST_TIME_TO_REFRESH_HOURS = 42;
+    private static final Instant BAD_CERT_START = Instant.now().minus(Duration.ofDays(2));
+    private static final Instant BAD_CERT_END = Instant.now().plus(Duration.ofDays(2));
     private static final String TEST_URL = "https://www.wonderifthisisvalid.combutjustincase";
 
     private byte[] encodeDataItem(DataItem toEncode) throws Exception {
@@ -95,6 +100,18 @@
                           new UnsignedInteger(TEST_TIME_TO_REFRESH_HOURS))
                      .put(new UnicodeString(CborUtils.PROVISIONING_URL),
                           new UnicodeString(TEST_URL));
+
+        mDeviceConfigWithBadCertInfo = new Map();
+        mDeviceConfigWithBadCertInfo.put(new UnicodeString(CborUtils.EXTRA_KEYS),
+                          new UnsignedInteger(TEST_EXTRA_KEYS))
+                     .put(new UnicodeString(CborUtils.TIME_TO_REFRESH),
+                          new UnsignedInteger(TEST_TIME_TO_REFRESH_HOURS))
+                     .put(new UnicodeString(CborUtils.PROVISIONING_URL),
+                          new UnicodeString(TEST_URL))
+                     .put(new UnicodeString(CborUtils.LAST_BAD_CERT_TIME_START_MILLIS),
+                          new UnsignedInteger(BAD_CERT_START.toEpochMilli()))
+                     .put(new UnicodeString(CborUtils.LAST_BAD_CERT_TIME_END_MILLIS),
+                          new UnsignedInteger(BAD_CERT_END.toEpochMilli()));
     }
 
     @Presubmit
@@ -184,6 +201,34 @@
         assertEquals(TEST_URL, resp.provisioningUrl);
     }
 
+    @Presubmit
+    @Test
+    public void testParseGeekResponseFakeDataWithBadCertTimeRange() throws Exception {
+        new CborEncoder(mBaos).encode(new CborBuilder()
+                .addArray()
+                    .addArray()                                       // GEEK Curve to Chains
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_25519))
+                            .add(mGeekChain1)
+                            .end()
+                        .addArray()
+                            .add(new UnsignedInteger(CborUtils.EC_CURVE_P256))
+                            .add(mGeekChain2)
+                            .end()
+                        .end()
+                    .add(CHALLENGE)
+                    .add(mDeviceConfigWithBadCertInfo)
+                    .end()
+                .build());
+        GeekResponse resp = CborUtils.parseGeekResponse(mBaos.toByteArray());
+        mBaos.reset();
+        assertEquals(TEST_EXTRA_KEYS, resp.numExtraAttestationKeys);
+        assertEquals(TEST_TIME_TO_REFRESH_HOURS, resp.timeToRefresh.toHours());
+        assertEquals(TEST_URL, resp.provisioningUrl);
+        assertEquals(BAD_CERT_START.toEpochMilli(), resp.lastBadCertTimeStart.toEpochMilli());
+        assertEquals(BAD_CERT_END.toEpochMilli(), resp.lastBadCertTimeEnd.toEpochMilli());
+    }
+
     @Test
     public void testExtraDeviceConfigEntriesDontFail() throws Exception {
         new CborEncoder(mBaos).encode(new CborBuilder()
@@ -238,6 +283,8 @@
         assertEquals(GeekResponse.NO_EXTRA_KEY_UPDATE, resp.numExtraAttestationKeys);
         assertNull(resp.timeToRefresh);
         assertNull(resp.provisioningUrl);
+        assertNull(resp.lastBadCertTimeStart);
+        assertNull(resp.lastBadCertTimeEnd);
     }
 
     @Test
@@ -266,6 +313,8 @@
         assertArrayEquals(CHALLENGE, resp.getChallenge());
         assertEquals(TEST_EXTRA_KEYS, resp.numExtraAttestationKeys);
         assertNull(resp.timeToRefresh);
+        assertNull(resp.lastBadCertTimeStart);
+        assertNull(resp.lastBadCertTimeEnd);
         assertEquals(TEST_URL, resp.provisioningUrl);
     }
 
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 0bccd28..b726aeb 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/PeriodicProvisionerTests.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/PeriodicProvisionerTests.java
@@ -53,8 +53,10 @@
 import org.junit.After;
 import org.junit.Assume;
 import org.junit.Before;
+import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mockito;
 
 import java.time.Duration;
 import java.time.Instant;
@@ -70,37 +72,41 @@
     private static final RkpKey FAKE_RKP_KEY = new RkpKey(new byte[1], new byte[2], new Array(),
             "fake-hal", new byte[3]);
 
+    private static Context sContext;
     private PeriodicProvisioner mProvisioner;
-    private Context mContext;
+
+    @BeforeClass
+    public static void init() {
+        sContext = Mockito.spy(ApplicationProvider.getApplicationContext());
+    }
 
     @Before
     public void setUp() {
-        mContext = ApplicationProvider.getApplicationContext();
-
         Assume.assumeFalse(Settings.getDefaultUrl().isEmpty());
 
-        RkpdDatabase.getDatabase(mContext).provisionedKeyDao().deleteAllKeys();
+        RkpdDatabase.getDatabase(sContext).provisionedKeyDao().deleteAllKeys();
         mProvisioner = TestWorkerBuilder.from(
-                mContext,
+                sContext,
                 PeriodicProvisioner.class,
                 Executors.newSingleThreadExecutor()).build();
 
         Configuration config = new Configuration.Builder()
                 .setExecutor(new SynchronousExecutor())
                 .build();
-        WorkManagerTestInitHelper.initializeTestWorkManager(mContext, config);
-        Settings.clearPreferences(mContext);
+        WorkManagerTestInitHelper.initializeTestWorkManager(sContext, config);
+        Settings.clearPreferences(sContext);
+        Utils.mockConnectivityState(sContext, Utils.ConnectivityState.CONNECTED);
     }
 
     @After
     public void tearDown() {
-        RkpdDatabase.getDatabase(mContext).provisionedKeyDao().deleteAllKeys();
+        RkpdDatabase.getDatabase(sContext).provisionedKeyDao().deleteAllKeys();
         ServiceManagerInterface.setInstances(null);
-        Settings.clearPreferences(mContext);
+        Settings.clearPreferences(sContext);
     }
 
     private WorkInfo getProvisionerWorkInfo() throws ExecutionException, InterruptedException {
-        WorkManager workManager = WorkManager.getInstance(mContext);
+        WorkManager workManager = WorkManager.getInstance(sContext);
         List<WorkInfo> infos = workManager.getWorkInfosForUniqueWork(
                 PeriodicProvisioner.UNIQUE_WORK_NAME).get();
         assertThat(infos.size()).isEqualTo(1);
@@ -110,20 +116,20 @@
     @Test
     public void provisionWithNoHals() throws Exception {
         // setup work with boot receiver
-        new BootReceiver().onReceive(mContext, null);
+        new BootReceiver().onReceive(sContext, null);
 
         WorkInfo worker = getProvisionerWorkInfo();
         assertThat(worker.getState()).isEqualTo(WorkInfo.State.ENQUEUED);
 
         ServiceManagerInterface.setInstances(new SystemInterface[0]);
-        WorkManagerTestInitHelper.getTestDriver(mContext).setAllConstraintsMet(worker.getId());
+        WorkManagerTestInitHelper.getTestDriver(sContext).setAllConstraintsMet(worker.getId());
 
         // the worker should uninstall itself once it realizes it's not needed on this system
         worker = getProvisionerWorkInfo();
         assertThat(worker.getState()).isEqualTo(WorkInfo.State.CANCELLED);
 
         // verify the worker doesn't run again
-        WorkManagerTestInitHelper.getTestDriver(mContext).setAllConstraintsMet(worker.getId());
+        WorkManagerTestInitHelper.getTestDriver(sContext).setAllConstraintsMet(worker.getId());
         worker = getProvisionerWorkInfo();
         assertThat(worker.getState()).isEqualTo(WorkInfo.State.CANCELLED);
     }
@@ -131,14 +137,14 @@
     @Test
     public void provisionWithNoHostNameWithoutServerUrl() throws Exception {
         // setup work with boot receiver
-        new BootReceiver().onReceive(mContext, null);
+        new BootReceiver().onReceive(sContext, null);
 
         try (SystemPropertySetter ignored = SystemPropertySetter.setHostname("")) {
             SystemInterface mockHal = mock(SystemInterface.class);
             ServiceManagerInterface.setInstances(new SystemInterface[]{mockHal});
 
             WorkInfo worker = getProvisionerWorkInfo();
-            WorkManagerTestInitHelper.getTestDriver(mContext).setAllConstraintsMet(worker.getId());
+            WorkManagerTestInitHelper.getTestDriver(sContext).setAllConstraintsMet(worker.getId());
         }
 
         WorkInfo worker = getProvisionerWorkInfo();
@@ -148,16 +154,16 @@
     @Test
     public void provisionWithNoHostNameWithServerUrl() throws Exception {
         // setup work with boot receiver
-        new BootReceiver().onReceive(mContext, null);
+        new BootReceiver().onReceive(sContext, null);
 
         try (SystemPropertySetter ignored = SystemPropertySetter.setHostname("")) {
             SystemInterface mockHal = mock(SystemInterface.class);
             ServiceManagerInterface.setInstances(new SystemInterface[]{mockHal});
-            Settings.setDeviceConfig(mContext, Settings.EXTRA_SIGNED_KEYS_AVAILABLE_DEFAULT,
+            Settings.setDeviceConfig(sContext, Settings.EXTRA_SIGNED_KEYS_AVAILABLE_DEFAULT,
                     Duration.ofDays(3), "https://notsure.whetherthisworks.combutjustincase");
 
             WorkInfo worker = getProvisionerWorkInfo();
-            WorkManagerTestInitHelper.getTestDriver(mContext).setAllConstraintsMet(worker.getId());
+            WorkManagerTestInitHelper.getTestDriver(sContext).setAllConstraintsMet(worker.getId());
         }
 
         WorkInfo worker = getProvisionerWorkInfo();
@@ -186,7 +192,7 @@
                 FakeRkpServer.Response.INTERNAL_ERROR,
                 FakeRkpServer.Response.SIGN_CERTS_OK_VALID_CBOR)) {
             saveUrlInSettings(fakeRkpServer);
-            Settings.setMaxRequestTime(mContext, 100);
+            Settings.setMaxRequestTime(sContext, 100);
             SystemInterface mockHal = mock(SystemInterface.class);
             ServiceManagerInterface.setInstances(new SystemInterface[]{mockHal});
             assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.failure());
@@ -198,7 +204,7 @@
 
     @Test
     public void fetchEekDisablesRkp() throws Exception {
-        ProvisionedKeyDao dao = RkpdDatabase.getDatabase(mContext).provisionedKeyDao();
+        ProvisionedKeyDao dao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
         ProvisionedKey fakeKey = new ProvisionedKey(new byte[42], "fake-irpc", new byte[3],
                 new byte[2], Instant.now().plusSeconds(120));
         dao.insertKeys(List.of(fakeKey));
@@ -222,7 +228,7 @@
 
     @Test
     public void provisioningExpiresOldKeys() throws Exception {
-        ProvisionedKeyDao dao = RkpdDatabase.getDatabase(mContext).provisionedKeyDao();
+        ProvisionedKeyDao dao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
         ProvisionedKey oldKey = new ProvisionedKey(new byte[1], "fake-irpc", new byte[2],
                 new byte[3],
                 Instant.now().minus(RegistrationBinder.MIN_KEY_LIFETIME.multipliedBy(2)));
@@ -324,6 +330,6 @@
     }
 
     private void saveUrlInSettings(FakeRkpServer server) {
-        Settings.setDeviceConfig(mContext, 1, Duration.ofSeconds(10), server.getUrl());
+        Settings.setDeviceConfig(sContext, 1, Duration.ofSeconds(10), server.getUrl());
     }
 }
diff --git a/app/tests/unit/src/com/android/rkpdapp/unittest/ProvisionerTest.java b/app/tests/unit/src/com/android/rkpdapp/unittest/ProvisionerTest.java
index 384b172..4acab0d 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/ProvisionerTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/ProvisionerTest.java
@@ -16,6 +16,8 @@
 
 package com.android.rkpdapp.unittest;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.notNull;
@@ -34,6 +36,7 @@
 
 import com.android.rkpdapp.GeekResponse;
 import com.android.rkpdapp.RkpdException;
+import com.android.rkpdapp.database.ProvisionedKey;
 import com.android.rkpdapp.database.ProvisionedKeyDao;
 import com.android.rkpdapp.database.RkpKey;
 import com.android.rkpdapp.database.RkpdDatabase;
@@ -43,23 +46,35 @@
 import com.android.rkpdapp.testutil.FakeRkpServer;
 import com.android.rkpdapp.utils.Settings;
 
+import com.google.crypto.tink.subtle.Random;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.security.KeyPair;
 import java.time.Duration;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.List;
 
 import co.nstant.in.cbor.model.Array;
 
 @RunWith(AndroidJUnit4.class)
 public class ProvisionerTest {
-    private static final RkpKey FAKE_RKP_KEY = new RkpKey(new byte[1], new byte[2], new Array(),
-            "hal", new byte[3]);
+    private static final byte[] FAKE_RKP_KEY_BLOB_1 = Random.randBytes(10);
+    private static final byte[] FAKE_RKP_KEY_BLOB_2 = Random.randBytes(10);
+    private static final byte[] FAKE_RKP_KEY_BLOB_3 = Random.randBytes(10);
+    private static final Instant NOW = Instant.now().truncatedTo(ChronoUnit.SECONDS);
+
+    private static final RkpKey FAKE_RKP_KEY = new RkpKey(FAKE_RKP_KEY_BLOB_1, new byte[2],
+            new Array(), "hal", new byte[3]);
 
     private static Context sContext;
     private Provisioner mProvisioner;
+    private ProvisionedKeyDao mKeyDao;
 
     @BeforeClass
     public static void init() {
@@ -70,10 +85,10 @@
     public void setUp() {
         Settings.clearPreferences(sContext);
 
-        ProvisionedKeyDao keyDao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
-        keyDao.deleteAllKeys();
+        mKeyDao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
+        mKeyDao.deleteAllKeys();
 
-        mProvisioner = new Provisioner(sContext, keyDao, false);
+        mProvisioner = new Provisioner(sContext, mKeyDao, false);
     }
 
     @After
@@ -121,4 +136,86 @@
                     mProvisioner.provisionKeys(atom, mockSystem, geekResponse));
         }
     }
+
+    private byte[] generateCertificateChain(Instant rootCreationTime, Instant leafCreationTime)
+            throws Exception {
+        KeyPair rootKey = Utils.generateEcdsaKeyPair();
+        KeyPair leafKey = Utils.generateEcdsaKeyPair();
+        // Just so that we don't get expired certificates by default.
+        Instant expirationTime = NOW.plus(Duration.ofDays(1));
+        byte[] rootCertEncoded = Utils.signPublicKey(rootKey, rootKey.getPublic(), rootCreationTime,
+                expirationTime).getEncoded();
+        byte[] leafCertEncoded = Utils.signPublicKey(rootKey, leafKey.getPublic(), leafCreationTime,
+                expirationTime).getEncoded();
+
+        byte[] encodedCertChain = new byte[leafCertEncoded.length + rootCertEncoded.length];
+        System.arraycopy(leafCertEncoded, 0, encodedCertChain, 0, leafCertEncoded.length);
+        System.arraycopy(rootCertEncoded, 0, encodedCertChain, leafCertEncoded.length,
+                rootCertEncoded.length);
+        return encodedCertChain;
+    }
+
+    private void setUpClearAttestationKeyTests(Instant failureStart, Instant failureEnd)
+            throws Exception {
+        Instant expiration = NOW.plus(Duration.ofDays(1));
+        Instant rootCreationTime = failureStart.minus(Duration.ofDays(10));
+
+        // add a fake key to the database with certificate time that is in the bad cert range.
+        ProvisionedKey keyBeforeFailure = new ProvisionedKey(
+                FAKE_RKP_KEY_BLOB_1,
+                "fakeHal1",
+                new byte[0],
+                generateCertificateChain(rootCreationTime, failureStart.minus(Duration.ofDays(1))),
+                expiration);
+        ProvisionedKey keyBadCert = new ProvisionedKey(
+                FAKE_RKP_KEY_BLOB_2,
+                "fakeHal2",
+                new byte[0],
+                generateCertificateChain(rootCreationTime, failureStart.plus(Duration.ofHours(1))),
+                expiration);
+        ProvisionedKey keyAfterFailure = new ProvisionedKey(
+                FAKE_RKP_KEY_BLOB_3,
+                "fakeHal3",
+                new byte[0],
+                generateCertificateChain(rootCreationTime, failureEnd.plus(Duration.ofDays(1))),
+                expiration);
+        mKeyDao.insertKeys(List.of(keyBeforeFailure, keyBadCert, keyAfterFailure));
+    }
+
+    @Test
+    public void testProvisionerClearsAttestationKeysOnResponse() throws Exception {
+        Instant failureTimeStart = NOW.minus(Duration.ofDays(5));
+        Instant failureTimeEnd = NOW.minus(Duration.ofDays(2));
+
+        setUpClearAttestationKeyTests(failureTimeStart, failureTimeEnd);
+
+        assertThat(mKeyDao.getAllKeys()).hasSize(3);
+
+        GeekResponse resp = new GeekResponse();
+        resp.lastBadCertTimeStart = failureTimeStart;
+        resp.lastBadCertTimeEnd = failureTimeEnd;
+
+        mProvisioner.clearBadAttestationKeys(resp);
+
+        assertThat(mKeyDao.getAllKeys()).hasSize(2);
+    }
+
+    @Test
+    public void testProvisionerClearsAttestationKeysOnlyOnce() throws Exception {
+        Instant failureTimeStart = NOW.minus(Duration.ofDays(5));
+        Instant failureTimeEnd = NOW.minus(Duration.ofDays(2));
+
+        setUpClearAttestationKeyTests(failureTimeStart, failureTimeEnd);
+
+        assertThat(mKeyDao.getAllKeys()).hasSize(3);
+
+        GeekResponse resp = new GeekResponse();
+        resp.lastBadCertTimeStart = failureTimeStart;
+        resp.lastBadCertTimeEnd = failureTimeEnd;
+        Settings.setLastBadCertTimeRange(sContext, failureTimeStart, failureTimeEnd);
+
+        mProvisioner.clearBadAttestationKeys(resp);
+
+        assertThat(mKeyDao.getAllKeys()).hasSize(3);
+    }
 }
diff --git a/app/tests/unit/src/com/android/rkpdapp/unittest/RkpdDatabaseTest.java b/app/tests/unit/src/com/android/rkpdapp/unittest/RkpdDatabaseTest.java
index e9b4a9e..a21deb5 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/RkpdDatabaseTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/RkpdDatabaseTest.java
@@ -32,8 +32,6 @@
 import com.android.rkpdapp.database.ProvisionedKey;
 import com.android.rkpdapp.database.ProvisionedKeyDao;
 import com.android.rkpdapp.database.RkpdDatabase;
-import com.android.rkpdapp.testutil.TestDatabase;
-import com.android.rkpdapp.testutil.TestProvisionedKeyDao;
 
 import org.junit.After;
 import org.junit.Before;
@@ -64,8 +62,6 @@
 
     private ProvisionedKeyDao mKeyDao;
     private RkpdDatabase mDatabase;
-    private TestDatabase mTestDatabase;
-    private TestProvisionedKeyDao mTestDao;
 
     @Before
     public void setUp() {
@@ -73,8 +69,6 @@
         mDatabase = Room.databaseBuilder(context, RkpdDatabase.class, DB_NAME).build();
         mKeyDao = mDatabase.provisionedKeyDao();
         mKeyDao.deleteAllKeys();
-        mTestDatabase = Room.databaseBuilder(context, TestDatabase.class, DB_NAME).build();
-        mTestDao = mTestDatabase.dao();
         mProvisionedKey1 = new ProvisionedKey(TEST_KEY_BLOB_1, TEST_HAL_1, TEST_KEY_BLOB_1,
                 TEST_KEY_BLOB_1, TEST_KEY_EXPIRY);
         mProvisionedKey2 = new ProvisionedKey(TEST_KEY_BLOB_2, TEST_HAL_2, TEST_KEY_BLOB_2,
@@ -84,13 +78,12 @@
     @After
     public void tearDown() {
         mDatabase.close();
-        mTestDatabase.close();
     }
 
     @Test
     public void testWriteToTable() {
         mKeyDao.insertKeys(List.of(mProvisionedKey1));
-        List<ProvisionedKey> keysInDatabase = mTestDao.getAllKeys();
+        List<ProvisionedKey> keysInDatabase = mKeyDao.getAllKeys();
 
         assertThat(keysInDatabase).containsExactly(mProvisionedKey1);
     }
@@ -105,7 +98,7 @@
             assertThat(ex).hasMessageThat().contains("UNIQUE constraint failed");
         }
 
-        List<ProvisionedKey> unassignedKeys = mTestDao.getAllKeys();
+        List<ProvisionedKey> unassignedKeys = mKeyDao.getAllKeys();
         assertThat(unassignedKeys).isEmpty();
     }
 
@@ -116,12 +109,12 @@
 
         mKeyDao.insertKeys(List.of(mProvisionedKey1, mProvisionedKey2));
 
-        List<ProvisionedKey> keysInDatabase = mTestDao.getAllKeys();
+        List<ProvisionedKey> keysInDatabase = mKeyDao.getAllKeys();
         assertThat(keysInDatabase).hasSize(2);
 
         mKeyDao.deleteExpiringKeys(Instant.now());
 
-        keysInDatabase = mTestDao.getAllKeys();
+        keysInDatabase = mKeyDao.getAllKeys();
         assertThat(keysInDatabase).containsExactly(mProvisionedKey2);
     }
 
@@ -144,7 +137,7 @@
     public void testUpdate() {
         mKeyDao.insertKeys(List.of(mProvisionedKey1));
 
-        List<ProvisionedKey> keysInDatabase = mTestDao.getAllKeys();
+        List<ProvisionedKey> keysInDatabase = mKeyDao.getAllKeys();
         ProvisionedKey key = keysInDatabase.get(0);
         assertThat(keysInDatabase).hasSize(1);
         assertThat(key.expirationTime).isEqualTo(
@@ -154,7 +147,7 @@
                 .minus(1000, ChronoUnit.MINUTES);
         key.expirationTime = expiredInstant;
         mKeyDao.updateKey(key);
-        keysInDatabase = mTestDao.getAllKeys();
+        keysInDatabase = mKeyDao.getAllKeys();
         assertThat(keysInDatabase).containsExactly(key);
         assertThat(keysInDatabase.get(0).expirationTime).isEqualTo(expiredInstant);
     }
@@ -163,18 +156,32 @@
     public void testUpdateWithNonExistentKey() {
         mKeyDao.updateKey(mProvisionedKey1);
 
-        assertThat(mTestDao.getAllKeys()).isEmpty();
+        assertThat(mKeyDao.getAllKeys()).isEmpty();
     }
 
     @Test
     public void testDeleteAllKeys() {
         mKeyDao.insertKeys(List.of(mProvisionedKey1, mProvisionedKey2));
 
-        List<ProvisionedKey> keysInDatabase = mTestDao.getAllKeys();
+        List<ProvisionedKey> keysInDatabase = mKeyDao.getAllKeys();
         assertThat(keysInDatabase).hasSize(2);
 
         mKeyDao.deleteAllKeys();
-        assertThat(mTestDao.getAllKeys()).isEmpty();
+        assertThat(mKeyDao.getAllKeys()).isEmpty();
+    }
+
+    @Test
+    public void testDeleteSingleKey() {
+        mKeyDao.insertKeys(List.of(mProvisionedKey1, mProvisionedKey2));
+        List<ProvisionedKey> keysInDatabase = mKeyDao.getAllKeys();
+        assertThat(keysInDatabase).hasSize(2);
+
+        mKeyDao.deleteKey(mProvisionedKey1.keyBlob);
+        keysInDatabase = mKeyDao.getAllKeys();
+        assertThat(keysInDatabase).hasSize(1);
+
+        ProvisionedKey key = keysInDatabase.get(0);
+        assertThat(key.keyBlob).isEqualTo(mProvisionedKey2.keyBlob);
     }
 
     @Test
@@ -234,14 +241,14 @@
         mProvisionedKey1.clientUid = FAKE_CLIENT_UID;
         mKeyDao.insertKeys(List.of(mProvisionedKey1));
 
-        ProvisionedKey databaseKey = mTestDao.getAllKeys().get(0);
+        ProvisionedKey databaseKey = mKeyDao.getAllKeys().get(0);
         assertThat(databaseKey.keyBlob).isEqualTo(TEST_KEY_BLOB_1);
         assertThat(mKeyDao.upgradeKeyBlob(FAKE_CLIENT_UID_2, TEST_KEY_BLOB_1, TEST_KEY_BLOB_2))
                 .isEqualTo(0);
         assertThat(mKeyDao.upgradeKeyBlob(FAKE_CLIENT_UID, TEST_KEY_BLOB_1, TEST_KEY_BLOB_2))
                 .isEqualTo(1);
 
-        databaseKey = mTestDao.getAllKeys().get(0);
+        databaseKey = mKeyDao.getAllKeys().get(0);
         assertThat(databaseKey.keyBlob).isEqualTo(TEST_KEY_BLOB_2);
     }
 
@@ -251,12 +258,12 @@
         mProvisionedKey1.clientUid = FAKE_CLIENT_UID;
         mKeyDao.insertKeys(List.of(mProvisionedKey1));
 
-        ProvisionedKey databaseKey = mTestDao.getAllKeys().get(0);
+        ProvisionedKey databaseKey = mKeyDao.getAllKeys().get(0);
         assertThat(databaseKey.keyBlob).isEqualTo(TEST_KEY_BLOB_1);
         assertThat(mKeyDao.upgradeKeyBlob(FAKE_CLIENT_UID_2, TEST_KEY_BLOB_1, TEST_KEY_BLOB_2))
                 .isEqualTo(0);
 
-        databaseKey = mTestDao.getAllKeys().get(0);
+        databaseKey = mKeyDao.getAllKeys().get(0);
         assertThat(databaseKey.keyBlob).isEqualTo(TEST_KEY_BLOB_1);
     }
 
@@ -282,7 +289,7 @@
         mProvisionedKey2.irpcHal = TEST_HAL_1;
         mKeyDao.insertKeys(List.of(mProvisionedKey1, mProvisionedKey2));
 
-        List<ProvisionedKey> keysPersisted = mTestDao.getAllKeys();
+        List<ProvisionedKey> keysPersisted = mKeyDao.getAllKeys();
         for (ProvisionedKey databaseKey : keysPersisted) {
             assertThat(databaseKey.keyId).isNull();
             assertThat(databaseKey.clientUid).isNull();
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 ed1f64a..12ab92d 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/ServerInterfaceTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/ServerInterfaceTest.java
@@ -22,8 +22,6 @@
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
-import android.net.ConnectivityManager;
-import android.net.NetworkCapabilities;
 import android.util.Base64;
 
 import androidx.test.core.app.ApplicationProvider;
@@ -64,6 +62,7 @@
     public void setUp() {
         Settings.clearPreferences(sContext);
         mServerInterface = new ServerInterface(sContext, false);
+        Utils.mockConnectivityState(sContext, Utils.ConnectivityState.CONNECTED);
     }
 
     @After
@@ -243,9 +242,6 @@
             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);
             mServerInterface.fetchGeek(metrics);
             assertWithMessage("Network transaction should not have proceeded.").fail();
         } catch (RkpdException e) {
@@ -268,7 +264,7 @@
 
             // We are okay in mocking connectivity failure since network check is the first thing
             // to happen.
-            mockConnectivityFailure(ConnectivityState.DISCONNECTED);
+            Utils.mockConnectivityState(sContext, Utils.ConnectivityState.DISCONNECTED);
             mServerInterface.fetchGeek(metrics);
             assertWithMessage("Network transaction should not have proceeded.").fail();
         } catch (RkpdException e) {
@@ -440,24 +436,4 @@
         fakeApplicationInfo.enabled = false;
         assertThat(ServerInterface.assumeNetworkConsent(mockedContext)).isTrue();
     }
-
-    private void mockConnectivityFailure(ConnectivityState state) {
-        ConnectivityManager mockedConnectivityManager = Mockito.mock(ConnectivityManager.class);
-
-        Mockito.when(sContext.getSystemService(ConnectivityManager.class))
-                .thenReturn(mockedConnectivityManager);
-        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 {
-        DISCONNECTED,
-        CONNECTED
-    }
-
 }
diff --git a/app/tests/unit/src/com/android/rkpdapp/unittest/SettingsTest.java b/app/tests/unit/src/com/android/rkpdapp/unittest/SettingsTest.java
index 3c1029f..ff8383a 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/SettingsTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/SettingsTest.java
@@ -20,6 +20,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 import android.content.Context;
@@ -38,6 +39,7 @@
 
 import java.time.Duration;
 import java.time.Instant;
+import java.time.temporal.ChronoUnit;
 
 @RunWith(AndroidJUnit4.class)
 public class SettingsTest {
@@ -214,4 +216,14 @@
         Settings.setMaxRequestTime(sContext, 100);
         assertEquals(100, Settings.getMaxRequestTime(sContext));
     }
+
+    @Test
+    public void testLastBadCertTimeRangeSetting() {
+        assertNull(Settings.getLastBadCertTimeStart(sContext));
+        assertNull(Settings.getLastBadCertTimeEnd(sContext));
+        Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS);
+        Settings.setLastBadCertTimeRange(sContext, now, now);
+        assertEquals(now, Settings.getLastBadCertTimeStart(sContext));
+        assertEquals(now, Settings.getLastBadCertTimeEnd(sContext));
+    }
 }
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 bbce2d9..7ef3a78 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/Utils.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/Utils.java
@@ -20,6 +20,10 @@
 import static com.google.crypto.tink.subtle.EllipticCurves.EcdsaEncoding.IEEE_P1363;
 import static com.google.crypto.tink.subtle.Enums.HashType.SHA256;
 
+import android.content.Context;
+import android.net.ConnectivityManager;
+import android.net.NetworkCapabilities;
+
 import com.google.crypto.tink.subtle.EcdsaSignJce;
 import com.google.crypto.tink.subtle.Ed25519Sign;
 import com.google.crypto.tink.subtle.EllipticCurves;
@@ -28,6 +32,7 @@
 import org.bouncycastle.asn1.x509.Extension;
 import org.bouncycastle.asn1.x509.KeyUsage;
 import org.bouncycastle.x509.X509V3CertificateGenerator;
+import org.mockito.Mockito;
 
 import java.io.ByteArrayOutputStream;
 import java.math.BigInteger;
@@ -150,15 +155,23 @@
 
     public static X509Certificate signPublicKey(KeyPair issuerKeyPair, PublicKey publicKeyToSign,
             Instant expirationInstant) throws Exception {
+        Instant now = Instant.now();
+        return signPublicKey(issuerKeyPair, publicKeyToSign, now, expirationInstant);
+    }
+
+    /**
+     * Generates a certificate for given key and issuer.
+     */
+    public static X509Certificate signPublicKey(KeyPair issuerKeyPair, PublicKey publicKeyToSign,
+            Instant creationInstant, Instant expirationInstant) throws Exception {
         X500Principal issuer = new X500Principal("CN=TEE");
         BigInteger serial = BigInteger.ONE;
         X500Principal subject = new X500Principal("CN=TEE");
 
-        Instant now = Instant.now();
         X509V3CertificateGenerator certificateBuilder = new X509V3CertificateGenerator();
         certificateBuilder.setIssuerDN(issuer);
         certificateBuilder.setSerialNumber(serial);
-        certificateBuilder.setNotBefore(Date.from(now));
+        certificateBuilder.setNotBefore(Date.from(creationInstant));
         certificateBuilder.setNotAfter(Date.from(expirationInstant));
         certificateBuilder.setSignatureAlgorithm("SHA256WITHECDSA");
         certificateBuilder.setSubjectDN(subject);
@@ -292,4 +305,26 @@
                 .build());
         return baos.toByteArray();
     }
+
+    /**
+     * Mocks out the connectivity status for unit tests.
+     */
+    public static void mockConnectivityState(Context context, ConnectivityState state) {
+        ConnectivityManager mockedConnectivityManager = Mockito.mock(ConnectivityManager.class);
+
+        Mockito.when(context.getSystemService(ConnectivityManager.class))
+                .thenReturn(mockedConnectivityManager);
+        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());
+    }
+
+    public enum ConnectivityState {
+        DISCONNECTED,
+        CONNECTED
+    }
 }
diff --git a/app/tests/util/src/com/android/rkpdapp/testutil/FakeRkpServer.java b/app/tests/util/src/com/android/rkpdapp/testutil/FakeRkpServer.java
index 3023619..6cf0d40 100644
--- a/app/tests/util/src/com/android/rkpdapp/testutil/FakeRkpServer.java
+++ b/app/tests/util/src/com/android/rkpdapp/testutil/FakeRkpServer.java
@@ -70,7 +70,6 @@
                     + "8L01k/PGu1lOXvneIQcUo7ako4uPgpaWugNYHQAAAYBINcxrASC0rWP9VTSO7LdABvcdkv7W2vh+"
                     + "onV0aW1lX3RvX3JlZnJlc2hfaG91cnMYSHgabnVtX2V4dHJhX2F0dGVzdGF0aW9uX2tleXMA";
 
-
     public enum Response {
         // canned responses for :fetchEekChain
         FETCH_EEK_OK(EEK_RESPONSE_OK),
diff --git a/app/tests/util/src/com/android/rkpdapp/testutil/TestDatabase.java b/app/tests/util/src/com/android/rkpdapp/testutil/TestDatabase.java
deleted file mode 100644
index f6b0eb6..0000000
--- a/app/tests/util/src/com/android/rkpdapp/testutil/TestDatabase.java
+++ /dev/null
@@ -1,30 +0,0 @@
-/*
- * 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.rkpdapp.testutil;
-
-import androidx.room.Database;
-import androidx.room.RoomDatabase;
-import androidx.room.TypeConverters;
-
-import com.android.rkpdapp.database.InstantConverter;
-import com.android.rkpdapp.database.ProvisionedKey;
-
-@Database(entities = {ProvisionedKey.class}, exportSchema = false, version = 1)
-@TypeConverters({InstantConverter.class})
-public abstract class TestDatabase extends RoomDatabase {
-    public abstract TestProvisionedKeyDao dao();
-}
diff --git a/app/tests/util/src/com/android/rkpdapp/testutil/TestProvisionedKeyDao.java b/app/tests/util/src/com/android/rkpdapp/testutil/TestProvisionedKeyDao.java
deleted file mode 100644
index cee1c53..0000000
--- a/app/tests/util/src/com/android/rkpdapp/testutil/TestProvisionedKeyDao.java
+++ /dev/null
@@ -1,30 +0,0 @@
-/*
- * 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.rkpdapp.testutil;
-
-import androidx.room.Dao;
-import androidx.room.Query;
-
-import com.android.rkpdapp.database.ProvisionedKey;
-
-import java.util.List;
-
-@Dao
-public abstract class TestProvisionedKeyDao {
-    @Query("SELECT * FROM provisioned_keys")
-    public abstract List<ProvisionedKey> getAllKeys();
-}