diff --git a/compos/benchmark/AndroidTest.xml b/compos/benchmark/AndroidTest.xml
index 0b6d0b1..f98b743 100644
--- a/compos/benchmark/AndroidTest.xml
+++ b/compos/benchmark/AndroidTest.xml
@@ -27,11 +27,6 @@
         <option name="force-root" value="true" />
     </target_preparer>
 
-    <target_preparer class="com.android.tradefed.targetprep.DeviceSetup">
-        <!-- Run in single thread to avoid nondeterministics. -->
-        <option name="set-property" key="dalvik.vm.boot-dex2oat-threads" value="1" />
-    </target_preparer>
-
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="com.android.compos.benchmark" />
         <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
diff --git a/microdroid/vm_payload/include/vm_payload.h b/microdroid/vm_payload/include/vm_payload.h
index 4ed07b8..6e065a5 100644
--- a/microdroid/vm_payload/include/vm_payload.h
+++ b/microdroid/vm_payload/include/vm_payload.h
@@ -49,8 +49,7 @@
 /**
  * Get the VM's DICE attestation chain.
  *
- * TODO: don't expose the contained privacy breaking identifiers to the payload
- * TODO: keep the DICE chain as an internal detail for as long as possible
+ * This function will fail if the use of restricted APIs is not permitted.
  *
  * \param data pointer to size bytes where the chain is written.
  * \param size number of bytes that can be written to data.
@@ -63,7 +62,7 @@
 /**
  * Get the VM's DICE attestation CDI.
  *
- * TODO: don't expose the raw CDI, only derived values
+ * This function will fail if the use of restricted APIs is not permitted.
  *
  * \param data pointer to size bytes where the CDI is written.
  * \param size number of bytes that can be written to data.
diff --git a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
index d3ebe5c..4dd3db6 100644
--- a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
+++ b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
@@ -39,18 +39,23 @@
     /**
      * Gets the DICE attestation chain for the VM.
      *
-     * STOPSHIP:
-     * TODO: don't expose this to untrusted payloads as it contains privacy breaking identifiers.
+     * The DICE chain must not be made available to all VMs as it contains privacy breaking
+     * identifiers.
+     *
+     * @return the VM's raw DICE certificate chain.
+     * @throws SecurityException if the use of test APIs is not permitted.
      */
     byte[] getDiceAttestationChain();
 
     /**
      * Gets the DICE attestation CDI for the VM.
      *
-     * STOPSHIP:
-     * TODO: A better API would handle key derivation on behalf of the payload so they can't forget
-     * to do it themselves. It also means the payload doesn't get the raw CDI so there's less chance
-     * of it leaking.
+     * The raw attestation CDI isn't very useful but is used for smoke tests. A better API would
+     * handle key derivation on behalf of the payload so they can't forget to do it themselves and
+     * would also mean the payload doesn't get the raw CDI which reduces the chance of it leaking.
+     *
+     * @return the VM's raw attestation CDI.
+     * @throws SecurityException if the use of test APIs is not permitted.
      */
     byte[] getDiceAttestationCdi();
 }
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 748b497..00c3dce 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -404,6 +404,13 @@
     )
     .context("Failed to run zipfuse")?;
 
+    // Restricted APIs are only allowed to be used by platform or test components. Infer this from
+    // the use of a VM config file since those can only be used by platform and test components.
+    let allow_restricted_apis = match payload_metadata {
+        PayloadMetadata::config_path(_) => true,
+        PayloadMetadata::config(_) => false,
+    };
+
     let config = load_config(payload_metadata).context("Failed to load payload metadata")?;
 
     let task = config
@@ -439,7 +446,7 @@
     // Wait until zipfuse has mounted the APK so we can access the payload
     wait_for_property_true(APK_MOUNT_DONE_PROP).context("Failed waiting for APK mount done")?;
 
-    register_vm_payload_service(service.clone(), dice)?;
+    register_vm_payload_service(allow_restricted_apis, service.clone(), dice)?;
     ProcessState::start_thread_pool();
 
     system_properties::write("dev.bootcomplete", "1").context("set dev.bootcomplete")?;
diff --git a/microdroid_manager/src/vm_payload_service.rs b/microdroid_manager/src/vm_payload_service.rs
index 4b47ad9..159bf67 100644
--- a/microdroid_manager/src/vm_payload_service.rs
+++ b/microdroid_manager/src/vm_payload_service.rs
@@ -26,6 +26,7 @@
 
 /// Implementation of `IVmPayloadService`.
 struct VmPayloadService {
+    allow_restricted_apis: bool,
     virtual_machine_service: Strong<dyn IVirtualMachineService>,
     dice: DiceContext,
 }
@@ -54,10 +55,12 @@
     }
 
     fn getDiceAttestationChain(&self) -> binder::Result<Vec<u8>> {
+        self.check_restricted_apis_allowed()?;
         Ok(self.dice.bcc.clone())
     }
 
     fn getDiceAttestationCdi(&self) -> binder::Result<Vec<u8>> {
+        self.check_restricted_apis_allowed()?;
         Ok(self.dice.cdi_attest.to_vec())
     }
 }
@@ -66,18 +69,32 @@
 
 impl VmPayloadService {
     /// Creates a new `VmPayloadService` instance from the `IVirtualMachineService` reference.
-    fn new(vm_service: Strong<dyn IVirtualMachineService>, dice: DiceContext) -> Self {
-        Self { virtual_machine_service: vm_service, dice }
+    fn new(
+        allow_restricted_apis: bool,
+        vm_service: Strong<dyn IVirtualMachineService>,
+        dice: DiceContext,
+    ) -> Self {
+        Self { allow_restricted_apis, virtual_machine_service: vm_service, dice }
+    }
+
+    fn check_restricted_apis_allowed(&self) -> binder::Result<()> {
+        if self.allow_restricted_apis {
+            Ok(())
+        } else {
+            error!("Use of restricted APIs is not allowed");
+            Err(Status::new_exception_str(ExceptionCode::SECURITY, Some("Use of restricted APIs")))
+        }
     }
 }
 
 /// Registers the `IVmPayloadService` service.
 pub(crate) fn register_vm_payload_service(
+    allow_restricted_apis: bool,
     vm_service: Strong<dyn IVirtualMachineService>,
     dice: DiceContext,
 ) -> Result<()> {
     let vm_payload_binder = BnVmPayloadService::new_binder(
-        VmPayloadService::new(vm_service, dice),
+        VmPayloadService::new(allow_restricted_apis, vm_service, dice),
         BinderFeatures::default(),
     );
     add_service(VM_PAYLOAD_SERVICE_NAME, vm_payload_binder.as_binder())
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index 23d546d..db74358 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -89,7 +89,7 @@
     private boolean canBootMicrodroidWithMemory(int mem)
             throws VirtualMachineException, InterruptedException, IOException {
         VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                .setPayloadBinaryPath("MicrodroidIdleNativeLib.so")
                 .setDebugLevel(DEBUG_LEVEL_NONE)
                 .setMemoryMib(mem)
                 .build();
@@ -148,7 +148,7 @@
 
             // To grab boot events from log, set debug mode to FULL
             VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                    .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                    .setPayloadBinaryPath("MicrodroidIdleNativeLib.so")
                     .setDebugLevel(DEBUG_LEVEL_FULL)
                     .setMemoryMib(256)
                     .build();
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 297341b..5c9cf42 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -261,9 +261,10 @@
                                     vm.connectToVsockServer(ITestService.SERVICE_PORT));
                             vmCdis.cdiAttest = testService.insecurelyExposeAttestationCdi();
                             vmCdis.instanceSecret = testService.insecurelyExposeVmInstanceSecret();
-                            forceStop(vm);
                         } catch (Exception e) {
                             exception.complete(e);
+                        } finally {
+                            forceStop(vm);
                         }
                     }
                 };
@@ -280,8 +281,9 @@
     public void instancesOfSameVmHaveDifferentCdis() throws Exception {
         assumeSupportedKernel();
 
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
         VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                .setPayloadConfigPath("assets/vm_config.json")
                 .setDebugLevel(DEBUG_LEVEL_FULL)
                 .build();
         mInner.forceCreateNewVirtualMachine("test_vm_a", normalConfig);
@@ -304,8 +306,9 @@
     public void sameInstanceKeepsSameCdis() throws Exception {
         assumeSupportedKernel();
 
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
         VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                .setPayloadConfigPath("assets/vm_config.json")
                 .setDebugLevel(DEBUG_LEVEL_FULL)
                 .build();
         mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
@@ -326,8 +329,9 @@
     public void bccIsSuperficiallyWellFormed() throws Exception {
         assumeSupportedKernel();
 
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
         VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                .setPayloadConfigPath("assets/vm_config.json")
                 .setDebugLevel(DEBUG_LEVEL_FULL)
                 .build();
         VirtualMachine vm = mInner.forceCreateNewVirtualMachine("bcc_vm", normalConfig);
@@ -341,9 +345,10 @@
                             ITestService testService = ITestService.Stub.asInterface(
                                     vm.connectToVsockServer(ITestService.SERVICE_PORT));
                             bcc.complete(testService.getBcc());
-                            forceStop(vm);
                         } catch (Exception e) {
                             exception.complete(e);
+                        } finally {
+                            forceStop(vm);
                         }
                     }
                 };