Create a Rust wrapper for vm_payload

And use it in place of directly calling the bindgen-generated
interface in our current clients.

Bug: 340857915
Test: atest VmAttestationTestApp
  composd_cmd test-compile
Change-Id: I51f1c1ab6a4dce09d9160731aacd83ebb9c0ce07
diff --git a/vm_payload/wrapper/attestation.rs b/vm_payload/wrapper/attestation.rs
new file mode 100644
index 0000000..e0055d5
--- /dev/null
+++ b/vm_payload/wrapper/attestation.rs
@@ -0,0 +1,288 @@
+/*
+ * 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.
+ */
+
+use std::error::Error;
+use std::ffi::{c_void, CStr};
+use std::fmt::{self, Display};
+use std::iter::FusedIterator;
+use std::ptr::{self, NonNull};
+
+use vm_payload_bindgen::{
+    AVmAttestationResult, AVmAttestationResult_free, AVmAttestationResult_getCertificateAt,
+    AVmAttestationResult_getCertificateCount, AVmAttestationResult_getPrivateKey,
+    AVmAttestationResult_sign, AVmAttestationStatus, AVmAttestationStatus_toString,
+    AVmPayload_requestAttestation, AVmPayload_requestAttestationForTesting,
+};
+
+/// Holds the result of a successful Virtual Machine attestation request.
+/// See [`request_attestation`].
+#[derive(Debug)]
+pub struct AttestationResult {
+    result: NonNull<AVmAttestationResult>,
+}
+
+/// Error type that can be returned from an unsuccessful Virtual Machine attestation request.
+/// See [`request_attestation`].
+#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
+pub enum AttestationError {
+    /// The challenge size was not between 0 and 64 bytes (inclusive).
+    InvalidChallenge,
+    /// The attempt to attest the VM failed. A subsequent request may succeed.
+    AttestationFailed,
+    /// VM attestation is not supported in the current environment.
+    AttestationUnsupported,
+}
+
+impl Error for AttestationError {}
+
+impl Display for AttestationError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
+        let status = match self {
+            Self::InvalidChallenge => AVmAttestationStatus::ATTESTATION_ERROR_INVALID_CHALLENGE,
+            Self::AttestationFailed => AVmAttestationStatus::ATTESTATION_ERROR_ATTESTATION_FAILED,
+            Self::AttestationUnsupported => AVmAttestationStatus::ATTESTATION_ERROR_UNSUPPORTED,
+        };
+        // SAFETY: AVmAttestationStatus_toString always returns a non-null pointer to a
+        // nul-terminated C string with static lifetime (which is valid UTF-8).
+        let c_str = unsafe { CStr::from_ptr(AVmAttestationStatus_toString(status)) };
+        let str = c_str.to_str().expect("Invalid UTF-8 for AVmAttestationStatus");
+        f.write_str(str)
+    }
+}
+
+impl Drop for AttestationResult {
+    fn drop(&mut self) {
+        let ptr = self.result.as_ptr();
+
+        // SAFETY: The `result` field is private, and only populated with a successful call to
+        // `AVmPayload_requestAttestation`, and not freed elsewhere.
+        unsafe { AVmAttestationResult_free(ptr) };
+    }
+}
+
+// SAFETY: The API functions that accept the `AVmAttestationResult` pointer are all safe to call
+// from any thread, including `AVmAttestationResult_free` which is called only on drop.
+unsafe impl Send for AttestationResult {}
+
+// SAFETY: There is no interior mutation here; any future functions that might mutate the data would
+// require a non-const pointer and hence need `&mut self` here. The only existing such function is
+// `AVmAttestationResult_free` where we take a mutable reference guaranteeing no other references
+// exist. The raw API functions are safe to call from any thread.
+unsafe impl Sync for AttestationResult {}
+
+/// Requests the remote attestation of this VM.
+///
+/// On success the supplied [`challenge`] will be included in the certificate chain accessible from
+/// the [`AttestationResult`]; this can be used as proof of the freshness of the attestation.
+///
+/// The challenge should be no more than 64 bytes long or the request will fail.
+pub fn request_attestation(challenge: &[u8]) -> Result<AttestationResult, AttestationError> {
+    let mut result: *mut AVmAttestationResult = ptr::null_mut();
+    // SAFETY: 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 result,
+        )
+    };
+    AttestationResult::new(status, result)
+}
+
+/// A variant of [`request_attestation`] used for testing purposes. This should not be used by
+/// normal VMs, and is not available to app owned VMs.
+pub fn request_attestation_for_testing(
+    challenge: &[u8],
+) -> Result<AttestationResult, AttestationError> {
+    let mut result: *mut AVmAttestationResult = ptr::null_mut();
+    // SAFETY: 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 result,
+        )
+    };
+    AttestationResult::new(status, result)
+}
+
+impl AttestationResult {
+    fn new(
+        status: AVmAttestationStatus,
+        result: *mut AVmAttestationResult,
+    ) -> Result<AttestationResult, AttestationError> {
+        match status {
+            AVmAttestationStatus::ATTESTATION_ERROR_INVALID_CHALLENGE => {
+                Err(AttestationError::InvalidChallenge)
+            }
+            AVmAttestationStatus::ATTESTATION_ERROR_ATTESTATION_FAILED => {
+                Err(AttestationError::AttestationFailed)
+            }
+            AVmAttestationStatus::ATTESTATION_ERROR_UNSUPPORTED => {
+                Err(AttestationError::AttestationUnsupported)
+            }
+            AVmAttestationStatus::ATTESTATION_OK => {
+                let result = NonNull::new(result)
+                    .expect("Attestation succeeded but the attestation result is null");
+                Ok(AttestationResult { result })
+            }
+        }
+    }
+
+    fn as_const_ptr(&self) -> *const AVmAttestationResult {
+        self.result.as_ptr().cast_const()
+    }
+
+    /// Returns the attested private key. This is the ECDSA P-256 private key corresponding to the
+    /// public key described by the leaf certificate in the attested
+    /// [certificate chain](AttestationResult::certificate_chain). It is a DER-encoded
+    /// `ECPrivateKey` structure as specified in
+    /// [RFC 5915 s3](https://datatracker.ietf.org/doc/html/rfc5915#section-3).
+    ///
+    /// Note: The [`sign_message`](AttestationResult::sign_message) method allows signing with the
+    /// key without retrieving it.
+    pub fn private_key(&self) -> Vec<u8> {
+        let ptr = self.as_const_ptr();
+
+        let size =
+            // SAFETY: We own the `AVmAttestationResult` pointer, so it is valid. The function
+            // writes no data since we pass a zero size, and null is explicitly allowed for the
+            // destination in that case.
+            unsafe { AVmAttestationResult_getPrivateKey(ptr, ptr::null_mut(), 0) };
+
+        let mut private_key = vec![0u8; size];
+        // SAFETY: We own the `AVmAttestationResult` pointer, so it is valid. The function only
+        // writes within the bounds of `private_key`, which we just allocated so cannot be aliased.
+        let size = unsafe {
+            AVmAttestationResult_getPrivateKey(
+                ptr,
+                private_key.as_mut_ptr() as *mut c_void,
+                private_key.len(),
+            )
+        };
+        assert_eq!(size, private_key.len());
+        private_key
+    }
+
+    /// Signs the given message using the attested private key. The signature uses ECDSA P-256; the
+    /// message is first hashed with SHA-256 and then it is signed with the attested EC P-256
+    /// [private key](AttestationResult::private_key).
+    ///
+    /// The signature is a DER-encoded `ECDSASignature`` structure as described in
+    /// [RFC 6979](https://datatracker.ietf.org/doc/html/rfc6979).
+    pub fn sign_message(&self, message: &[u8]) -> Vec<u8> {
+        let ptr = self.as_const_ptr();
+
+        // SAFETY: We own the `AVmAttestationResult` pointer, so it is valid. The function
+        // writes no data since we pass a zero size, and null is explicitly allowed for the
+        // destination in that case.
+        let size = unsafe {
+            AVmAttestationResult_sign(
+                ptr,
+                message.as_ptr() as *const c_void,
+                message.len(),
+                ptr::null_mut(),
+                0,
+            )
+        };
+
+        let mut signature = vec![0u8; size];
+        // SAFETY: We own the `AVmAttestationResult` pointer, so it is valid. The function only
+        // writes within the bounds of `signature`, which we just allocated so cannot be aliased.
+        let size = unsafe {
+            AVmAttestationResult_sign(
+                ptr,
+                message.as_ptr() as *const c_void,
+                message.len(),
+                signature.as_mut_ptr() as *mut c_void,
+                signature.len(),
+            )
+        };
+        assert!(size <= signature.len());
+        signature.truncate(size);
+        signature
+    }
+
+    /// Returns an iterator over the certificates forming the certificate chain for the VM, and its
+    /// public key, obtained by the attestation process.
+    ///
+    /// The certificate chain consists of a sequence of DER-encoded X.509 certificates that form
+    /// the attestation key's certificate chain. It starts with the leaf certificate covering the
+    /// attested public key and ends with the root certificate.
+    pub fn certificate_chain(&self) -> CertIterator {
+        // SAFETY: We own the `AVmAttestationResult` pointer, so it is valid.
+        let count = unsafe { AVmAttestationResult_getCertificateCount(self.as_const_ptr()) };
+
+        CertIterator { result: self, count, current: 0 }
+    }
+
+    fn certificate(&self, index: usize) -> Vec<u8> {
+        let ptr = self.as_const_ptr();
+
+        let size =
+            // SAFETY: We own the `AVmAttestationResult` pointer, so it is valid. The function
+            // writes no data since we pass a zero size, and null is explicitly allowed for the
+            // destination in that case. The function will panic if `index` is out of range (which
+            // is safe).
+            unsafe { AVmAttestationResult_getCertificateAt(ptr, index, ptr::null_mut(), 0) };
+
+        let mut cert = vec![0u8; size];
+        // SAFETY: We own the `AVmAttestationResult` pointer, so it is valid. The function only
+        // writes within the bounds of `cert`, which we just allocated so cannot be aliased.
+        let size = unsafe {
+            AVmAttestationResult_getCertificateAt(
+                ptr,
+                index,
+                cert.as_mut_ptr() as *mut c_void,
+                cert.len(),
+            )
+        };
+        assert_eq!(size, cert.len());
+        cert
+    }
+}
+
+/// An iterator over the DER-encoded X.509 certificates containin in an [`AttestationResult`].
+/// See [`certificate_chain`](AttestationResult::certificate_chain) for more details.
+pub struct CertIterator<'a> {
+    result: &'a AttestationResult,
+    count: usize,
+    current: usize, // Invariant: current <= count
+}
+
+impl<'a> Iterator for CertIterator<'a> {
+    type Item = Vec<u8>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.current < self.count {
+            let cert = self.result.certificate(self.current);
+            self.current += 1;
+            Some(cert)
+        } else {
+            None
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let size = self.count - self.current;
+        (size, Some(size))
+    }
+}
+
+impl<'a> ExactSizeIterator for CertIterator<'a> {}
+impl<'a> FusedIterator for CertIterator<'a> {}
diff --git a/vm_payload/wrapper/lib.rs b/vm_payload/wrapper/lib.rs
new file mode 100644
index 0000000..bc26802
--- /dev/null
+++ b/vm_payload/wrapper/lib.rs
@@ -0,0 +1,188 @@
+/*
+ * 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.
+ */
+
+//! Rust wrapper for the VM Payload API, allowing virtual machine payload code to be written in
+//! Rust. This wraps the raw C API, accessed via bindgen, into a more idiomatic Rust interface.
+//!
+//! See `https://cs.android.com/android/platform/superproject/main/+/main:packages/modules/Virtualization/vm_payload/README.md`
+//! for more information on the VM Payload API.
+
+mod attestation;
+
+pub use attestation::{
+    request_attestation, request_attestation_for_testing, AttestationError, AttestationResult,
+};
+use binder::unstable_api::AsNative;
+use binder::{FromIBinder, Strong};
+use std::ffi::{c_void, CStr, OsStr};
+use std::os::unix::ffi::OsStrExt;
+use std::path::Path;
+use std::ptr;
+use vm_payload_bindgen::{
+    AIBinder, AVmPayload_getApkContentsPath, AVmPayload_getEncryptedStoragePath,
+    AVmPayload_getVmInstanceSecret, AVmPayload_notifyPayloadReady, AVmPayload_runVsockRpcServer,
+};
+
+/// Marks the main function of the VM payload.
+///
+/// When the VM is run, this function is called. If it returns, the VM ends normally with a 0 exit
+/// code.
+///
+/// Example:
+///
+/// ```rust
+/// use log::info;
+///
+/// vm_payload::main!(vm_main);
+///
+/// fn vm_main() {
+///     android_logger::init_once(
+///          android_logger::Config::default()
+///             .with_tag("example_vm_payload")
+///             .with_max_level(log::LevelFilter::Info),
+///     );
+///     info!("Hello world");
+/// }
+/// ```
+#[macro_export]
+macro_rules! main {
+    ($name:path) => {
+        // Export a symbol with a name matching the extern declaration below.
+        #[export_name = "rust_main"]
+        fn __main() {
+            // Ensure that the main function provided by the application has the correct type.
+            $name()
+        }
+    };
+}
+
+// This is the real C entry point for the VM; we just forward to the Rust entry point.
+#[allow(non_snake_case)]
+#[no_mangle]
+extern "C" fn AVmPayload_main() {
+    extern "Rust" {
+        fn rust_main();
+    }
+
+    // SAFETY: rust_main is provided by the application using the `main!` macro above, which makes
+    // sure it has the right type.
+    unsafe { rust_main() }
+}
+
+/// Notifies the host that the payload is ready.
+///
+/// If the host app has set a `VirtualMachineCallback` for the VM, its
+/// `onPayloadReady` method will be called.
+///
+/// Note that subsequent calls to this function after the first have no effect;
+/// `onPayloadReady` is never called more than once.
+pub fn notify_payload_ready() {
+    // SAFETY: Invokes a method from the bindgen library `vm_payload_bindgen` which is safe to
+    // call at any time.
+    unsafe { AVmPayload_notifyPayloadReady() };
+}
+
+/// Runs a binder RPC server, serving the supplied binder service implementation on the given vsock
+/// port.
+///
+/// If and when the server is ready for connections (i.e. it is listening on the port),
+/// [`notify_payload_ready`] is called to notify the host that the server is ready. This is
+/// appropriate for VM payloads that serve a single binder service - which is common.
+///
+/// Note that this function does not return. The calling thread joins the binder
+/// thread pool to handle incoming messages.
+pub fn run_single_vsock_service<T>(service: Strong<T>, port: u32) -> !
+where
+    T: FromIBinder + ?Sized,
+{
+    extern "C" fn on_ready(_param: *mut c_void) {
+        notify_payload_ready();
+    }
+
+    let mut service = service.as_binder();
+    // The cast here is needed because the compiler doesn't know that our vm_payload_bindgen
+    // AIBinder is the same type as binder_ndk_sys::AIBinder.
+    let service = service.as_native_mut() as *mut AIBinder;
+    let param = ptr::null_mut();
+    // SAFETY: We have a strong reference to the service, so the raw pointer remains valid. It is
+    // safe for on_ready to be invoked at any time, with any parameter.
+    unsafe { AVmPayload_runVsockRpcServer(service, port, Some(on_ready), param) }
+}
+
+/// Gets the path to the contents of the APK containing the VM payload. It is a directory, under
+/// which are the unzipped contents of the APK containing the payload, all read-only
+/// but accessible to the payload.
+pub fn apk_contents_path() -> &'static Path {
+    // SAFETY: AVmPayload_getApkContentsPath always returns a non-null pointer to a
+    // nul-terminated C string with static lifetime.
+    let c_str = unsafe { CStr::from_ptr(AVmPayload_getApkContentsPath()) };
+    Path::new(OsStr::from_bytes(c_str.to_bytes()))
+}
+
+/// Gets the path to the encrypted persistent storage for the VM, if any. This is
+/// a directory under which any files or directories created will be stored on
+/// behalf of the VM by the host app. All data is encrypted using a key known
+/// only to the VM, so the host cannot decrypt it, but may delete it.
+///
+/// Returns `None` if no encrypted storage was requested in the VM configuration.
+pub fn encrypted_storage_path() -> Option<&'static Path> {
+    // SAFETY: AVmPayload_getEncryptedStoragePath returns either null or a pointer to a
+    // nul-terminated C string with static lifetime.
+    let ptr = unsafe { AVmPayload_getEncryptedStoragePath() };
+    if ptr.is_null() {
+        None
+    } else {
+        // SAFETY: We know the pointer is not null, and so it is a valid C string.
+        let c_str = unsafe { CStr::from_ptr(ptr) };
+        Some(Path::new(OsStr::from_bytes(c_str.to_bytes())))
+    }
+}
+
+/// Retrieves all or part of a 32-byte secret that is bound to this unique VM
+/// instance and the supplied identifier. The secret can be used e.g. as an
+/// encryption key.
+///
+/// Every VM has a secret that is derived from a device-specific value known to
+/// the hypervisor, the code that runs in the VM and its non-modifiable
+/// configuration; it is not made available to the host OS.
+///
+/// This function performs a further derivation from the VM secret and the
+/// supplied identifier. As long as the VM identity doesn't change the same value
+/// will be returned for the same identifier, even if the VM is stopped &
+/// restarted or the device rebooted.
+///
+/// If multiple secrets are required for different purposes, a different
+/// identifier should be used for each. The identifiers otherwise are arbitrary
+/// byte sequences and do not need to be kept secret; typically they are
+/// hardcoded in the calling code.
+///
+/// The secret is returned in [`secret`], truncated to its size, which must be between
+/// 1 and 32 bytes (inclusive) or the function will panic.
+pub fn get_vm_instance_secret<const N: usize>(identifier: &[u8], secret: &mut [u8]) {
+    let secret_size = secret.len();
+    assert!((1..=32).contains(&secret_size), "VM instance secrets can be up to 32 bytes long");
+
+    // SAFETY: The function only reads from `[identifier]` within its bounds, and only writes to
+    // `[secret]` within its bounds. Neither reference is retained, and we know neither is null.
+    unsafe {
+        AVmPayload_getVmInstanceSecret(
+            identifier.as_ptr() as *const c_void,
+            identifier.len(),
+            secret.as_mut_ptr() as *mut c_void,
+            secret_size,
+        )
+    }
+}