Merge "Enable again tests booting microdroid with vendor partition" into main
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 4da96c8..f146b4e 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -12,6 +12,9 @@
       "name": "MicrodroidTestApp"
     },
     {
+      "name": "VmAttestationTestApp"
+    },
+    {
       "name": "CustomPvmfwHostTestCases"
     },
     {
diff --git a/javalib/api/test-current.txt b/javalib/api/test-current.txt
index 5aff93f..3ea50e2 100644
--- a/javalib/api/test-current.txt
+++ b/javalib/api/test-current.txt
@@ -2,6 +2,7 @@
 package android.system.virtualmachine {
 
   public class VirtualMachine implements java.lang.AutoCloseable {
+    method @FlaggedApi("RELEASE_AVF_ENABLE_REMOTE_ATTESTATION") @RequiresPermission(android.system.virtualmachine.VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION) public void enableTestAttestation() throws android.system.virtualmachine.VirtualMachineException;
     method @NonNull @WorkerThread public java.io.OutputStream getConsoleInput() throws android.system.virtualmachine.VirtualMachineException;
     method @NonNull public java.io.File getRootDir();
   }
@@ -26,6 +27,7 @@
     method @RequiresPermission(android.system.virtualmachine.VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION) public boolean isFeatureEnabled(String) throws android.system.virtualmachine.VirtualMachineException;
     field public static final String FEATURE_DICE_CHANGES = "com.android.kvm.DICE_CHANGES";
     field public static final String FEATURE_MULTI_TENANT = "com.android.kvm.MULTI_TENANT";
+    field public static final String FEATURE_REMOTE_ATTESTATION = "com.android.kvm.REMOTE_ATTESTATION";
     field public static final String FEATURE_VENDOR_MODULES = "com.android.kvm.VENDOR_MODULES";
   }
 
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 5025e88..b4ba00b 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -42,6 +42,7 @@
 
 import static java.util.Objects.requireNonNull;
 
+import android.annotation.FlaggedApi;
 import android.annotation.CallbackExecutor;
 import android.annotation.IntDef;
 import android.annotation.IntRange;
@@ -1202,6 +1203,28 @@
     }
 
     /**
+     * Enables the VM to request attestation in testing mode.
+     *
+     * <p>This function provisions a key pair for the VM attestation testing, a fake certificate
+     * will be associated to the fake key pair when the VM requests attestation in testing mode.
+     *
+     * <p>The provisioned key pair can only be used in subsequent calls to {@link
+     * AVmPayload_requestAttestationForTesting} within a running VM.
+     *
+     * @hide
+     */
+    @TestApi
+    @RequiresPermission(USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION)
+    @FlaggedApi("RELEASE_AVF_ENABLE_REMOTE_ATTESTATION")
+    public void enableTestAttestation() throws VirtualMachineException {
+        try {
+            mVirtualizationService.getBinder().enableTestAttestation();
+        } catch (RemoteException e) {
+            throw e.rethrowAsRuntimeException();
+        }
+    }
+
+    /**
      * Captures the current state of the VM in a {@link VirtualMachineDescriptor} instance. The VM
      * needs to be stopped to avoid inconsistency in its state representation.
      *
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
index 1607c0a..1a4b53a 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
@@ -116,7 +116,12 @@
     @Retention(RetentionPolicy.SOURCE)
     @StringDef(
             prefix = "FEATURE_",
-            value = {FEATURE_DICE_CHANGES, FEATURE_MULTI_TENANT, FEATURE_VENDOR_MODULES})
+            value = {
+                FEATURE_DICE_CHANGES,
+                FEATURE_MULTI_TENANT,
+                FEATURE_REMOTE_ATTESTATION,
+                FEATURE_VENDOR_MODULES
+            })
     public @interface Features {}
 
     /**
@@ -136,6 +141,15 @@
     public static final String FEATURE_MULTI_TENANT = IVirtualizationService.FEATURE_MULTI_TENANT;
 
     /**
+     * Feature to allow remote attestation in Microdroid.
+     *
+     * @hide
+     */
+    @TestApi
+    public static final String FEATURE_REMOTE_ATTESTATION =
+            IVirtualizationService.FEATURE_REMOTE_ATTESTATION;
+
+    /**
      * Feature to allow vendor modules in Microdroid.
      *
      * @hide
diff --git a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
index 4813b35..b7a539b 100644
--- a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
+++ b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
@@ -107,9 +107,13 @@
      * serving as proof of the freshness of the result.
      *
      * @param challenge the maximum supported challenge size is 64 bytes.
+     * @param testMode whether the attestation is only for testing purposes. If testMode is true,
+     * caller must invoke {@link VirtualMachineManager#enableTestAttestation} prior to
+     * calling this method to provision a key pair to sign the attested result, and the returned
+     * certificate chain will not be RKP server rooted.
      *
      * @return An {@link AttestationResult} parcelable containing an attested key pair and its
      *         certification chain.
      */
-    AttestationResult requestAttestation(in byte[] challenge);
+    AttestationResult requestAttestation(in byte[] challenge, in boolean testMode);
 }
diff --git a/microdroid_manager/src/vm_payload_service.rs b/microdroid_manager/src/vm_payload_service.rs
index 20a1b89..959197a 100644
--- a/microdroid_manager/src/vm_payload_service.rs
+++ b/microdroid_manager/src/vm_payload_service.rs
@@ -68,7 +68,11 @@
         Ok(self.secret.dice_artifacts().cdi_attest().to_vec())
     }
 
-    fn requestAttestation(&self, challenge: &[u8]) -> binder::Result<AttestationResult> {
+    fn requestAttestation(
+        &self,
+        challenge: &[u8],
+        test_mode: bool,
+    ) -> binder::Result<AttestationResult> {
         self.check_restricted_apis_allowed()?;
         let ClientVmAttestationData { private_key, csr } =
             generate_attestation_key_and_csr(challenge, self.secret.dice_artifacts())
@@ -88,7 +92,7 @@
                 )
             })
             .with_log()?;
-        let cert_chain = self.virtual_machine_service.requestAttestation(&csr)?;
+        let cert_chain = self.virtual_machine_service.requestAttestation(&csr, test_mode)?;
         Ok(AttestationResult {
             privateKey: private_key.as_slice().to_vec(),
             certificateChain: cert_chain,
diff --git a/service_vm/test_apk/Android.bp b/service_vm/test_apk/Android.bp
new file mode 100644
index 0000000..8f5fb41
--- /dev/null
+++ b/service_vm/test_apk/Android.bp
@@ -0,0 +1,44 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_test {
+    name: "VmAttestationTestApp",
+    test_suites: [
+        "general-tests",
+    ],
+    srcs: ["src/java/**/*.java"],
+    static_libs: [
+        "MicrodroidDeviceTestHelper",
+        "androidx.test.runner",
+        "androidx.test.ext.junit",
+        "com.android.virt.vm_attestation.testservice-java",
+        "truth",
+    ],
+    jni_libs: ["libvm_attestation_test_payload"],
+    jni_uses_platform_apis: true,
+    use_embedded_native_libs: true,
+    sdk_version: "test_current",
+    compile_multilib: "first",
+}
+
+rust_defaults {
+    name: "vm_attestation_test_payload_defaults",
+    crate_name: "vm_attestation_test_payload",
+    defaults: ["avf_build_flags_rust"],
+    srcs: ["src/native/main.rs"],
+    prefer_rlib: true,
+    rustlibs: [
+        "com.android.virt.vm_attestation.testservice-rust",
+        "libandroid_logger",
+        "libanyhow",
+        "libavflog",
+        "liblog_rust",
+        "libvm_payload_bindgen",
+    ],
+}
+
+rust_ffi {
+    name: "libvm_attestation_test_payload",
+    defaults: ["vm_attestation_test_payload_defaults"],
+}
diff --git a/service_vm/test_apk/AndroidManifest.xml b/service_vm/test_apk/AndroidManifest.xml
new file mode 100644
index 0000000..b998b7f
--- /dev/null
+++ b/service_vm/test_apk/AndroidManifest.xml
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+      package="com.android.virt.vm_attestation.testapp">
+     <uses-permission android:name="android.permission.MANAGE_VIRTUAL_MACHINE" />
+     <uses-permission android:name="android.permission.USE_CUSTOM_VIRTUAL_MACHINE" />
+
+     <application />
+     <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+        android:targetPackage="com.android.virt.vm_attestation.testapp"
+        android:label="Microdroid VM attestation" />
+</manifest>
diff --git a/service_vm/test_apk/AndroidTest.xml b/service_vm/test_apk/AndroidTest.xml
new file mode 100644
index 0000000..18b4e46
--- /dev/null
+++ b/service_vm/test_apk/AndroidTest.xml
@@ -0,0 +1,34 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<configuration description="Runs Microdroid VM remote attestation tests.">
+    <option name="config-descriptor:metadata" key="component" value="security" />
+    <option name="config-descriptor:metadata" key="parameter" value="not_instant_app" />
+    <option name="config-descriptor:metadata" key="parameter" value="not_multi_abi" />
+    <option name="config-descriptor:metadata" key="parameter" value="secondary_user" />
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="test-file-name" value="VmAttestationTestApp.apk" />
+    </target_preparer>
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="run-command" value="mkdir -p /data/local/tmp/cts/microdroid" />
+        <option name="teardown-command" value="rm -rf /data/local/tmp/cts/microdroid" />
+    </target_preparer>
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.virt.vm_attestation.testapp" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+        <option name="shell-timeout" value="300000" />
+        <option name="test-timeout" value="300000" />
+    </test>
+</configuration>
diff --git a/service_vm/test_apk/aidl/Android.bp b/service_vm/test_apk/aidl/Android.bp
new file mode 100644
index 0000000..3ecce46
--- /dev/null
+++ b/service_vm/test_apk/aidl/Android.bp
@@ -0,0 +1,18 @@
+package {
+    default_team: "trendy_team_virtualization",
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+aidl_interface {
+    name: "com.android.virt.vm_attestation.testservice",
+    srcs: ["com/android/virt/vm_attestation/testservice/**/*.aidl"],
+    unstable: true,
+    backend: {
+        java: {
+            gen_rpc: true,
+        },
+        rust: {
+            enabled: true,
+        },
+    },
+}
diff --git a/service_vm/test_apk/aidl/com/android/virt/vm_attestation/testservice/IAttestationService.aidl b/service_vm/test_apk/aidl/com/android/virt/vm_attestation/testservice/IAttestationService.aidl
new file mode 100644
index 0000000..94a7b8d
--- /dev/null
+++ b/service_vm/test_apk/aidl/com/android/virt/vm_attestation/testservice/IAttestationService.aidl
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2024 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.virt.vm_attestation.testservice;
+
+/** {@hide} */
+interface IAttestationService {
+    const int PORT = 5679;
+
+    /**
+     * Requests attestation for testing.
+     *
+     * A fake key pair should be provisioned with the call to
+     * {@link VirtualMachine#enableTestAttestation()} before calling this method.
+     *
+     * The attestation result will be cached in the VM and can be validated with
+     * {@link #validateAttestationResult}.
+     */
+    void requestAttestationForTesting();
+
+    /**
+     * Validates the attestation result returned by the last call to
+     * {@link #requestAttestationForTesting}.
+     */
+    void validateAttestationResult();
+}
diff --git a/service_vm/test_apk/assets/config.json b/service_vm/test_apk/assets/config.json
new file mode 100644
index 0000000..caae3ce
--- /dev/null
+++ b/service_vm/test_apk/assets/config.json
@@ -0,0 +1,10 @@
+{
+    "os": {
+      "name": "microdroid"
+    },
+    "task": {
+      "type": "microdroid_launcher",
+      "command": "libvm_attestation_test_payload.so"
+    },
+    "export_tombstones": true
+  }
\ No newline at end of file
diff --git a/service_vm/test_apk/src/java/com/android/virt/vm_attestation/testapp/VmAttestationTests.java b/service_vm/test_apk/src/java/com/android/virt/vm_attestation/testapp/VmAttestationTests.java
new file mode 100644
index 0000000..7771e83
--- /dev/null
+++ b/service_vm/test_apk/src/java/com/android/virt/vm_attestation/testapp/VmAttestationTests.java
@@ -0,0 +1,108 @@
+/*
+ * Copyright (C) 2024 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.virt.vm_attestation.testapp;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.TruthJUnit.assume;
+import static android.system.virtualmachine.VirtualMachineConfig.DEBUG_LEVEL_FULL;
+
+import android.system.virtualmachine.VirtualMachine;
+import android.system.virtualmachine.VirtualMachineConfig;
+import android.system.virtualmachine.VirtualMachineManager;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import com.android.microdroid.test.device.MicrodroidDeviceTestBase;
+import com.android.virt.vm_attestation.testservice.IAttestationService;
+
+@RunWith(Parameterized.class)
+public class VmAttestationTests extends MicrodroidDeviceTestBase {
+    private static final String TAG = "VmAttestationTest";
+    private static final String DEFAULT_CONFIG = "assets/config.json";
+
+    @Parameterized.Parameter(0)
+    public String mGki;
+
+    @Parameterized.Parameters(name = "gki={0}")
+    public static Collection<Object[]> params() {
+        List<Object[]> ret = new ArrayList<>();
+        ret.add(new Object[] {null /* use microdroid kernel */});
+        for (String gki : SUPPORTED_GKI_VERSIONS) {
+            ret.add(new Object[] {gki});
+        }
+        return ret;
+    }
+
+    @Before
+    public void setup() throws IOException {
+        grantPermission(VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION);
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
+        prepareTestSetup(true /* protectedVm */, mGki);
+        setMaxPerformanceTaskProfile();
+    }
+
+    @Test
+    public void requestingAttestationSucceeds() throws Exception {
+        assume().withMessage("Remote attestation is not supported on CF.")
+                .that(isCuttlefish())
+                .isFalse();
+        assumeFeatureEnabled(VirtualMachineManager.FEATURE_REMOTE_ATTESTATION);
+
+        VirtualMachineConfig.Builder builder =
+                newVmConfigBuilderWithPayloadConfig(DEFAULT_CONFIG)
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .setVmOutputCaptured(true);
+        VirtualMachineConfig config = builder.build();
+        VirtualMachine vm = forceCreateNewVirtualMachine("attestation_client", config);
+
+        vm.enableTestAttestation();
+        CompletableFuture<Exception> exception = new CompletableFuture<>();
+        CompletableFuture<Boolean> payloadReady = new CompletableFuture<>();
+        VmEventListener listener =
+                new VmEventListener() {
+                    @Override
+                    public void onPayloadReady(VirtualMachine vm) {
+                        payloadReady.complete(true);
+                        try {
+                            IAttestationService service =
+                                    IAttestationService.Stub.asInterface(
+                                            vm.connectToVsockServer(IAttestationService.PORT));
+                            android.os.Trace.beginSection("runningVmRequestsAttestation");
+                            service.requestAttestationForTesting();
+                            android.os.Trace.endSection();
+                            service.validateAttestationResult();
+                        } catch (Exception e) {
+                            exception.complete(e);
+                        } finally {
+                            forceStop(vm);
+                        }
+                    }
+                };
+
+        listener.runToFinish(TAG, vm);
+        assertThat(payloadReady.getNow(false)).isTrue();
+        assertThat(exception.getNow(null)).isNull();
+    }
+}
diff --git a/service_vm/test_apk/src/native/main.rs b/service_vm/test_apk/src/native/main.rs
new file mode 100644
index 0000000..26b8062
--- /dev/null
+++ b/service_vm/test_apk/src/native/main.rs
@@ -0,0 +1,271 @@
+// Copyright 2024, 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.
+
+//! Main executable of VM attestation for end-to-end testing.
+
+use anyhow::{anyhow, ensure, Result};
+use avflog::LogResult;
+use com_android_virt_vm_attestation_testservice::{
+    aidl::com::android::virt::vm_attestation::testservice::IAttestationService::{
+        BnAttestationService, IAttestationService, PORT,
+    },
+    binder::{self, unstable_api::AsNative, BinderFeatures, Interface, IntoBinderResult, Strong},
+};
+use log::{error, info};
+use std::{
+    ffi::{c_void, CStr},
+    panic,
+    ptr::{self, NonNull},
+    result,
+    sync::{Arc, Mutex},
+};
+use vm_payload_bindgen::{
+    attestation_status_t, AIBinder, AVmAttestationResult, AVmAttestationResult_free,
+    AVmAttestationResult_getCertificateAt, AVmAttestationResult_getCertificateCount,
+    AVmAttestationResult_getPrivateKey, AVmAttestationResult_resultToString,
+    AVmAttestationResult_sign, AVmPayload_notifyPayloadReady,
+    AVmPayload_requestAttestationForTesting, AVmPayload_runVsockRpcServer,
+};
+
+/// Entry point of the Service VM client.
+#[allow(non_snake_case)]
+#[no_mangle]
+pub extern "C" fn AVmPayload_main() {
+    android_logger::init_once(
+        android_logger::Config::default()
+            .with_tag("service_vm_client")
+            .with_min_level(log::Level::Debug),
+    );
+    // Redirect panic messages to logcat.
+    panic::set_hook(Box::new(|panic_info| {
+        error!("{}", panic_info);
+    }));
+    if let Err(e) = try_main() {
+        error!("failed with {:?}", e);
+        std::process::exit(1);
+    }
+}
+
+fn try_main() -> Result<()> {
+    info!("Welcome to Service VM Client!");
+
+    let mut service = AttestationService::new_binder().as_binder();
+    let service = service.as_native_mut() as *mut AIBinder;
+    let param = ptr::null_mut();
+    // SAFETY: We hold a strong pointer, so the raw pointer remains valid. The bindgen AIBinder
+    // is the same type as `sys::AIBinder`. It is safe for `on_ready` to be invoked at any time,
+    // with any parameter.
+    unsafe { AVmPayload_runVsockRpcServer(service, PORT.try_into()?, Some(on_ready), param) };
+}
+
+extern "C" fn on_ready(_param: *mut c_void) {
+    // SAFETY: It is safe to call `AVmPayload_notifyPayloadReady` at any time.
+    unsafe { AVmPayload_notifyPayloadReady() };
+}
+
+struct AttestationService {
+    res: Arc<Mutex<Option<AttestationResult>>>,
+}
+
+impl Interface for AttestationService {}
+
+impl AttestationService {
+    fn new_binder() -> Strong<dyn IAttestationService> {
+        let res = Arc::new(Mutex::new(None));
+        BnAttestationService::new_binder(AttestationService { res }, BinderFeatures::default())
+    }
+}
+
+impl IAttestationService for AttestationService {
+    fn requestAttestationForTesting(&self) -> binder::Result<()> {
+        // The data below is only a placeholder generated randomly with urandom
+        let challenge = &[
+            0x6c, 0xad, 0x52, 0x50, 0x15, 0xe7, 0xf4, 0x1d, 0xa5, 0x60, 0x7e, 0xd2, 0x7d, 0xf1,
+            0x51, 0x67, 0xc3, 0x3e, 0x73, 0x9b, 0x30, 0xbd, 0x04, 0x20, 0x2e, 0xde, 0x3b, 0x1d,
+            0xc8, 0x07, 0x11, 0x7b,
+        ];
+        let res = AttestationResult::request_attestation(challenge)
+            .map_err(|e| anyhow!("Unexpected status: {:?}", status_to_cstr(e)))
+            .with_log()
+            .or_service_specific_exception(-1)?;
+        *self.res.lock().unwrap() = Some(res);
+        Ok(())
+    }
+
+    fn validateAttestationResult(&self) -> binder::Result<()> {
+        // TODO(b/191073073): Returns the attestation result to the host for validation.
+        self.res.lock().unwrap().as_ref().unwrap().log().or_service_specific_exception(-1)
+    }
+}
+
+#[derive(Debug)]
+struct AttestationResult(NonNull<AVmAttestationResult>);
+
+// Safety: `AttestationResult` is not `Send` because it contains a raw pointer to a C struct.
+unsafe impl Send for AttestationResult {}
+
+impl AttestationResult {
+    fn request_attestation(challenge: &[u8]) -> result::Result<Self, attestation_status_t> {
+        let mut res: *mut AVmAttestationResult = ptr::null_mut();
+        // SAFETY: It is safe as we only read the challenge within its bounds and the
+        // function does not retain any reference to it.
+        let status = unsafe {
+            AVmPayload_requestAttestationForTesting(
+                challenge.as_ptr() as *const c_void,
+                challenge.len(),
+                &mut res,
+            )
+        };
+        if status == attestation_status_t::ATTESTATION_OK {
+            info!("Attestation succeeds. Status: {:?}", status_to_cstr(status));
+            let res = NonNull::new(res).expect("The attestation result is null");
+            Ok(Self(res))
+        } else {
+            Err(status)
+        }
+    }
+
+    fn certificate_chain(&self) -> Result<Vec<Box<[u8]>>> {
+        let num_certs = get_certificate_count(self.as_ref());
+        let mut certs = Vec::with_capacity(num_certs);
+        for i in 0..num_certs {
+            certs.push(get_certificate_at(self.as_ref(), i)?);
+        }
+        Ok(certs)
+    }
+
+    fn private_key(&self) -> Result<Box<[u8]>> {
+        get_private_key(self.as_ref())
+    }
+
+    fn sign(&self, message: &[u8]) -> Result<Box<[u8]>> {
+        sign_with_attested_key(self.as_ref(), message)
+    }
+
+    fn log(&self) -> Result<()> {
+        let cert_chain = self.certificate_chain()?;
+        info!("Attestation result certificateChain = {:?}", cert_chain);
+
+        let private_key = self.private_key()?;
+        info!("Attestation result privateKey = {:?}", private_key);
+
+        let message = b"Hello from Service VM client";
+        info!("Signing message: {:?}", message);
+        let signature = self.sign(message)?;
+        info!("Signature: {:?}", signature);
+        Ok(())
+    }
+}
+
+impl AsRef<AVmAttestationResult> for AttestationResult {
+    fn as_ref(&self) -> &AVmAttestationResult {
+        // SAFETY: This field is private, and only populated with a successful call to
+        // `AVmPayload_requestAttestation`.
+        unsafe { self.0.as_ref() }
+    }
+}
+
+impl Drop for AttestationResult {
+    fn drop(&mut self) {
+        // SAFETY: This field is private, and only populated with a successful call to
+        // `AVmPayload_requestAttestation`, and not freed elsewhere.
+        unsafe { AVmAttestationResult_free(self.0.as_ptr()) };
+    }
+}
+
+fn get_certificate_count(res: &AVmAttestationResult) -> usize {
+    // SAFETY: The result is returned by `AVmPayload_requestAttestation` and should be valid
+    // before getting freed.
+    unsafe { AVmAttestationResult_getCertificateCount(res) }
+}
+
+fn get_certificate_at(res: &AVmAttestationResult, index: usize) -> Result<Box<[u8]>> {
+    let size =
+        // SAFETY: The result is returned by `AVmPayload_requestAttestation` and should be valid
+        // before getting freed.
+        unsafe { AVmAttestationResult_getCertificateAt(res, index, ptr::null_mut(), 0) };
+    let mut cert = vec![0u8; size];
+    // SAFETY: The result is returned by `AVmPayload_requestAttestation` and should be valid
+    // before getting freed. This function only writes within the bounds of `cert`.
+    // And `cert` cannot overlap `res` because we just allocated it.
+    let size = unsafe {
+        AVmAttestationResult_getCertificateAt(
+            res,
+            index,
+            cert.as_mut_ptr() as *mut c_void,
+            cert.len(),
+        )
+    };
+    ensure!(size == cert.len());
+    Ok(cert.into_boxed_slice())
+}
+
+fn get_private_key(res: &AVmAttestationResult) -> Result<Box<[u8]>> {
+    let size =
+        // SAFETY: The result is returned by `AVmPayload_requestAttestation` and should be valid
+        // before getting freed.
+        unsafe { AVmAttestationResult_getPrivateKey(res, ptr::null_mut(), 0) };
+    let mut private_key = vec![0u8; size];
+    // SAFETY: The result is returned by `AVmPayload_requestAttestation` and should be valid
+    // before getting freed. This function only writes within the bounds of `private_key`.
+    // And `private_key` cannot overlap `res` because we just allocated it.
+    let size = unsafe {
+        AVmAttestationResult_getPrivateKey(
+            res,
+            private_key.as_mut_ptr() as *mut c_void,
+            private_key.len(),
+        )
+    };
+    ensure!(size == private_key.len());
+    Ok(private_key.into_boxed_slice())
+}
+
+fn sign_with_attested_key(res: &AVmAttestationResult, message: &[u8]) -> Result<Box<[u8]>> {
+    // SAFETY: The result is returned by `AVmPayload_requestAttestation` and should be valid
+    // before getting freed.
+    let size = unsafe {
+        AVmAttestationResult_sign(
+            res,
+            message.as_ptr() as *const c_void,
+            message.len(),
+            ptr::null_mut(),
+            0,
+        )
+    };
+    let mut signature = vec![0u8; size];
+    // SAFETY: The result is returned by `AVmPayload_requestAttestation` and should be valid
+    // before getting freed. This function only writes within the bounds of `signature`.
+    // And `signature` cannot overlap `res` because we just allocated it.
+    let size = unsafe {
+        AVmAttestationResult_sign(
+            res,
+            message.as_ptr() as *const c_void,
+            message.len(),
+            signature.as_mut_ptr() as *mut c_void,
+            signature.len(),
+        )
+    };
+    ensure!(size <= signature.len());
+    signature.truncate(size);
+    Ok(signature.into_boxed_slice())
+}
+
+fn status_to_cstr(status: attestation_status_t) -> &'static CStr {
+    // SAFETY: The function only reads the given enum status and returns a pointer to a
+    // static string.
+    let message = unsafe { AVmAttestationResult_resultToString(status) };
+    // SAFETY: The pointer returned by `AVmAttestationResult_resultToString` is guaranteed to
+    // point to a valid C String that lives forever.
+    unsafe { CStr::from_ptr(message) }
+}
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index d9d10ea..62df47d 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -19,9 +19,9 @@
 use crate::composite::make_composite_image;
 use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VmContext, VmInstance, VmState};
 use crate::debug_config::DebugConfig;
-use crate::dt_overlay::{create_device_tree_overlay, DtAddition, VM_REFERENCE_DT_ON_HOST_PATH, VM_DT_OVERLAY_MAX_SIZE, VM_DT_OVERLAY_PATH};
 use crate::payload::{add_microdroid_payload_images, add_microdroid_system_images, add_microdroid_vendor_image};
 use crate::selinux::{getfilecon, SeContext};
+use crate::reference_dt;
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
 use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::{
     Certificate::Certificate,
@@ -38,6 +38,7 @@
     IVirtualizationService::FEATURE_MULTI_TENANT,
     IVirtualizationService::FEATURE_VENDOR_MODULES,
     IVirtualizationService::FEATURE_DICE_CHANGES,
+    IVirtualizationService::FEATURE_REMOTE_ATTESTATION,
     MemoryTrimLevel::MemoryTrimLevel,
     Partition::Partition,
     PartitionType::PartitionType,
@@ -68,7 +69,6 @@
     Status, StatusCode, Strong,
     IntoBinderResult,
 };
-use cstr::cstr;
 use disk::QcowFile;
 use glob::glob;
 use lazy_static::lazy_static;
@@ -80,7 +80,6 @@
 use semver::VersionReq;
 use std::collections::HashSet;
 use std::convert::TryInto;
-use std::fs;
 use std::ffi::CStr;
 use std::fs::{canonicalize, read_dir, remove_file, File, OpenOptions};
 use std::io::{BufRead, BufReader, Error, ErrorKind, Seek, SeekFrom, Write};
@@ -309,6 +308,7 @@
         match feature {
             FEATURE_DICE_CHANGES => Ok(cfg!(dice_changes)),
             FEATURE_MULTI_TENANT => Ok(cfg!(multi_tenant)),
+            FEATURE_REMOTE_ATTESTATION => Ok(cfg!(remote_attestation)),
             FEATURE_VENDOR_MODULES => Ok(cfg!(vendor_modules)),
             _ => {
                 warn!("unknown feature {feature}");
@@ -316,6 +316,10 @@
             }
         }
     }
+
+    fn enableTestAttestation(&self) -> binder::Result<()> {
+        GLOBAL_SERVICE.enableTestAttestation()
+    }
 }
 
 impl VirtualizationService {
@@ -378,35 +382,13 @@
             check_gdb_allowed(config)?;
         }
 
-        // Currently, VirtMgr adds the host copy of reference DT & an untrusted prop (instance-id)
-        let mut dt_additions = Vec::with_capacity(2);
-        let host_ref_dt = Path::new(VM_REFERENCE_DT_ON_HOST_PATH);
-        if host_ref_dt.exists()
-            && read_dir(host_ref_dt).or_service_specific_exception(-1)?.next().is_some()
-        {
-            dt_additions.push(DtAddition::FromPath(host_ref_dt));
-        } else {
-            warn!("VM reference DT doesn't exist in host DT");
+        let reference_dt = reference_dt::parse_reference_dt(&temporary_directory)
+            .context("Failed to create VM reference DT")
+            .or_service_specific_exception(-1)?;
+        if reference_dt.is_none() {
+            warn!("VM reference DT doesn't exist");
         }
 
-        if cfg!(llpvm_changes) {
-            // TODO(b/291213394): Replace this with a per-VM instance Id.
-            let instance_id = b"sixtyfourbyteslonghardcoded_indeed_sixtyfourbyteslonghardcoded_h";
-            dt_additions.push(DtAddition::AvfUntrustedProp(cstr!("instance-id"), &instance_id[..]));
-        }
-
-        let device_tree_overlay = if !dt_additions.is_empty() {
-            let dt_output = temporary_directory.join(VM_DT_OVERLAY_PATH);
-            let mut data = [0_u8; VM_DT_OVERLAY_MAX_SIZE];
-            let fdt = create_device_tree_overlay(&mut data, &dt_additions)
-                .map_err(|e| anyhow!("Failed to create DT overlay, {e:?}"))
-                .or_service_specific_exception(-1)?;
-            fs::write(&dt_output, fdt.as_slice()).or_service_specific_exception(-1)?;
-            Some(File::open(dt_output).or_service_specific_exception(-1)?)
-        } else {
-            None
-        };
-
         let debug_level = match config {
             VirtualMachineConfig::AppConfig(config) => config.debugLevel,
             _ => DebugLevel::NONE,
@@ -555,7 +537,7 @@
             gdb_port,
             vfio_devices,
             dtbo,
-            device_tree_overlay,
+            reference_dt,
         };
         let instance = Arc::new(
             VmInstance::new(
@@ -1415,8 +1397,8 @@
         Ok(sk.map(|s| BnSecretkeeper::new_binder(SecretkeeperProxy(s), BinderFeatures::default())))
     }
 
-    fn requestAttestation(&self, csr: &[u8]) -> binder::Result<Vec<Certificate>> {
-        GLOBAL_SERVICE.requestAttestation(csr, get_calling_uid() as i32)
+    fn requestAttestation(&self, csr: &[u8], test_mode: bool) -> binder::Result<Vec<Certificate>> {
+        GLOBAL_SERVICE.requestAttestation(csr, get_calling_uid() as i32, test_mode)
     }
 }
 
diff --git a/virtualizationmanager/src/crosvm.rs b/virtualizationmanager/src/crosvm.rs
index 2c23441..84c60bd 100644
--- a/virtualizationmanager/src/crosvm.rs
+++ b/virtualizationmanager/src/crosvm.rs
@@ -118,7 +118,7 @@
     pub gdb_port: Option<NonZeroU16>,
     pub vfio_devices: Vec<VfioDevice>,
     pub dtbo: Option<File>,
-    pub device_tree_overlay: Option<File>,
+    pub reference_dt: Option<File>,
 }
 
 /// A disk image to pass to crosvm for a VM.
@@ -896,8 +896,10 @@
         .arg("--socket")
         .arg(add_preserved_fd(&mut preserved_fds, &control_server_socket.as_raw_descriptor()));
 
-    if let Some(dt_overlay) = &config.device_tree_overlay {
-        command.arg("--device-tree-overlay").arg(add_preserved_fd(&mut preserved_fds, dt_overlay));
+    if let Some(reference_dt) = &config.reference_dt {
+        command
+            .arg("--device-tree-overlay")
+            .arg(add_preserved_fd(&mut preserved_fds, reference_dt));
     }
 
     append_platform_devices(&mut command, &mut preserved_fds, &config)?;
diff --git a/virtualizationmanager/src/dt_overlay.rs b/virtualizationmanager/src/dt_overlay.rs
deleted file mode 100644
index c2d3c07..0000000
--- a/virtualizationmanager/src/dt_overlay.rs
+++ /dev/null
@@ -1,133 +0,0 @@
-// Copyright 2024, 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.
-
-//! This module support creating AFV related overlays, that can then be appended to DT by VM.
-
-use anyhow::{anyhow, Result};
-use cstr::cstr;
-use fsfdt::FsFdt;
-use libfdt::Fdt;
-use std::ffi::CStr;
-use std::path::Path;
-
-pub(crate) const AVF_NODE_NAME: &CStr = cstr!("avf");
-pub(crate) const UNTRUSTED_NODE_NAME: &CStr = cstr!("untrusted");
-pub(crate) const VM_REFERENCE_DT_ON_HOST_PATH: &str = "/proc/device-tree/avf/reference";
-pub(crate) const VM_DT_OVERLAY_PATH: &str = "vm_dt_overlay.dtbo";
-pub(crate) const VM_DT_OVERLAY_MAX_SIZE: usize = 2000;
-
-/// Provide ways to modify the device tree.
-#[derive(PartialEq, Eq)]
-pub(crate) enum DtAddition<'a> {
-    /// Include the device tree at given path.
-    FromPath(&'a Path),
-    /// Include a property in /avf/untrusted node. This node is used to specify host provided
-    /// properties such as `instance-id`.
-    /// pVM firmware does minimal validation of properties in this node.
-    AvfUntrustedProp(&'a CStr, &'a [u8]),
-}
-
-/// Given a list of `dt_additions`, return a Device tree overlay containing those!
-/// Example: with `create_device_tree_overlay(_, DtAddition::AvfUntrustedProp("instance-id", _))`
-/// ```
-///   {
-///     fragment@0 {
-///         target-path = "/";
-///         __overlay__ {
-///             avf {
-///                 untrusted { instance-id = [0x01 0x23 .. ] }
-///               }
-///             };
-///         };
-///     };
-/// };
-/// ```
-pub(crate) fn create_device_tree_overlay<'a>(
-    buffer: &'a mut [u8],
-    dt_additions: &[DtAddition],
-) -> Result<&'a mut Fdt> {
-    if dt_additions.is_empty() {
-        return Err(anyhow!("Expected non empty list of device tree additions"));
-    }
-
-    let (additional_properties, additional_paths): (Vec<_>, _) =
-        dt_additions.iter().partition(|o| matches!(o, DtAddition::AvfUntrustedProp(_, _)));
-
-    let fdt =
-        Fdt::create_empty_tree(buffer).map_err(|e| anyhow!("Failed to create empty Fdt: {e:?}"))?;
-    let mut root = fdt.root_mut().map_err(|e| anyhow!("Failed to get root: {e:?}"))?;
-    let mut node =
-        root.add_subnode(cstr!("fragment@0")).map_err(|e| anyhow!("Failed to fragment: {e:?}"))?;
-    node.setprop(cstr!("target-path"), b"/\0")
-        .map_err(|e| anyhow!("Failed to set target-path: {e:?}"))?;
-    let mut node = node
-        .add_subnode(cstr!("__overlay__"))
-        .map_err(|e| anyhow!("Failed to __overlay__ node: {e:?}"))?;
-
-    if !additional_properties.is_empty() {
-        let mut node = node
-            .add_subnode(AVF_NODE_NAME)
-            .map_err(|e| anyhow!("Failed to add avf node: {e:?}"))?;
-        let mut node = node
-            .add_subnode(UNTRUSTED_NODE_NAME)
-            .map_err(|e| anyhow!("Failed to add /avf/untrusted node: {e:?}"))?;
-        for prop in additional_properties {
-            if let DtAddition::AvfUntrustedProp(name, value) = prop {
-                node.setprop(name, value).map_err(|e| anyhow!("Failed to set property: {e:?}"))?;
-            }
-        }
-    }
-
-    for path in additional_paths {
-        if let DtAddition::FromPath(path) = path {
-            fdt.append(cstr!("/fragment@0/__overlay__"), path)?;
-        }
-    }
-    fdt.pack().map_err(|e| anyhow!("Failed to pack DT overlay, {e:?}"))?;
-
-    Ok(fdt)
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn empty_overlays_not_allowed() {
-        let mut buffer = vec![0_u8; VM_DT_OVERLAY_MAX_SIZE];
-        let res = create_device_tree_overlay(&mut buffer, &[]);
-        assert!(res.is_err());
-    }
-
-    #[test]
-    fn untrusted_prop_test() {
-        let mut buffer = vec![0_u8; VM_DT_OVERLAY_MAX_SIZE];
-        let prop_name = cstr!("XOXO");
-        let prop_val_input = b"OXOX";
-        let fdt = create_device_tree_overlay(
-            &mut buffer,
-            &[DtAddition::AvfUntrustedProp(prop_name, prop_val_input)],
-        )
-        .unwrap();
-
-        let prop_value_dt = fdt
-            .node(cstr!("/fragment@0/__overlay__/avf/untrusted"))
-            .unwrap()
-            .expect("/avf/untrusted node doesn't exist")
-            .getprop(prop_name)
-            .unwrap()
-            .expect("Prop not found!");
-        assert_eq!(prop_value_dt, prop_val_input, "Unexpected property value");
-    }
-}
diff --git a/virtualizationmanager/src/main.rs b/virtualizationmanager/src/main.rs
index b2a734a..2e542c3 100644
--- a/virtualizationmanager/src/main.rs
+++ b/virtualizationmanager/src/main.rs
@@ -19,8 +19,8 @@
 mod composite;
 mod crosvm;
 mod debug_config;
-mod dt_overlay;
 mod payload;
+mod reference_dt;
 mod selinux;
 
 use crate::aidl::{GLOBAL_SERVICE, VirtualizationService};
diff --git a/virtualizationmanager/src/reference_dt.rs b/virtualizationmanager/src/reference_dt.rs
new file mode 100644
index 0000000..797ee3c
--- /dev/null
+++ b/virtualizationmanager/src/reference_dt.rs
@@ -0,0 +1,93 @@
+// Copyright 2024, 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.
+
+//! Functions for VM reference DT
+
+use anyhow::{anyhow, Result};
+use cstr::cstr;
+use fsfdt::FsFdt;
+use libfdt::Fdt;
+use std::fs;
+use std::fs::File;
+use std::path::Path;
+
+const VM_REFERENCE_DT_ON_HOST_PATH: &str = "/proc/device-tree/avf/reference";
+const VM_REFERENCE_DT_NAME: &str = "vm_reference_dt.dtbo";
+const VM_REFERENCE_DT_MAX_SIZE: usize = 2000;
+
+// Parses to VM reference if exists.
+// TODO(b/318431695): Allow to parse from custom VM reference DT
+pub(crate) fn parse_reference_dt(out_dir: &Path) -> Result<Option<File>> {
+    parse_reference_dt_internal(
+        Path::new(VM_REFERENCE_DT_ON_HOST_PATH),
+        &out_dir.join(VM_REFERENCE_DT_NAME),
+    )
+}
+
+fn parse_reference_dt_internal(dir_path: &Path, fdt_path: &Path) -> Result<Option<File>> {
+    if !dir_path.exists() || fs::read_dir(dir_path)?.next().is_none() {
+        return Ok(None);
+    }
+
+    let mut data = vec![0_u8; VM_REFERENCE_DT_MAX_SIZE];
+
+    let fdt = Fdt::create_empty_tree(&mut data)
+        .map_err(|e| anyhow!("Failed to create an empty DT, {e:?}"))?;
+    let mut root = fdt.root_mut().map_err(|e| anyhow!("Failed to find the DT root, {e:?}"))?;
+    let mut fragment = root
+        .add_subnode(cstr!("fragment@0"))
+        .map_err(|e| anyhow!("Failed to create the fragment@0, {e:?}"))?;
+    fragment
+        .setprop(cstr!("target-path"), b"/\0")
+        .map_err(|e| anyhow!("Failed to set target-path, {e:?}"))?;
+    fragment
+        .add_subnode(cstr!("__overlay__"))
+        .map_err(|e| anyhow!("Failed to create the __overlay__, {e:?}"))?;
+
+    fdt.append(cstr!("/fragment@0/__overlay__"), dir_path)?;
+
+    fdt.pack().map_err(|e| anyhow!("Failed to pack VM reference DT, {e:?}"))?;
+    fs::write(fdt_path, fdt.as_slice())?;
+
+    Ok(Some(File::open(fdt_path)?))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_reference_dt_from_empty_dir() {
+        let empty_dir = tempfile::TempDir::new().unwrap();
+        let test_dir = tempfile::TempDir::new().unwrap();
+
+        let empty_dir_path = empty_dir.path();
+        let fdt_path = test_dir.path().join("test.dtb");
+
+        let fdt_file = parse_reference_dt_internal(empty_dir_path, &fdt_path).unwrap();
+
+        assert!(fdt_file.is_none());
+    }
+
+    #[test]
+    fn test_parse_reference_dt_from_empty_reference() {
+        let fdt_file = parse_reference_dt_internal(
+            Path::new("/this/path/would/not/exists"),
+            Path::new("test.dtb"),
+        )
+        .unwrap();
+
+        assert!(fdt_file.is_none());
+    }
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
index 92a5812..7962bc3 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
@@ -24,6 +24,7 @@
 interface IVirtualizationService {
     const String FEATURE_DICE_CHANGES = "com.android.kvm.DICE_CHANGES";
     const String FEATURE_MULTI_TENANT = "com.android.kvm.MULTI_TENANT";
+    const String FEATURE_REMOTE_ATTESTATION = "com.android.kvm.REMOTE_ATTESTATION";
     const String FEATURE_VENDOR_MODULES = "com.android.kvm.VENDOR_MODULES";
 
     /**
@@ -73,4 +74,10 @@
 
     /** Returns whether given feature is enabled. */
     boolean isFeatureEnabled(in String feature);
+
+    /**
+     * Provisions a key pair for the VM attestation testing, a fake certificate will be
+     * associated to the fake key pair when the VM requests attestation in testing mode.
+     */
+    void enableTestAttestation();
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
index dd94526..abfc45a 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
@@ -61,10 +61,20 @@
      *                     attested is owned by this app.
      *                     The uniqueness of the UID ensures that no two VMs owned by different apps
      *                     are able to correlate keys.
+     * @param testMode Whether the request is for testing purposes.
      * @return A sequence of DER-encoded X.509 certificates that make up the attestation
      *         key's certificate chain. The attestation key is provided in the CSR.
      */
-    Certificate[] requestAttestation(in byte[] csr, int requesterUid);
+    Certificate[] requestAttestation(in byte[] csr, int requesterUid, in boolean testMode);
+
+    /**
+     * Provisions a key pair for the VM attestation testing, a fake certificate will be
+     * associated to the fake key pair when the VM requests attestation in testing mode.
+     *
+     * The provisioned key pair will be used in the subsequent call to {@link #requestAttestation}
+     * with testMode set to true.
+     */
+    void enableTestAttestation();
 
     /**
      * Get a list of assignable devices.
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
index cf91302..6806a5c 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -51,10 +51,11 @@
      * Requests a certificate chain for the provided certificate signing request (CSR).
      *
      * @param csr The certificate signing request.
+     * @param testMode Whether the request is for test purposes.
      * @return A sequence of DER-encoded X.509 certificates that make up the attestation
      *         key's certificate chain. The attestation key is provided in the CSR.
      */
-    Certificate[] requestAttestation(in byte[] csr);
+    Certificate[] requestAttestation(in byte[] csr, in boolean testMode);
 
     /**
      * Request connection to Secretkeeper. This is used by pVM to store Anti-Rollback protected
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index a1a1fb9..d0c5d4a 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -16,7 +16,8 @@
 
 use crate::{get_calling_pid, get_calling_uid, REMOTELY_PROVISIONED_COMPONENT_SERVICE_NAME};
 use crate::atom::{forward_vm_booted_atom, forward_vm_creation_atom, forward_vm_exited_atom};
-use crate::rkpvm::request_attestation;
+use crate::rkpvm::{request_attestation, generate_ecdsa_p256_key_pair};
+use crate::remote_provisioning;
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
 use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::Certificate::Certificate;
 use android_system_virtualizationservice::{
@@ -38,6 +39,7 @@
 use anyhow::{anyhow, ensure, Context, Result};
 use avflog::LogResult;
 use binder::{self, wait_for_interface, BinderFeatures, ExceptionCode, Interface, LazyServiceGuard, Status, Strong, IntoBinderResult};
+use service_vm_comm::Response;
 use lazy_static::lazy_static;
 use libc::VMADDR_CID_HOST;
 use log::{error, info, warn};
@@ -73,7 +75,73 @@
 
 const CHUNK_RECV_MAX_LEN: usize = 1024;
 
+/// The fake certificate is used for testing only when a client VM requests attestation in test
+/// mode, it is a single certificate extracted on an unregistered device for testing.
+/// Here is the snapshot of the certificate:
+///
+/// ```
+/// Certificate:
+/// Data:
+/// Version: 3 (0x2)
+/// Serial Number:
+///     59:ae:50:98:95:e1:34:25:f1:21:93:c0:4c:e5:24:66
+/// Signature Algorithm: ecdsa-with-SHA256
+/// Issuer: CN = Droid Unregistered Device CA, O = Google Test LLC
+/// Validity
+///     Not Before: Feb  5 14:39:39 2024 GMT
+///     Not After : Feb 14 14:39:39 2024 GMT
+/// Subject: CN = 59ae509895e13425f12193c04ce52466, O = TEE
+/// Subject Public Key Info:
+///     Public Key Algorithm: id-ecPublicKey
+///         Public-Key: (256 bit)
+///         pub:
+///             04:30:32:cd:95:12:b0:71:8b:b7:14:44:26:58:d5:
+///             82:8c:25:55:2c:6d:ef:98:e3:4f:88:d0:74:82:09:
+///             3e:8d:6c:f0:f2:18:d5:83:0e:0d:f2:ce:c5:15:38:
+///             e5:6a:e6:4d:4d:95:15:b7:24:e7:cb:4b:63:42:21:
+///             bc:36:c6:0a:d8
+///         ASN1 OID: prime256v1
+///         NIST CURVE: P-256
+/// X509v3 extensions:
+///  ...
+/// ```
+const FAKE_CERTIFICATE_FOR_TESTING: &[u8] = &[
+    0x30, 0x82, 0x01, 0xee, 0x30, 0x82, 0x01, 0x94, 0xa0, 0x03, 0x02, 0x01, 0x02, 0x02, 0x10, 0x59,
+    0xae, 0x50, 0x98, 0x95, 0xe1, 0x34, 0x25, 0xf1, 0x21, 0x93, 0xc0, 0x4c, 0xe5, 0x24, 0x66, 0x30,
+    0x0a, 0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x04, 0x03, 0x02, 0x30, 0x41, 0x31, 0x25, 0x30,
+    0x23, 0x06, 0x03, 0x55, 0x04, 0x03, 0x13, 0x1c, 0x44, 0x72, 0x6f, 0x69, 0x64, 0x20, 0x55, 0x6e,
+    0x72, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x65, 0x64, 0x20, 0x44, 0x65, 0x76, 0x69, 0x63,
+    0x65, 0x20, 0x43, 0x41, 0x31, 0x18, 0x30, 0x16, 0x06, 0x03, 0x55, 0x04, 0x0a, 0x13, 0x0f, 0x47,
+    0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x20, 0x54, 0x65, 0x73, 0x74, 0x20, 0x4c, 0x4c, 0x43, 0x30, 0x1e,
+    0x17, 0x0d, 0x32, 0x34, 0x30, 0x32, 0x30, 0x35, 0x31, 0x34, 0x33, 0x39, 0x33, 0x39, 0x5a, 0x17,
+    0x0d, 0x32, 0x34, 0x30, 0x32, 0x31, 0x34, 0x31, 0x34, 0x33, 0x39, 0x33, 0x39, 0x5a, 0x30, 0x39,
+    0x31, 0x29, 0x30, 0x27, 0x06, 0x03, 0x55, 0x04, 0x03, 0x13, 0x20, 0x35, 0x39, 0x61, 0x65, 0x35,
+    0x30, 0x39, 0x38, 0x39, 0x35, 0x65, 0x31, 0x33, 0x34, 0x32, 0x35, 0x66, 0x31, 0x32, 0x31, 0x39,
+    0x33, 0x63, 0x30, 0x34, 0x63, 0x65, 0x35, 0x32, 0x34, 0x36, 0x36, 0x31, 0x0c, 0x30, 0x0a, 0x06,
+    0x03, 0x55, 0x04, 0x0a, 0x13, 0x03, 0x54, 0x45, 0x45, 0x30, 0x59, 0x30, 0x13, 0x06, 0x07, 0x2a,
+    0x86, 0x48, 0xce, 0x3d, 0x02, 0x01, 0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07,
+    0x03, 0x42, 0x00, 0x04, 0x30, 0x32, 0xcd, 0x95, 0x12, 0xb0, 0x71, 0x8b, 0xb7, 0x14, 0x44, 0x26,
+    0x58, 0xd5, 0x82, 0x8c, 0x25, 0x55, 0x2c, 0x6d, 0xef, 0x98, 0xe3, 0x4f, 0x88, 0xd0, 0x74, 0x82,
+    0x09, 0x3e, 0x8d, 0x6c, 0xf0, 0xf2, 0x18, 0xd5, 0x83, 0x0e, 0x0d, 0xf2, 0xce, 0xc5, 0x15, 0x38,
+    0xe5, 0x6a, 0xe6, 0x4d, 0x4d, 0x95, 0x15, 0xb7, 0x24, 0xe7, 0xcb, 0x4b, 0x63, 0x42, 0x21, 0xbc,
+    0x36, 0xc6, 0x0a, 0xd8, 0xa3, 0x76, 0x30, 0x74, 0x30, 0x1d, 0x06, 0x03, 0x55, 0x1d, 0x0e, 0x04,
+    0x16, 0x04, 0x14, 0x39, 0x81, 0x41, 0x0a, 0xb9, 0xf3, 0xf4, 0x5b, 0x75, 0x97, 0x4a, 0x46, 0xd6,
+    0x30, 0x9e, 0x1d, 0x7a, 0x3b, 0xec, 0xa8, 0x30, 0x1f, 0x06, 0x03, 0x55, 0x1d, 0x23, 0x04, 0x18,
+    0x30, 0x16, 0x80, 0x14, 0x82, 0xbd, 0x00, 0xde, 0xcb, 0xc5, 0xe7, 0x72, 0x87, 0x3d, 0x1c, 0x0a,
+    0x1e, 0x78, 0x4f, 0xf5, 0xd3, 0xc1, 0x3e, 0xb8, 0x30, 0x0f, 0x06, 0x03, 0x55, 0x1d, 0x13, 0x01,
+    0x01, 0xff, 0x04, 0x05, 0x30, 0x03, 0x01, 0x01, 0xff, 0x30, 0x0e, 0x06, 0x03, 0x55, 0x1d, 0x0f,
+    0x01, 0x01, 0xff, 0x04, 0x04, 0x03, 0x02, 0x02, 0x04, 0x30, 0x11, 0x06, 0x0a, 0x2b, 0x06, 0x01,
+    0x04, 0x01, 0xd6, 0x79, 0x02, 0x01, 0x1e, 0x04, 0x03, 0xa1, 0x01, 0x08, 0x30, 0x0a, 0x06, 0x08,
+    0x2a, 0x86, 0x48, 0xce, 0x3d, 0x04, 0x03, 0x02, 0x03, 0x48, 0x00, 0x30, 0x45, 0x02, 0x21, 0x00,
+    0xae, 0xd8, 0x40, 0x9e, 0x37, 0x3e, 0x5c, 0x9c, 0xe2, 0x93, 0x3d, 0x8c, 0xf7, 0x05, 0x10, 0xe7,
+    0xd1, 0x2b, 0x87, 0x8a, 0xee, 0xd6, 0x1e, 0x6c, 0x3b, 0xd2, 0x91, 0x3e, 0xa5, 0xdf, 0x91, 0x20,
+    0x02, 0x20, 0x7f, 0x0f, 0x29, 0x54, 0x60, 0x80, 0x07, 0x50, 0x5f, 0x56, 0x6b, 0x9f, 0xe0, 0x94,
+    0xb4, 0x3f, 0x3b, 0x0f, 0x61, 0xa0, 0x33, 0x40, 0xe6, 0x1a, 0x42, 0xda, 0x4b, 0xa4, 0xfd, 0x92,
+    0xb9, 0x0f,
+];
+
 lazy_static! {
+    static ref FAKE_PROVISIONED_KEY_BLOB_FOR_TESTING: Mutex<Option<Vec<u8>>> = Mutex::new(None);
     static ref VFIO_SERVICE: Strong<dyn IVfioHandler> =
         wait_for_interface(<BpVfioHandler as IVfioHandler>::get_descriptor())
             .expect("Could not connect to VfioHandler");
@@ -169,10 +237,41 @@
         Ok(cids)
     }
 
+    fn enableTestAttestation(&self) -> binder::Result<()> {
+        check_manage_access()?;
+        check_use_custom_virtual_machine()?;
+        if !cfg!(remote_attestation) {
+            return Err(Status::new_exception_str(
+                ExceptionCode::UNSUPPORTED_OPERATION,
+                Some(
+                    "enableTestAttestation is not supported with the remote_attestation \
+                     feature disabled",
+                ),
+            ))
+            .with_log();
+        }
+        let res = generate_ecdsa_p256_key_pair()
+            .context("Failed to generate ECDSA P-256 key pair for testing")
+            .with_log()
+            .or_service_specific_exception(-1)?;
+        match res {
+            Response::GenerateEcdsaP256KeyPair(key_pair) => {
+                FAKE_PROVISIONED_KEY_BLOB_FOR_TESTING
+                    .lock()
+                    .unwrap()
+                    .replace(key_pair.key_blob.to_vec());
+                Ok(())
+            }
+            _ => Err(remote_provisioning::to_service_specific_error(res)),
+        }
+        .with_log()
+    }
+
     fn requestAttestation(
         &self,
         csr: &[u8],
         requester_uid: i32,
+        test_mode: bool,
     ) -> binder::Result<Vec<Certificate>> {
         check_manage_access()?;
         if !cfg!(remote_attestation) {
@@ -186,14 +285,31 @@
             .with_log();
         }
         info!("Received csr. Requestting attestation...");
-        let attestation_key = get_rkpd_attestation_key(
-            REMOTELY_PROVISIONED_COMPONENT_SERVICE_NAME,
-            requester_uid as u32,
-        )
-        .context("Failed to retrieve the remotely provisioned keys")
-        .with_log()
-        .or_service_specific_exception(-1)?;
-        let mut certificate_chain = split_x509_certificate_chain(&attestation_key.encodedCertChain)
+        let (key_blob, certificate_chain) = if test_mode {
+            check_use_custom_virtual_machine()?;
+            info!("Using the fake key blob for testing...");
+            (
+                FAKE_PROVISIONED_KEY_BLOB_FOR_TESTING
+                    .lock()
+                    .unwrap()
+                    .clone()
+                    .ok_or_else(|| anyhow!("No key blob for testing"))
+                    .with_log()
+                    .or_service_specific_exception(-1)?,
+                FAKE_CERTIFICATE_FOR_TESTING.to_vec(),
+            )
+        } else {
+            info!("Retrieving the remotely provisioned keys from RKPD...");
+            let attestation_key = get_rkpd_attestation_key(
+                REMOTELY_PROVISIONED_COMPONENT_SERVICE_NAME,
+                requester_uid as u32,
+            )
+            .context("Failed to retrieve the remotely provisioned keys")
+            .with_log()
+            .or_service_specific_exception(-1)?;
+            (attestation_key.keyBlob, attestation_key.encodedCertChain)
+        };
+        let mut certificate_chain = split_x509_certificate_chain(&certificate_chain)
             .context("Failed to split the remotely provisioned certificate chain")
             .with_log()
             .or_service_specific_exception(-1)?;
@@ -206,7 +322,7 @@
         }
         let certificate = request_attestation(
             csr.to_vec(),
-            attestation_key.keyBlob,
+            key_blob,
             certificate_chain[0].encodedCertificate.clone(),
         )
         .context("Failed to request attestation")
diff --git a/virtualizationservice/src/remote_provisioning.rs b/virtualizationservice/src/remote_provisioning.rs
index 40f54db..c2c04df 100644
--- a/virtualizationservice/src/remote_provisioning.rs
+++ b/virtualizationservice/src/remote_provisioning.rs
@@ -145,7 +145,7 @@
     }
 }
 
-fn to_service_specific_error(response: Response) -> Status {
+pub(crate) fn to_service_specific_error(response: Response) -> Status {
     match response {
         Response::Err(e) => match e {
             RequestProcessingError::InvalidMac => {
diff --git a/vm_payload/include-restricted/vm_payload_restricted.h b/vm_payload/include-restricted/vm_payload_restricted.h
index 15c37ed..d7324a8 100644
--- a/vm_payload/include-restricted/vm_payload_restricted.h
+++ b/vm_payload/include-restricted/vm_payload_restricted.h
@@ -55,4 +55,25 @@
  */
 size_t AVmPayload_getDiceAttestationCdi(void* _Nullable data, size_t size);
 
+/**
+ * Requests attestation for the VM for testing only.
+ *
+ * This function is only for testing and will not return a real RKP server backed
+ * certificate chain.
+ *
+ * Prior to calling this function, the caller must provision a key pair to be used in
+ * this function with `VirtualMachineManager#enableTestAttestation`.
+ *
+ * \param challenge A pointer to the challenge buffer.
+ * \param challenge_size size of the challenge. The maximum supported challenge size is
+ *          64 bytes. The status ATTESTATION_ERROR_INVALID_CHALLENGE will be returned if
+ *          an invalid challenge is passed.
+ * \param result The remote attestation result will be filled here if the attestation
+ *               succeeds. The result remains valid until it is freed with
+ *              `AVmPayload_freeAttestationResult`.
+ */
+attestation_status_t AVmPayload_requestAttestationForTesting(
+        const void* _Nonnull challenge, size_t challenge_size,
+        struct AVmAttestationResult* _Nullable* _Nonnull result) __INTRODUCED_IN(__ANDROID_API_V__);
+
 __END_DECLS
diff --git a/vm_payload/include/vm_payload.h b/vm_payload/include/vm_payload.h
index 3483e1d..af755c9 100644
--- a/vm_payload/include/vm_payload.h
+++ b/vm_payload/include/vm_payload.h
@@ -211,7 +211,7 @@
  * If `size` is smaller than the total size of the signature, the signature will be
  * truncated to this `size`.
  *
- * \return The total size of the signature.
+ * \return The size of the signature, or the size needed if the supplied buffer is too small.
  *
  * [RFC 6979]: https://datatracker.ietf.org/doc/html/rfc6979
  */
diff --git a/vm_payload/libvm_payload.map.txt b/vm_payload/libvm_payload.map.txt
index 975a5a3..caf8f84 100644
--- a/vm_payload/libvm_payload.map.txt
+++ b/vm_payload/libvm_payload.map.txt
@@ -8,6 +8,7 @@
     AVmPayload_getApkContentsPath;       # systemapi introduced=UpsideDownCake
     AVmPayload_getEncryptedStoragePath;  # systemapi introduced=UpsideDownCake
     AVmPayload_requestAttestation;       # systemapi introduced=VanillaIceCream
+    AVmPayload_requestAttestationForTesting; # systemapi introduced=VanillaIceCream
     AVmAttestationResult_getPrivateKey;  # systemapi introduced=VanillaIceCream
     AVmAttestationResult_sign;           # systemapi introduced=VanillaIceCream
     AVmAttestationResult_free;           # systemapi introduced=VanillaIceCream
diff --git a/vm_payload/src/lib.rs b/vm_payload/src/lib.rs
index 7978059..6188b21 100644
--- a/vm_payload/src/lib.rs
+++ b/vm_payload/src/lib.rs
@@ -39,6 +39,9 @@
 };
 use vm_payload_status_bindgen::attestation_status_t;
 
+/// Maximum size of an ECDSA signature for EC P-256 key is 72 bytes.
+const MAX_ECDSA_P256_SIGNATURE_SIZE: usize = 72;
+
 lazy_static! {
     static ref VM_APK_CONTENTS_PATH_C: CString =
         CString::new(VM_APK_CONTENTS_PATH).expect("CString::new failed");
@@ -273,9 +276,6 @@
 /// Behavior is undefined if any of the following conditions are violated:
 ///
 /// * `challenge` must be [valid] for reads of `challenge_size` bytes.
-/// * `res` must be [valid] to write the attestation result.
-/// * The region of memory beginning at `challenge` with `challenge_size` bytes must not
-///  overlap with the region of memory `res` points to.
 ///
 /// [valid]: ptr#safety
 #[no_mangle]
@@ -284,6 +284,60 @@
     challenge_size: usize,
     res: &mut *mut AttestationResult,
 ) -> attestation_status_t {
+    // SAFETY: The caller guarantees that `challenge` is valid for reads and `res` is valid
+    // for writes.
+    unsafe {
+        request_attestation(
+            challenge,
+            challenge_size,
+            false, // test_mode
+            res,
+        )
+    }
+}
+
+/// Requests the remote attestation of the client VM for testing.
+///
+/// # Safety
+///
+/// Behavior is undefined if any of the following conditions are violated:
+///
+/// * `challenge` must be [valid] for reads of `challenge_size` bytes.
+///
+/// [valid]: ptr#safety
+#[no_mangle]
+pub unsafe extern "C" fn AVmPayload_requestAttestationForTesting(
+    challenge: *const u8,
+    challenge_size: usize,
+    res: &mut *mut AttestationResult,
+) -> attestation_status_t {
+    // SAFETY: The caller guarantees that `challenge` is valid for reads and `res` is valid
+    // for writes.
+    unsafe {
+        request_attestation(
+            challenge,
+            challenge_size,
+            true, // test_mode
+            res,
+        )
+    }
+}
+
+/// Requests the remote attestation of the client VM.
+///
+/// # Safety
+///
+/// Behavior is undefined if any of the following conditions are violated:
+///
+/// * `challenge` must be [valid] for reads of `challenge_size` bytes.
+///
+/// [valid]: ptr#safety
+unsafe fn request_attestation(
+    challenge: *const u8,
+    challenge_size: usize,
+    test_mode: bool,
+    res: &mut *mut AttestationResult,
+) -> attestation_status_t {
     initialize_logging();
     const MAX_CHALLENGE_SIZE: usize = 64;
     if challenge_size > MAX_CHALLENGE_SIZE {
@@ -297,7 +351,7 @@
         unsafe { std::slice::from_raw_parts(challenge, challenge_size) }
     };
     let service = unwrap_or_abort(get_vm_payload_service());
-    match service.requestAttestation(challenge) {
+    match service.requestAttestation(challenge, test_mode) {
         Ok(attestation_res) => {
             *res = Box::into_raw(Box::new(attestation_res));
             attestation_status_t::ATTESTATION_OK
@@ -400,27 +454,36 @@
     data: *mut u8,
     size: usize,
 ) -> usize {
+    // A DER-encoded ECDSA signature can have varying sizes even with the same EC Key and message,
+    // due to the encoding of the random values r and s that are part of the signature.
+    if size == 0 {
+        return MAX_ECDSA_P256_SIGNATURE_SIZE;
+    }
     if message_size == 0 {
         panic!("Message to be signed must not be empty.")
     }
     // SAFETY: See the requirements on `message` above.
     let message = unsafe { std::slice::from_raw_parts(message, message_size) };
     let signature = unwrap_or_abort(try_ecdsa_sign(message, &res.privateKey));
-    if size != 0 {
-        let data = NonNull::new(data).expect("data must not be null when size > 0");
-        // SAFETY: See the requirements on `data` above. The number of bytes copied doesn't exceed
-        // the length of either buffer, and the caller ensures that `signature` cannot overlap
-        // `data`. We allow data to be null, which is never valid, but only if size == 0
-        // which is checked above.
-        unsafe {
-            ptr::copy_nonoverlapping(
-                signature.as_ptr(),
-                data.as_ptr(),
-                std::cmp::min(signature.len(), size),
-            )
-        };
+    let data = NonNull::new(data).expect("data must not be null when size > 0");
+    // SAFETY: See the requirements on `data` above. The number of bytes copied doesn't exceed
+    // the length of either buffer, and the caller ensures that `signature` cannot overlap
+    // `data`. We allow data to be null, which is never valid, but only if size == 0
+    // which is checked above.
+    unsafe {
+        ptr::copy_nonoverlapping(
+            signature.as_ptr(),
+            data.as_ptr(),
+            usize::min(signature.len(), size),
+        )
+    };
+    if size < signature.len() {
+        // If the buffer is too small, return the maximum size of the signature to allow the caller
+        // to allocate a buffer large enough to call this function again.
+        MAX_ECDSA_P256_SIGNATURE_SIZE
+    } else {
+        signature.len()
     }
-    signature.len()
 }
 
 fn try_ecdsa_sign(message: &[u8], der_encoded_ec_private_key: &[u8]) -> Result<Vec<u8>> {