Snap for 11173240 from b7a5d0973ffc4bb3feb8cd6896199f5aa93ac7cd to mainline-sdkext-releaseaml_sdk_341410000android14-mainline-sdkext-release

Change-Id: Ic41d2cc25b38c11c1440e9054719de9639aa1482

22 files changed, 186 insertions, 45 deletions

diff --git a/OWNERS b/OWNERS
index 01955db..2cbf527 100644
--- a/OWNERS
+++ b/OWNERS
@@ -1,3 +1,4 @@
+# Bug component: 1084908
 sethmo@google.com
 vikramgaur@google.com
 
diff --git a/app/TEST_MAPPING b/app/TEST_MAPPING
index cb4d40b..824fc9e 100644
--- a/app/TEST_MAPPING
+++ b/app/TEST_MAPPING
@@ -4,13 +4,15 @@
       "name": "RkpdAppUnitTests"
     },
     {
-      "name": "RkpdAppGoogleUnitTests"
+      "name": "RkpdAppGoogleUnitTests",
+      "keywords": ["internal"]
     },
     {
       "name": "RkpdAppIntegrationTests"
     },
     {
-      "name": "RkpdAppGoogleIntegrationTests"
+      "name": "RkpdAppGoogleIntegrationTests",
+      "keywords": ["internal"]
     }
   ],
   "mainline-presubmit": [
diff --git a/app/src/com/android/rkpdapp/interfaces/ServerInterface.java b/app/src/com/android/rkpdapp/interfaces/ServerInterface.java
index 173a2a8..8ce4917 100644
--- a/app/src/com/android/rkpdapp/interfaces/ServerInterface.java
+++ b/app/src/com/android/rkpdapp/interfaces/ServerInterface.java
@@ -55,9 +55,9 @@ import java.util.UUID;
  */
 public class ServerInterface {
 
+    private static final int SYNC_CONNECT_TIMEOUT_MS = 1000;
     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";
@@ -65,6 +65,7 @@ public class ServerInterface {
     private static final String CHALLENGE_PARAMETER = "challenge";
     private static final String REQUEST_ID_PARAMETER = "request_id";
     private final Context mContext;
+    private final boolean mIsAsync;
 
     private enum Operation {
         FETCH_GEEK(1),
@@ -108,8 +109,13 @@ public class ServerInterface {
         }
     }
 
-    public ServerInterface(Context context) {
+    public ServerInterface(Context context, boolean isAsync) {
         this.mContext = context;
+        this.mIsAsync = isAsync;
+    }
+
+    private int getConnectTimeoutMs() {
+        return mIsAsync ? TIMEOUT_MS : SYNC_CONNECT_TIMEOUT_MS;
     }
 
     /**
@@ -179,7 +185,7 @@ public class ServerInterface {
     /**
      * Calls out to the specified backend servers to retrieve an Endpoint Encryption Key and
      * corresponding certificate chain to provide to KeyMint. This public key will be used to
-     * perform an ECDH computation, using the shared secret to encrypt privacy sensitive components
+     * perform an ECDH computation, using the shared secret to encrypt privacy-sensitive components
      * in the bundle that the server needs from the device in order to provision certificates.
      *
      * A challenge is also returned from the server so that it can check freshness of the follow-up
@@ -358,7 +364,7 @@ public class ServerInterface {
                 }
                 // Only RkpdExceptions should get longer retries.
                 if (retryTimer.getElapsedMillis() > Settings.getMaxRequestTime(mContext)
-                        || (attempt >= SHORT_RETRY_COUNT && lastSeenRkpdException == null)) {
+                        || lastSeenRkpdException == null) {
                     break;
                 }
                 Thread.sleep(backoff_time);
@@ -381,7 +387,7 @@ public class ServerInterface {
         try (StopWatch serverWaitTimer = metrics.startServerWait()) {
             HttpURLConnection con = (HttpURLConnection) url.openConnection();
             con.setRequestMethod("POST");
-            con.setConnectTimeout(TIMEOUT_MS);
+            con.setConnectTimeout(getConnectTimeoutMs());
             con.setReadTimeout(TIMEOUT_MS);
             con.setDoOutput(true);
 
diff --git a/app/src/com/android/rkpdapp/interfaces/ServiceManagerInterface.java b/app/src/com/android/rkpdapp/interfaces/ServiceManagerInterface.java
index 38766b1..3ba716c 100644
--- a/app/src/com/android/rkpdapp/interfaces/ServiceManagerInterface.java
+++ b/app/src/com/android/rkpdapp/interfaces/ServiceManagerInterface.java
@@ -19,8 +19,11 @@ package com.android.rkpdapp.interfaces;
 import android.annotation.TestApi;
 import android.hardware.security.keymint.IRemotelyProvisionedComponent;
 import android.os.ServiceManager;
+import android.util.Log;
 
 import java.util.Arrays;
+import java.util.Map;
+import java.util.Objects;
 
 /**
  * Provides convenience methods for interfacing with ServiceManager class and its static functions.
@@ -28,33 +31,54 @@ import java.util.Arrays;
 public class ServiceManagerInterface {
     private static final String TAG = "RkpdSvcManagerInterface";
     private static SystemInterface[] sInstances;
+    private static Map<String, IRemotelyProvisionedComponent> sBinders;
 
     private ServiceManagerInterface() {
     }
 
-    private static SystemInterface createSystemInterface(String serviceName) {
+    private static SystemInterface tryCreateSystemInterface(IRemotelyProvisionedComponent binder,
+            String serviceName) {
+        try {
+            return new SystemInterface(binder, serviceName);
+        } catch (UnsupportedOperationException e) {
+            Log.i(TAG, serviceName + " is unsupported.");
+            return null;
+        }
+    }
+
+    private static IRemotelyProvisionedComponent getBinder(String serviceName) {
         IRemotelyProvisionedComponent binder = IRemotelyProvisionedComponent.Stub.asInterface(
                 ServiceManager.waitForDeclaredService(serviceName));
         if (binder == null) {
             throw new IllegalArgumentException("Cannot find any implementation for " + serviceName);
         }
-        return new SystemInterface(binder, serviceName);
+        return binder;
     }
 
     /**
      * Gets all the instances on this device for IRemotelyProvisionedComponent as an array. The
      * returned values each contain a binder for interacting with the instance.
      *
-     * For testing purposes, the instances may be overridden by setInstances
+     * For testing purposes, the instances may be overridden by either setInstances or setBinders.
      */
     public static SystemInterface[] getAllInstances() {
+        if (sBinders != null) {
+            return sBinders.entrySet().stream()
+                    .map(x -> tryCreateSystemInterface(x.getValue(), x.getKey()))
+                    .filter(Objects::nonNull)
+                    .toArray(SystemInterface[]::new);
+        }
         if (sInstances != null) {
             return sInstances;
         }
 
         String irpcInterface = IRemotelyProvisionedComponent.DESCRIPTOR;
         return Arrays.stream(ServiceManager.getDeclaredInstances(irpcInterface))
-                .map(x -> createSystemInterface(irpcInterface + "/" + x))
+                .map(x -> {
+                    String serviceName = irpcInterface + "/" + x;
+                    return tryCreateSystemInterface(getBinder(serviceName), serviceName);
+                })
+                .filter(Objects::nonNull)
                 .toArray(SystemInterface[]::new);
     }
 
@@ -62,10 +86,17 @@ public class ServiceManagerInterface {
      * Get a specific system interface instance for a given IRemotelyProvisionedComponent.
      * If the given serviceName does not map to a known IRemotelyProvisionedComponent, this
      * method throws IllegalArgumentException.
+     * If the given serviceName is not supported, this method throws UnsupportedOperationException.
      *
-     * For testing purposes, the instances may be overridden by setInstances.
+     * For testing purposes, the instances may be overridden by either setInstances or setBinders.
      */
     public static SystemInterface getInstance(String serviceName) {
+        if (sBinders != null) {
+            if (sBinders.containsKey(serviceName)) {
+                return new SystemInterface(sBinders.get(serviceName), serviceName);
+            }
+            throw new IllegalArgumentException("Cannot find any binder for " + serviceName);
+        }
         if (sInstances != null) {
             for (SystemInterface i : sInstances) {
                 if (i.getServiceName().equals(serviceName)) {
@@ -75,11 +106,16 @@ public class ServiceManagerInterface {
             throw new IllegalArgumentException("Cannot find any implementation for " + serviceName);
         }
 
-        return createSystemInterface(serviceName);
+        return new SystemInterface(getBinder(serviceName), serviceName);
     }
 
     @TestApi
     public static void setInstances(SystemInterface[] instances) {
         sInstances = instances;
     }
+
+    @TestApi
+    public static void setBinders(Map<String, IRemotelyProvisionedComponent> binders) {
+        sBinders = binders;
+    }
 }
diff --git a/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java b/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java
index 8c7179f..f3da031 100644
--- a/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java
+++ b/app/src/com/android/rkpdapp/provisioner/PeriodicProvisioner.java
@@ -49,6 +49,7 @@ import co.nstant.in.cbor.CborException;
 public class PeriodicProvisioner extends Worker {
     public static final String UNIQUE_WORK_NAME = "ProvisioningJob";
     private static final String TAG = "RkpdPeriodicProvisioner";
+    private static final boolean IS_ASYNC = true;
 
     private final Context mContext;
     private final ProvisionedKeyDao mKeyDao;
@@ -87,7 +88,7 @@ public class PeriodicProvisioner extends Worker {
             // Fetch geek from the server and figure out whether provisioning needs to be stopped.
             GeekResponse response;
             try {
-                response = new ServerInterface(mContext).fetchGeekAndUpdate(metrics);
+                response = new ServerInterface(mContext, IS_ASYNC).fetchGeekAndUpdate(metrics);
             } catch (InterruptedException | RkpdException e) {
                 Log.e(TAG, "Error fetching configuration from the RKP server", e);
                 return Result.failure();
@@ -104,7 +105,7 @@ public class PeriodicProvisioner extends Worker {
             }
 
             Log.i(TAG, "Total services found implementing IRPC: " + irpcs.length);
-            Provisioner provisioner = new Provisioner(mContext, mKeyDao);
+            Provisioner provisioner = new Provisioner(mContext, mKeyDao, IS_ASYNC);
             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 b8f6057..19db3a7 100644
--- a/app/src/com/android/rkpdapp/provisioner/Provisioner.java
+++ b/app/src/com/android/rkpdapp/provisioner/Provisioner.java
@@ -53,10 +53,13 @@ public class Provisioner {
 
     private final Context mContext;
     private final ProvisionedKeyDao mKeyDao;
+    private final boolean mIsAsync;
 
-    public Provisioner(final Context applicationContext, ProvisionedKeyDao keyDao) {
+    public Provisioner(final Context applicationContext, ProvisionedKeyDao keyDao,
+            boolean isAsync) {
         mContext = applicationContext;
         mKeyDao = keyDao;
+        mIsAsync = isAsync;
     }
 
     /**
@@ -154,7 +157,7 @@ public class Provisioner {
             throw new RkpdException(RkpdException.ErrorCode.INTERNAL_ERROR,
                     "Failed to serialize payload");
         }
-        return new ServerInterface(mContext).requestSignedCertificates(certRequest,
+        return new ServerInterface(mContext, mIsAsync).requestSignedCertificates(certRequest,
                 response.getChallenge(), metrics);
     }
 
diff --git a/app/src/com/android/rkpdapp/service/RemoteProvisioningService.java b/app/src/com/android/rkpdapp/service/RemoteProvisioningService.java
index c95bf7f..c92cbd7 100644
--- a/app/src/com/android/rkpdapp/service/RemoteProvisioningService.java
+++ b/app/src/com/android/rkpdapp/service/RemoteProvisioningService.java
@@ -39,6 +39,7 @@ import com.android.rkpdapp.utils.Settings;
 /** Provides the implementation for IRemoteProvisioning.aidl */
 public class RemoteProvisioningService extends Service {
     public static final String TAG = "com.android.rkpdapp";
+    private static final boolean IS_ASYNC = false;
     private final IRemoteProvisioning.Stub mBinder = new RemoteProvisioningBinder();
 
     @Override
@@ -75,9 +76,9 @@ public class RemoteProvisioningService extends Service {
                 }
 
                 ProvisionedKeyDao dao = RkpdDatabase.getDatabase(context).provisionedKeyDao();
-                Provisioner provisioner = new Provisioner(context, dao);
+                Provisioner provisioner = new Provisioner(context, dao, IS_ASYNC);
                 IRegistration.Stub registration = new RegistrationBinder(context, callerUid,
-                        systemInterface, dao, new ServerInterface(context), provisioner,
+                        systemInterface, dao, new ServerInterface(context, IS_ASYNC), provisioner,
                         ThreadPool.EXECUTOR);
                 metric.setResult(RkpdClientOperation.Result.SUCCESS);
                 callback.onSuccess(registration);
diff --git a/app/tests/e2e/Android.bp b/app/tests/e2e/Android.bp
index 208770c..baa003d 100644
--- a/app/tests/e2e/Android.bp
+++ b/app/tests/e2e/Android.bp
@@ -26,7 +26,7 @@ android_test {
         "androidx.test.rules",
         "androidx.work_work-testing",
         "platform-test-annotations",
-        "truth-prebuilt",
+        "truth",
     ],
     platform_apis: true,
     test_suites: [
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 46bd879..6f46896 100644
--- a/app/tests/e2e/src/com/android/rkpdapp/e2etest/KeystoreIntegrationTest.java
+++ b/app/tests/e2e/src/com/android/rkpdapp/e2etest/KeystoreIntegrationTest.java
@@ -46,6 +46,7 @@ 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.SystemInterfaceSelector;
 import com.android.rkpdapp.testutil.SystemPropertySetter;
 import com.android.rkpdapp.utils.Settings;
 import com.android.rkpdapp.utils.X509Utils;
@@ -129,7 +130,8 @@ public class KeystoreIntegrationTest {
         mKeyStore.load(null);
         mKeyDao.deleteAllKeys();
 
-        SystemInterface systemInterface = ServiceManagerInterface.getInstance(mServiceName);
+        SystemInterface systemInterface =
+                SystemInterfaceSelector.getSystemInterfaceForServiceName(mServiceName);
         ServiceManagerInterface.setInstances(new SystemInterface[] {systemInterface});
     }
 
@@ -228,9 +230,10 @@ public class KeystoreIntegrationTest {
         // Verify that if the system is set to rkp only, key creation fails when RKP is unable
         // to get keys.
 
-        try {
+        try (FakeRkpServer server = new FakeRkpServer(FakeRkpServer.Response.INTERNAL_ERROR,
+                FakeRkpServer.Response.INTERNAL_ERROR)) {
             Settings.setDeviceConfig(sContext, Settings.EXTRA_SIGNED_KEYS_AVAILABLE_DEFAULT,
-                    Duration.ofDays(1), "bad url");
+                    Duration.ofDays(1), server.getUrl());
             Settings.setMaxRequestTime(sContext, 100);
             createKeystoreKeyBackedByRkp();
             assertWithMessage("Should have gotten a KeyStoreException").fail();
@@ -249,14 +252,17 @@ public class KeystoreIntegrationTest {
                 .that(SystemProperties.getBoolean(getRkpOnlyProp(), false))
                 .isFalse();
 
-        Settings.setDeviceConfig(sContext, Settings.EXTRA_SIGNED_KEYS_AVAILABLE_DEFAULT,
-                Duration.ofDays(1), "bad url");
+        try (FakeRkpServer server = new FakeRkpServer(FakeRkpServer.Response.INTERNAL_ERROR,
+                FakeRkpServer.Response.INTERNAL_ERROR)) {
+            Settings.setDeviceConfig(sContext, Settings.EXTRA_SIGNED_KEYS_AVAILABLE_DEFAULT,
+                    Duration.ofDays(1), server.getUrl());
 
-        createKeystoreKey();
+            createKeystoreKey();
 
-        // Ensure the key has a cert, but it didn't come from rkpd.
-        assertThat(mKeyStore.getCertificateChain(getTestKeyAlias())).isNotEmpty();
-        assertThat(mKeyDao.getTotalKeysForIrpc(mServiceName)).isEqualTo(0);
+            // Ensure the key has a cert, but it didn't come from rkpd.
+            assertThat(mKeyStore.getCertificateChain(getTestKeyAlias())).isNotEmpty();
+            assertThat(mKeyDao.getTotalKeysForIrpc(mServiceName)).isEqualTo(0);
+        }
     }
 
     @Test
@@ -277,8 +283,9 @@ public class KeystoreIntegrationTest {
 
     @Test
     public void testRetryableRkpError() throws Exception {
-        try {
-            Settings.setDeviceConfig(sContext, 1, Duration.ofDays(1), "bad url");
+        try (FakeRkpServer server = new FakeRkpServer(FakeRkpServer.Response.INTERNAL_ERROR,
+                FakeRkpServer.Response.INTERNAL_ERROR)) {
+            Settings.setDeviceConfig(sContext, 1, Duration.ofDays(1), server.getUrl());
             Settings.setMaxRequestTime(sContext, 100);
             createKeystoreKeyBackedByRkp();
             Assert.fail("Expected a keystore exception");
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 0eb76be..c1762d5 100644
--- a/app/tests/e2e/src/com/android/rkpdapp/e2etest/RkpdHostTestHelperTests.java
+++ b/app/tests/e2e/src/com/android/rkpdapp/e2etest/RkpdHostTestHelperTests.java
@@ -42,6 +42,7 @@ 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.SystemInterfaceSelector;
 import com.android.rkpdapp.testutil.TestDatabase;
 import com.android.rkpdapp.testutil.TestProvisionedKeyDao;
 import com.android.rkpdapp.utils.Settings;
@@ -107,13 +108,15 @@ public class RkpdHostTestHelperTests {
         mRealDao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
         mRealDao.deleteAllKeys();
         mTestDao = Room.databaseBuilder(sContext, TestDatabase.class, DB_NAME).build().dao();
-        SystemInterface systemInterface = ServiceManagerInterface.getInstance(mServiceName);
-        ServiceManagerInterface.setInstances(new SystemInterface[] {systemInterface});
 
         mProvisioner = TestWorkerBuilder.from(
                 sContext,
                 PeriodicProvisioner.class,
                 Executors.newSingleThreadExecutor()).build();
+
+        SystemInterface systemInterface =
+                SystemInterfaceSelector.getSystemInterfaceForServiceName(mServiceName);
+        ServiceManagerInterface.setInstances(new SystemInterface[] {systemInterface});
     }
 
     @After
diff --git a/app/tests/hosttest/Android.bp b/app/tests/hosttest/Android.bp
index 282951c..c0bfd87 100644
--- a/app/tests/hosttest/Android.bp
+++ b/app/tests/hosttest/Android.bp
@@ -29,7 +29,7 @@ java_test_host {
         "host-libprotobuf-java-full",
         "platformprotos",
         "tradefed",
-        "truth-prebuilt",
+        "truth",
     ],
     static_libs: [
         "cts-statsd-atom-host-test-utils",
diff --git a/app/tests/hosttest/src/com/android/rkpdapp/hosttest/RkpdStatsTests.java b/app/tests/hosttest/src/com/android/rkpdapp/hosttest/RkpdStatsTests.java
index 823179d..e7bf7cc 100644
--- a/app/tests/hosttest/src/com/android/rkpdapp/hosttest/RkpdStatsTests.java
+++ b/app/tests/hosttest/src/com/android/rkpdapp/hosttest/RkpdStatsTests.java
@@ -40,6 +40,7 @@ import java.util.List;
 @RunWith(DeviceJUnit4ClassRunner.class)
 public final class RkpdStatsTests extends AtomsHostTest {
     private static final int NO_HTTP_STATUS_ERROR = 0;
+    private static final int HTTP_STATUS_SERVER_ERROR = 500;
     private static final int HTTPS_OK = 200;
     private static final String RPC_DEFAULT =
             "android.hardware.security.keymint.IRemotelyProvisionedComponent/default";
@@ -107,7 +108,7 @@ public final class RkpdStatsTests extends AtomsHostTest {
         assertThat(attempt.getUptime()).isNotEqualTo(UpTime.UPTIME_UNKNOWN);
         assertThat(attempt.getEnablement()).isEqualTo(Enablement.ENABLED_RKP_ONLY);
         assertThat(attempt.getStatus()).isEqualTo(
-                RemoteKeyProvisioningStatus.FETCH_GEEK_IO_EXCEPTION);
+                RemoteKeyProvisioningStatus.FETCH_GEEK_HTTP_ERROR);
 
         final RemoteKeyProvisioningTiming timing = getTimingMetric(data);
         assertThat(timing).isNotNull();
@@ -125,7 +126,7 @@ public final class RkpdStatsTests extends AtomsHostTest {
         assertThat(network).isNotNull();
         assertThat(network.getTransportType()).isEqualTo(timing.getTransportType());
         assertThat(network.getStatus()).isEqualTo(attempt.getStatus());
-        assertThat(network.getHttpStatusError()).isEqualTo(NO_HTTP_STATUS_ERROR);
+        assertThat(network.getHttpStatusError()).isEqualTo(HTTP_STATUS_SERVER_ERROR);
     }
 
     @Test
diff --git a/app/tests/stress/Android.bp b/app/tests/stress/Android.bp
index fa8f344..4320f2f 100644
--- a/app/tests/stress/Android.bp
+++ b/app/tests/stress/Android.bp
@@ -24,7 +24,7 @@ android_test {
         "androidx.test.core",
         "androidx.test.rules",
         "platform-test-annotations",
-        "truth-prebuilt",
+        "truth",
     ],
     platform_apis: true,
     test_suites: [
diff --git a/app/tests/stress/src/com/android/rkpdapp/stress/RegistrationBinderStressTest.java b/app/tests/stress/src/com/android/rkpdapp/stress/RegistrationBinderStressTest.java
index ce64d91..4ab5438 100644
--- a/app/tests/stress/src/com/android/rkpdapp/stress/RegistrationBinderStressTest.java
+++ b/app/tests/stress/src/com/android/rkpdapp/stress/RegistrationBinderStressTest.java
@@ -82,8 +82,10 @@ public class RegistrationBinderStressTest {
     }
 
     private RegistrationBinder createRegistrationBinder() {
+        boolean isAsync = false;
         return new RegistrationBinder(mContext, Process.myUid(), mIrpcHal, mKeyDao,
-                new ServerInterface(mContext), new Provisioner(mContext, mKeyDao), mExecutor);
+                new ServerInterface(mContext, isAsync), new Provisioner(mContext, mKeyDao, isAsync),
+                mExecutor);
     }
 
     private void getKeyHelper(int keyId) {
diff --git a/app/tests/unit/Android.bp b/app/tests/unit/Android.bp
index 13c5cbb..2a9ae30 100644
--- a/app/tests/unit/Android.bp
+++ b/app/tests/unit/Android.bp
@@ -30,7 +30,7 @@ android_test {
         "libnanohttpd",
         "mockito-target-extended-minus-junit4",
         "platform-test-annotations",
-        "truth-prebuilt",
+        "truth",
         "rkpdapp-tink-prebuilt-test-only",
         "bouncycastle-unbundled",
     ],
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 24a48f1..384b172 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/ProvisionerTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/ProvisionerTest.java
@@ -73,7 +73,7 @@ public class ProvisionerTest {
         ProvisionedKeyDao keyDao = RkpdDatabase.getDatabase(sContext).provisionedKeyDao();
         keyDao.deleteAllKeys();
 
-        mProvisioner = new Provisioner(sContext, keyDao);
+        mProvisioner = new Provisioner(sContext, keyDao, false);
     }
 
     @After
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 725b1a6..896a4b3 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/ServerInterfaceTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/ServerInterfaceTest.java
@@ -61,7 +61,7 @@ public class ServerInterfaceTest {
     @Before
     public void setUp() {
         Settings.clearPreferences(sContext);
-        mServerInterface = new ServerInterface(sContext);
+        mServerInterface = new ServerInterface(sContext, false);
     }
 
     @After
diff --git a/app/tests/unit/src/com/android/rkpdapp/unittest/SystemInterfaceTest.java b/app/tests/unit/src/com/android/rkpdapp/unittest/SystemInterfaceTest.java
index 92dc16f..c201e11 100644
--- a/app/tests/unit/src/com/android/rkpdapp/unittest/SystemInterfaceTest.java
+++ b/app/tests/unit/src/com/android/rkpdapp/unittest/SystemInterfaceTest.java
@@ -25,6 +25,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -53,8 +54,10 @@ import com.android.rkpdapp.interfaces.SystemInterface;
 import com.android.rkpdapp.metrics.ProvisioningAttempt;
 import com.android.rkpdapp.utils.CborUtils;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.crypto.tink.subtle.Ed25519Sign;
 
+import org.junit.After;
 import org.junit.Assert;
 import org.junit.Assume;
 import org.junit.Before;
@@ -93,6 +96,12 @@ public class SystemInterfaceTest {
         Assume.assumeTrue(ServiceManager.isDeclared(SERVICE));
     }
 
+    @After
+    public void cleanUp() {
+        ServiceManagerInterface.setBinders(null);
+        ServiceManagerInterface.setInstances(null);
+    }
+
     @Test
     public void testGetDeclaredInstances() {
         SystemInterface[] instances = ServiceManagerInterface.getAllInstances();
@@ -109,8 +118,32 @@ public class SystemInterfaceTest {
         try {
             ServiceManagerInterface.getInstance("non-existent");
             fail("Getting the declared service 'non-existent' should fail due to SEPolicy.");
-        } catch (RuntimeException e) {
-            assertThat(e).isInstanceOf(SecurityException.class);
+        } catch (IllegalArgumentException | SecurityException e) {
+            // Pre-android V, we'd get SecurityException. Post-V we get IllegalArgumentException
+        }
+    }
+
+    @Test
+    public void testGetAllInstancesWithUnsupportedService() throws RemoteException {
+        IRemotelyProvisionedComponent mockedBinder = mock(IRemotelyProvisionedComponent.class);
+        doThrow(new UnsupportedOperationException()).when(mockedBinder).getHardwareInfo();
+        ServiceManagerInterface.setBinders(ImmutableMap.of(SERVICE, mockedBinder));
+
+        SystemInterface[] instances = ServiceManagerInterface.getAllInstances();
+        assertThat(instances).isEmpty();
+    }
+
+    @Test
+    public void testGetInstanceWithUnsupportedService() throws RemoteException {
+        IRemotelyProvisionedComponent mockedBinder = mock(IRemotelyProvisionedComponent.class);
+        doThrow(new UnsupportedOperationException()).when(mockedBinder).getHardwareInfo();
+        ServiceManagerInterface.setBinders(ImmutableMap.of(SERVICE, mockedBinder));
+
+        try {
+            SystemInterface ignored = ServiceManagerInterface.getInstance(SERVICE);
+            fail("Expected UnsupportedOperationException");
+        } catch (UnsupportedOperationException e) {
+            // pass
         }
     }
 
diff --git a/app/tests/util/Android.bp b/app/tests/util/Android.bp
index 6b8daa8..c12f372 100644
--- a/app/tests/util/Android.bp
+++ b/app/tests/util/Android.bp
@@ -25,7 +25,7 @@ java_library {
         "androidx.room_room-runtime",
         "libnanohttpd",
         "libprotobuf-java-lite",
-        "truth-prebuilt",
+        "truth",
     ],
     plugins: [
         "androidx.room_room-compiler-plugin",
diff --git a/app/tests/util/src/com/android/rkpdapp/testutil/SystemInterfaceSelector.java b/app/tests/util/src/com/android/rkpdapp/testutil/SystemInterfaceSelector.java
new file mode 100644
index 0000000..0f49691
--- /dev/null
+++ b/app/tests/util/src/com/android/rkpdapp/testutil/SystemInterfaceSelector.java
@@ -0,0 +1,44 @@
+/*
+ * 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 static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.TruthJUnit.assume;
+
+import android.hardware.security.keymint.IRemotelyProvisionedComponent;
+
+import com.android.rkpdapp.interfaces.ServiceManagerInterface;
+import com.android.rkpdapp.interfaces.SystemInterface;
+
+public class SystemInterfaceSelector {
+    /**
+     * Gets the system interface object for provided service name.
+     */
+    public static SystemInterface getSystemInterfaceForServiceName(String serviceName) {
+        SystemInterface matchingInterface = null;
+        for (SystemInterface systemInterface: ServiceManagerInterface.getAllInstances()) {
+            if (systemInterface.getServiceName().equals(serviceName)) {
+                matchingInterface = systemInterface;
+            }
+        }
+        if (matchingInterface == null) {
+            assertThat(serviceName).isEqualTo(IRemotelyProvisionedComponent.DESCRIPTOR + "/avf");
+            assume().withMessage("AVF is not supported by this system").fail();
+        }
+        return matchingInterface;
+    }
+}
diff --git a/system-server/tests/unit/Android.bp b/system-server/tests/unit/Android.bp
index 3aac9e7..a44c774 100644
--- a/system-server/tests/unit/Android.bp
+++ b/system-server/tests/unit/Android.bp
@@ -24,7 +24,7 @@ android_test {
         "androidx.test.runner",
         "mockito-target",
         "service-rkp.impl",
-        "truth-prebuilt",
+        "truth",
     ],
     libs: [
         "android.test.mock",
diff --git a/system-server/tests/unit/AndroidTest.xml b/system-server/tests/unit/AndroidTest.xml
index 5c83c54..517d827 100644
--- a/system-server/tests/unit/AndroidTest.xml
+++ b/system-server/tests/unit/AndroidTest.xml
@@ -32,5 +32,6 @@
             class="com.android.tradefed.testtype.suite.module.MainlineTestModuleController">
         <option name="enable" value="true" />
         <option name="mainline-module-package-name" value="com.android.rkpd" />
+        <option name="mainline-module-package-name" value="com.google.android.rkpd" />
     </object>
 </configuration>