[client-vm] Build client VM CSR and sign the CSR with two keys

This cl builds the CSR that a client VM sends to the RKP VM for
remote attestation and adjusted the API accordingly as discussed
in the doc go/pvm-remote-attestation

The CSR payload is signed with both the CDI_Leaf_Priv of the
client VM's DICE chain and the attestation key. RKP VM should
verify the signature later with the CDI_Leaf_Pub extracted
from the same DICE chain in the CSR and the attestation public
key.

The new unit tests are added to config at cl/577763874.

Bug: 303807447
Test: run ServiceVmClientTestApp
Test: atest libservice_vm_comm.test
Test: atest microdroid_manager_test
Change-Id: Ic2c09e7339d9981edda028e2694fa551c911a274
diff --git a/service_vm/comm/Android.bp b/service_vm/comm/Android.bp
index 3a18052..6e05587 100644
--- a/service_vm/comm/Android.bp
+++ b/service_vm/comm/Android.bp
@@ -43,3 +43,31 @@
         "std",
     ],
 }
+
+rust_defaults {
+    name: "libservice_vm_comm_test_defaults",
+    crate_name: "diced_open_dice_test",
+    srcs: ["tests/*.rs"],
+    test_suites: ["general-tests"],
+    prefer_rlib: true,
+    rustlibs: [
+        "libdiced_sample_inputs",
+        "libdiced_open_dice",
+    ],
+}
+
+rust_test {
+    name: "libservice_vm_comm.test",
+    defaults: ["libservice_vm_comm_test_defaults"],
+    rustlibs: [
+        "libservice_vm_comm",
+    ],
+}
+
+rust_test {
+    name: "libservice_vm_comm_nostd.test",
+    defaults: ["libservice_vm_comm_test_defaults"],
+    rustlibs: [
+        "libservice_vm_comm_nostd",
+    ],
+}
diff --git a/service_vm/comm/TEST_MAPPING b/service_vm/comm/TEST_MAPPING
new file mode 100644
index 0000000..e677ba2
--- /dev/null
+++ b/service_vm/comm/TEST_MAPPING
@@ -0,0 +1,12 @@
+// When adding or removing tests here, don't forget to amend _all_modules list in
+// wireless/android/busytown/ath_config/configs/prod/avf/tests.gcl
+{
+  "avf-presubmit" : [
+    {
+      "name" : "libservice_vm_comm.test"
+    },
+    {
+      "name" : "libservice_vm_comm_nostd.test"
+    }
+  ]
+}
\ No newline at end of file
diff --git a/service_vm/comm/src/csr.rs b/service_vm/comm/src/csr.rs
new file mode 100644
index 0000000..5e1cbad
--- /dev/null
+++ b/service_vm/comm/src/csr.rs
@@ -0,0 +1,121 @@
+// Copyright 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.
+
+//! This module contains the structs related to the CSR (Certificate Signing Request)
+//! sent from the client VM to the service VM for attestation.
+
+use alloc::vec;
+use alloc::vec::Vec;
+use ciborium::Value;
+use coset::{self, CborSerializable, CoseError};
+
+/// Represents a CSR sent from the client VM to the service VM for attestation.
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct Csr {
+    /// The DICE certificate chain of the client VM.
+    pub dice_cert_chain: Vec<u8>,
+
+    /// The signed CSR payload in COSE_Sign structure, which includes two signatures:
+    /// - one by CDI_Leaf_Priv of the client VM's DICE chain,
+    /// - another by the private key corresponding to the public key.
+    pub signed_csr_payload: Vec<u8>,
+}
+
+impl Csr {
+    /// Serializes this object to a CBOR-encoded vector.
+    pub fn into_cbor_vec(self) -> coset::Result<Vec<u8>> {
+        let value = Value::Array(vec![
+            Value::Bytes(self.dice_cert_chain),
+            Value::Bytes(self.signed_csr_payload),
+        ]);
+        value.to_vec()
+    }
+
+    /// Creates an object instance from the provided CBOR-encoded slice.
+    pub fn from_cbor_slice(data: &[u8]) -> coset::Result<Self> {
+        let value = Value::from_slice(data)?;
+        let Value::Array(mut arr) = value else {
+            return Err(CoseError::UnexpectedItem(cbor_value_type(&value), "array"));
+        };
+        if arr.len() != 2 {
+            return Err(CoseError::UnexpectedItem("array", "array with 2 items"));
+        }
+        Ok(Self {
+            signed_csr_payload: try_as_bytes(arr.remove(1))?,
+            dice_cert_chain: try_as_bytes(arr.remove(0))?,
+        })
+    }
+}
+
+/// Represents the data to be signed and sent from the client VM to the service VM
+/// for attestation.
+///
+/// It will be signed by both CDI_Leaf_Priv of the client VM's DICE chain and
+/// the private key corresponding to the public key to be attested.
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct CsrPayload {
+    /// COSE_Key encoded EC P-256 public key to be attested.
+    pub public_key: Vec<u8>,
+
+    /// A random array with a length between 0 and 64.
+    /// It will be included in the certificate chain in the attestation result,
+    /// serving as proof of the freshness of the result.
+    pub challenge: Vec<u8>,
+}
+
+impl CsrPayload {
+    /// Serializes this object to a CBOR-encoded vector.
+    pub fn into_cbor_vec(self) -> coset::Result<Vec<u8>> {
+        let value = Value::Array(vec![Value::Bytes(self.public_key), Value::Bytes(self.challenge)]);
+        value.to_vec()
+    }
+
+    /// Creates an object instance from the provided CBOR-encoded slice.
+    pub fn from_cbor_slice(data: &[u8]) -> coset::Result<Self> {
+        let value = Value::from_slice(data)?;
+        let Value::Array(mut arr) = value else {
+            return Err(CoseError::UnexpectedItem(cbor_value_type(&value), "array"));
+        };
+        if arr.len() != 2 {
+            return Err(CoseError::UnexpectedItem("array", "array with 2 items"));
+        }
+        Ok(Self {
+            challenge: try_as_bytes(arr.remove(1))?,
+            public_key: try_as_bytes(arr.remove(0))?,
+        })
+    }
+}
+
+fn try_as_bytes(v: Value) -> coset::Result<Vec<u8>> {
+    if let Value::Bytes(data) = v {
+        Ok(data)
+    } else {
+        Err(CoseError::UnexpectedItem(cbor_value_type(&v), "bytes"))
+    }
+}
+
+fn cbor_value_type(v: &Value) -> &'static str {
+    match v {
+        Value::Integer(_) => "int",
+        Value::Bytes(_) => "bstr",
+        Value::Float(_) => "float",
+        Value::Text(_) => "tstr",
+        Value::Bool(_) => "bool",
+        Value::Null => "nul",
+        Value::Tag(_, _) => "tag",
+        Value::Array(_) => "array",
+        Value::Map(_) => "map",
+        _ => "other",
+    }
+}
diff --git a/service_vm/comm/src/lib.rs b/service_vm/comm/src/lib.rs
index d8f7bd7..0818f24 100644
--- a/service_vm/comm/src/lib.rs
+++ b/service_vm/comm/src/lib.rs
@@ -19,9 +19,11 @@
 
 extern crate alloc;
 
+mod csr;
 mod message;
 mod vsock;
 
+pub use csr::{Csr, CsrPayload};
 pub use message::{
     EcdsaP256KeyPair, GenerateCertificateRequestParams, Request, RequestProcessingError, Response,
     ServiceVmRequest,
diff --git a/service_vm/comm/tests/api_test.rs b/service_vm/comm/tests/api_test.rs
new file mode 100644
index 0000000..44a3ef9
--- /dev/null
+++ b/service_vm/comm/tests/api_test.rs
@@ -0,0 +1,49 @@
+/*
+ * 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.
+ */
+
+use diced_open_dice::DiceArtifacts;
+use service_vm_comm::{Csr, CsrPayload};
+
+/// The following test data are generated with urandom
+const DATA1: [u8; 32] = [
+    0x8b, 0x09, 0xc0, 0x7e, 0x20, 0x3c, 0xa2, 0x11, 0x7e, 0x7f, 0x0b, 0xdd, 0x2b, 0x68, 0x98, 0xb0,
+    0x2b, 0x34, 0xb5, 0x63, 0x39, 0x01, 0x90, 0x06, 0xaf, 0x5f, 0xdd, 0xb7, 0x81, 0xca, 0xc7, 0x46,
+];
+const DATA2: [u8; 16] = [
+    0x6c, 0xb9, 0x39, 0x86, 0x9b, 0x2f, 0x12, 0xd8, 0x45, 0x92, 0x57, 0x44, 0x65, 0xce, 0x94, 0x63,
+];
+
+#[test]
+fn csr_payload_cbor_serialization() {
+    let csr_payload = CsrPayload { public_key: DATA1.to_vec(), challenge: DATA2.to_vec() };
+    let expected_csr_payload = csr_payload.clone();
+    let cbor_vec = csr_payload.into_cbor_vec().unwrap();
+    let deserialized_csr_payload = CsrPayload::from_cbor_slice(&cbor_vec).unwrap();
+
+    assert_eq!(expected_csr_payload, deserialized_csr_payload);
+}
+
+#[test]
+fn csr_cbor_serialization() {
+    let dice_artifacts = diced_sample_inputs::make_sample_bcc_and_cdis().unwrap();
+    let dice_cert_chain = dice_artifacts.bcc().unwrap().to_vec();
+    let csr = Csr { signed_csr_payload: DATA1.to_vec(), dice_cert_chain };
+    let expected_csr = csr.clone();
+    let cbor_vec = csr.into_cbor_vec().unwrap();
+    let deserialized_csr = Csr::from_cbor_slice(&cbor_vec).unwrap();
+
+    assert_eq!(expected_csr, deserialized_csr);
+}
diff --git a/service_vm/test_apk/src/main.rs b/service_vm/test_apk/src/main.rs
index 7d8416f..ba65aca 100644
--- a/service_vm/test_apk/src/main.rs
+++ b/service_vm/test_apk/src/main.rs
@@ -14,10 +14,20 @@
 
 //! Main executable of Service VM client for manual testing.
 
-use anyhow::Result;
+use anyhow::{anyhow, ensure, Result};
 use log::{error, info};
-use std::{ffi::c_void, panic};
-use vm_payload_bindgen::AVmPayload_requestAttestation;
+use std::{
+    ffi::{c_void, CStr},
+    panic,
+    ptr::{self, NonNull},
+    result,
+};
+use vm_payload_bindgen::{
+    attestation_status_t, AVmAttestationResult, AVmAttestationResult_free,
+    AVmAttestationResult_getCertificateAt, AVmAttestationResult_getCertificateCount,
+    AVmAttestationResult_getPrivateKey, AVmAttestationResult_resultToString,
+    AVmAttestationResult_sign, AVmPayload_requestAttestation,
+};
 
 /// Entry point of the Service VM client.
 #[allow(non_snake_case)]
@@ -40,38 +50,180 @@
 
 fn try_main() -> Result<()> {
     info!("Welcome to Service VM Client!");
+
+    let too_big_challenge = &[0u8; 66];
+    let res = AttestationResult::request_attestation(too_big_challenge);
+    ensure!(res.is_err());
+    let status = res.unwrap_err();
+    ensure!(
+        status == attestation_status_t::ATTESTATION_ERROR_INVALID_CHALLENGE,
+        "Unexpected status: {:?}",
+        status
+    );
+    info!("Status: {:?}", status_to_cstr(status));
+
     // 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,
     ];
-    info!("Sending challenge: {:?}", challenge);
-    let certificate = request_attestation(challenge);
-    info!("Certificate: {:?}", certificate);
+    let res = AttestationResult::request_attestation(challenge)
+        .map_err(|e| anyhow!("Unexpected status: {:?}", status_to_cstr(e)))?;
+
+    let cert_chain = res.certificate_chain()?;
+    info!("Attestation result certificateChain = {:?}", cert_chain);
+
+    let private_key = res.private_key()?;
+    info!("Attestation result privateKey = {:?}", private_key);
+
+    let message = b"Hello from Service VM client";
+    info!("Signing message: {:?}", message);
+    let signature = res.sign(message)?;
+    info!("Signature: {:?}", signature);
+
     Ok(())
 }
 
-fn request_attestation(challenge: &[u8]) -> Vec<u8> {
-    // SAFETY: It is safe as we only request the size of the certificate in this call.
-    let certificate_size = unsafe {
-        AVmPayload_requestAttestation(
-            challenge.as_ptr() as *const c_void,
-            challenge.len(),
-            [].as_mut_ptr(),
+#[derive(Debug)]
+struct AttestationResult(NonNull<AVmAttestationResult>);
+
+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_requestAttestation(
+                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)
+    }
+}
+
+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 certificate = vec![0u8; certificate_size];
-    // SAFETY: It is safe as we only write the data into the given buffer within the buffer
-    // size in this call.
-    unsafe {
-        AVmPayload_requestAttestation(
-            challenge.as_ptr() as *const c_void,
-            challenge.len(),
-            certificate.as_mut_ptr() as *mut c_void,
-            certificate.len(),
-        );
+    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(),
+        )
     };
-    certificate
+    ensure!(size == signature.len());
+    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.
+    unsafe { CStr::from_ptr(message) }
 }