Merge "[VS]Wait for statsd to run before pushing atoms"
diff --git a/compos/compos_key_helper/Android.bp b/compos/compos_key_helper/Android.bp
index c9480fc..cffa1e3 100644
--- a/compos/compos_key_helper/Android.bp
+++ b/compos/compos_key_helper/Android.bp
@@ -24,6 +24,7 @@
     defaults: ["compos_key_defaults"],
     srcs: ["compos_key_main.cpp"],
 
+    header_libs: ["vm_payload_restricted_headers"],
     static_libs: [
         "libcompos_key",
     ],
diff --git a/compos/compos_key_helper/compos_key_main.cpp b/compos/compos_key_helper/compos_key_main.cpp
index 4fb0762..9417584 100644
--- a/compos/compos_key_helper/compos_key_main.cpp
+++ b/compos/compos_key_helper/compos_key_main.cpp
@@ -17,7 +17,7 @@
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <unistd.h>
-#include <vm_payload.h>
+#include <vm_payload_restricted.h>
 
 #include <string_view>
 #include <vector>
diff --git a/microdroid/kernel/README.md b/microdroid/kernel/README.md
index 1f6b2ee..38eca71 100644
--- a/microdroid/kernel/README.md
+++ b/microdroid/kernel/README.md
@@ -24,6 +24,7 @@
 
 For x86\_64,
 ```bash
+tools/bazel clean
 tools/bazel run --config=fast --lto=thin //common-modules/virtual-device:microdroid_x86_64_dist -- --dist_dir=out/dist
 ```
 
@@ -36,7 +37,15 @@
 
 ### Change the kernel configs
 
+For ARM64
+```bash
+tools/bazel run //common-modules/virtual-device:microdroid_aarch64_config menuconfig
+```
 
+For x86\_64
+```bash
+tools/bazel run //common-modules/virtual-device:microdroid_x86_64_config menuconfig
+```
 
 ## How to update Microdroid kernel prebuilts
 
diff --git a/microdroid/vm_payload/Android.bp b/microdroid/vm_payload/Android.bp
index 8d78444..e153f92 100644
--- a/microdroid/vm_payload/Android.bp
+++ b/microdroid/vm_payload/Android.bp
@@ -29,7 +29,7 @@
 
 rust_bindgen {
     name: "libvm_payload_bindgen",
-    wrapper_src: "include/vm_payload.h",
+    wrapper_src: "include-restricted/vm_payload_restricted.h",
     crate_name: "vm_payload_bindgen",
     source_stem: "bindings",
     apex_available: ["com.android.compos"],
@@ -41,5 +41,15 @@
 
 cc_library_headers {
     name: "vm_payload_headers",
+    apex_available: ["com.android.compos"],
     export_include_dirs: ["include"],
 }
+
+cc_library_headers {
+    name: "vm_payload_restricted_headers",
+    header_libs: ["vm_payload_headers"],
+    export_header_lib_headers: ["vm_payload_headers"],
+    export_include_dirs: ["include-restricted"],
+    apex_available: ["com.android.compos"],
+    visibility: ["//packages/modules/Virtualization:__subpackages__"],
+}
diff --git a/microdroid/vm_payload/include-restricted/vm_payload_restricted.h b/microdroid/vm_payload/include-restricted/vm_payload_restricted.h
new file mode 100644
index 0000000..8170a64
--- /dev/null
+++ b/microdroid/vm_payload/include-restricted/vm_payload_restricted.h
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+
+#pragma once
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <sys/cdefs.h>
+
+#include "vm_payload.h"
+
+// The functions declared here are restricted to VMs created with a config file;
+// they will fail if called in other VMs. The ability to create such VMs
+// requires the android.permission.USE_CUSTOM_VIRTUAL_MACHINE permission, and is
+// therefore not available to privileged or third party apps.
+
+// These functions can be used by tests, if the permission is granted via shell.
+
+__BEGIN_DECLS
+
+/**
+ * Get the VM's DICE attestation chain.
+ *
+ * \param data pointer to size bytes where the chain is written.
+ * \param size number of bytes that can be written to data.
+ * \param total outputs the total size of the chain if the function succeeds
+ *
+ * \return true on success and false on failure.
+ */
+bool AVmPayload_getDiceAttestationChain(void *data, size_t size, size_t *total);
+
+/**
+ * Get the VM's DICE attestation CDI.
+ *
+ * \param data pointer to size bytes where the CDI is written.
+ * \param size number of bytes that can be written to data.
+ * \param total outputs the total size of the CDI if the function succeeds
+ *
+ * \return true on success and false on failure.
+ */
+bool AVmPayload_getDiceAttestationCdi(void *data, size_t size, size_t *total);
+
+__END_DECLS
diff --git a/microdroid/vm_payload/include/vm_payload.h b/microdroid/vm_payload/include/vm_payload.h
index 2aeb44e..82dbd6d 100644
--- a/microdroid/vm_payload/include/vm_payload.h
+++ b/microdroid/vm_payload/include/vm_payload.h
@@ -18,12 +18,11 @@
 
 #include <stdbool.h>
 #include <stddef.h>
+#include <sys/cdefs.h>
 
 #include "vm_main.h"
 
-#ifdef __cplusplus
-extern "C" {
-#endif
+__BEGIN_DECLS
 
 struct AIBinder;
 typedef struct AIBinder AIBinder;
@@ -71,32 +70,6 @@
                                     size_t size);
 
 /**
- * Get the VM's DICE attestation chain.
- *
- * This function will fail if the use of restricted APIs is not permitted.
- *
- * \param data pointer to size bytes where the chain is written.
- * \param size number of bytes that can be written to data.
- * \param total outputs the total size of the chain if the function succeeds
- *
- * \return true on success and false on failure.
- */
-bool AVmPayload_getDiceAttestationChain(void *data, size_t size, size_t *total);
-
-/**
- * Get the VM's DICE attestation CDI.
- *
- * This function will fail if the use of restricted APIs is not permitted.
- *
- * \param data pointer to size bytes where the CDI is written.
- * \param size number of bytes that can be written to data.
- * \param total outputs the total size of the CDI if the function succeeds
- *
- * \return true on success and false on failure.
- */
-bool AVmPayload_getDiceAttestationCdi(void *data, size_t size, size_t *total);
-
-/**
  * Gets the path to the APK contents. It is a directory, under which are
  * the unzipped contents of the APK containing the payload, all read-only
  * but accessible to the payload.
@@ -107,6 +80,4 @@
  */
 const char *AVmPayload_getApkContentsPath(void);
 
-#ifdef __cplusplus
-} // extern "C"
-#endif
+__END_DECLS
diff --git a/microdroid/vm_payload/src/vm_payload_service.rs b/microdroid/vm_payload/src/vm_payload_service.rs
index b0dd891..098d246 100644
--- a/microdroid/vm_payload/src/vm_payload_service.rs
+++ b/microdroid/vm_payload/src/vm_payload_service.rs
@@ -15,12 +15,12 @@
 //! This module handles the interaction with virtual machine payload service.
 
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
-    IVmPayloadService, VM_PAYLOAD_SERVICE_NAME, VM_APK_CONTENTS_PATH};
+    IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME, VM_APK_CONTENTS_PATH};
 use anyhow::{Context, Result};
-use binder::{wait_for_interface, Strong, unstable_api::{AIBinder, new_spibinder}};
+use binder::{Strong, unstable_api::{AIBinder, new_spibinder}};
 use lazy_static::lazy_static;
 use log::{error, info, Level};
-use rpcbinder::run_vsock_rpc_server;
+use rpcbinder::{get_unix_domain_rpc_interface, run_vsock_rpc_server};
 use std::ffi::CString;
 use std::os::raw::{c_char, c_void};
 
@@ -203,6 +203,6 @@
 }
 
 fn get_vm_payload_service() -> Result<Strong<dyn IVmPayloadService>> {
-    wait_for_interface(VM_PAYLOAD_SERVICE_NAME)
-        .context(format!("Failed to connect to service: {}", VM_PAYLOAD_SERVICE_NAME))
+    get_unix_domain_rpc_interface(VM_PAYLOAD_SERVICE_SOCKET_NAME)
+        .context(format!("Failed to connect to service: {}", VM_PAYLOAD_SERVICE_SOCKET_NAME))
 }
diff --git a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
index 4823bb8..f8e7d34 100644
--- a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
+++ b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
@@ -21,8 +21,8 @@
  * Microdroid Manager for execution.
  */
 interface IVmPayloadService {
-    /** Name of the service IVmPayloadService. */
-    const String VM_PAYLOAD_SERVICE_NAME = "virtual_machine_payload_service";
+    /** Socket name of the service IVmPayloadService. */
+    const String VM_PAYLOAD_SERVICE_SOCKET_NAME = "vm_payload_service";
 
     /** Path to the APK contents path. */
     const String VM_APK_CONTENTS_PATH = "/mnt/apk";
diff --git a/microdroid_manager/microdroid_manager.rc b/microdroid_manager/microdroid_manager.rc
index 74a219d..cfa70bd 100644
--- a/microdroid_manager/microdroid_manager.rc
+++ b/microdroid_manager/microdroid_manager.rc
@@ -6,3 +6,4 @@
     oneshot
     # SYS_BOOT is required to exec kexecload from microdroid_manager
     capabilities AUDIT_CONTROL SYS_ADMIN SYS_BOOT
+    socket vm_payload_service stream 0666 system system
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 73c36aa..4b4f996 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -31,7 +31,7 @@
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::VM_APK_CONTENTS_PATH;
 use anyhow::{anyhow, bail, ensure, Context, Error, Result};
 use apkverify::{get_public_key_der, verify, V4Signature};
-use binder::{ProcessState, Strong};
+use binder::Strong;
 use diced_utils::cbor::{encode_header, encode_number};
 use glob::glob;
 use itertools::sorted;
@@ -397,7 +397,6 @@
     wait_for_property_true(APK_MOUNT_DONE_PROP).context("Failed waiting for APK mount done")?;
 
     register_vm_payload_service(allow_restricted_apis, service.clone(), dice)?;
-    ProcessState::start_thread_pool();
 
     system_properties::write("dev.bootcomplete", "1").context("set dev.bootcomplete")?;
     exec_task(task, service).context("Failed to run payload")
diff --git a/microdroid_manager/src/swap.rs b/microdroid_manager/src/swap.rs
index 0fd1468..d7916db 100644
--- a/microdroid_manager/src/swap.rs
+++ b/microdroid_manager/src/swap.rs
@@ -43,20 +43,21 @@
     let sysfs_size = format!("/sys/{}/size", dev);
     let len = read_to_string(&sysfs_size)?
         .trim()
-        .parse::<u32>()
-        .context(format!("No u32 in {}", &sysfs_size))?
-        * 512;
+        .parse::<u64>()
+        .context(format!("No u64 in {}", &sysfs_size))?
+        .checked_mul(512)
+        .ok_or_else(|| anyhow!("sysfs_size too large"))?;
 
-    let pagesize: libc::c_uint;
     // safe because we give a constant and known-valid sysconf parameter
-    unsafe {
-        pagesize = libc::sysconf(libc::_SC_PAGE_SIZE) as libc::c_uint;
-    }
+    let pagesize = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) as u64 };
 
     let mut f = OpenOptions::new().read(false).write(true).open(format!("/dev/{}", dev))?;
 
+    let last_page = len / pagesize - 1;
+
     // Write the info fields: [ version, last_page ]
-    let info: [u32; 2] = [1, (len / pagesize) - 1];
+    let info: [u32; 2] = [1, last_page.try_into().context("Number of pages out of range")?];
+
     f.seek(SeekFrom::Start(1024))?;
     f.write_all(&info.iter().flat_map(|v| v.to_ne_bytes()).collect::<Vec<u8>>())?;
 
diff --git a/microdroid_manager/src/vm_payload_service.rs b/microdroid_manager/src/vm_payload_service.rs
index 159bf67..fcfc79d 100644
--- a/microdroid_manager/src/vm_payload_service.rs
+++ b/microdroid_manager/src/vm_payload_service.rs
@@ -16,13 +16,17 @@
 
 use crate::dice::DiceContext;
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
-    BnVmPayloadService, IVmPayloadService, VM_PAYLOAD_SERVICE_NAME};
+    BnVmPayloadService, IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME};
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
-use anyhow::{Context, Result};
-use binder::{Interface, BinderFeatures, ExceptionCode, Status, Strong, add_service};
-use log::error;
+use anyhow::{bail, Result};
+use binder::{Interface, BinderFeatures, ExceptionCode, Status, Strong};
+use log::{error, info};
 use openssl::hkdf::hkdf;
 use openssl::md::Md;
+use rpcbinder::run_init_unix_domain_rpc_server;
+use std::sync::mpsc;
+use std::thread;
+use std::time::Duration;
 
 /// Implementation of `IVmPayloadService`.
 struct VmPayloadService {
@@ -97,8 +101,29 @@
         VmPayloadService::new(allow_restricted_apis, vm_service, dice),
         BinderFeatures::default(),
     );
-    add_service(VM_PAYLOAD_SERVICE_NAME, vm_payload_binder.as_binder())
-        .with_context(|| format!("Failed to register service {}", VM_PAYLOAD_SERVICE_NAME))?;
-    log::info!("{} is running", VM_PAYLOAD_SERVICE_NAME);
-    Ok(())
+    let (sender, receiver) = mpsc::channel();
+    thread::spawn(move || {
+        let retval = run_init_unix_domain_rpc_server(
+            vm_payload_binder.as_binder(),
+            VM_PAYLOAD_SERVICE_SOCKET_NAME,
+            || {
+                sender.send(()).unwrap();
+            },
+        );
+        if retval {
+            info!(
+                "The RPC server at '{}' has shut down gracefully.",
+                VM_PAYLOAD_SERVICE_SOCKET_NAME
+            );
+        } else {
+            error!("Premature termination of the RPC server '{}'.", VM_PAYLOAD_SERVICE_SOCKET_NAME);
+        }
+    });
+    match receiver.recv_timeout(Duration::from_millis(200)) {
+        Ok(()) => {
+            info!("The RPC server '{}' is running.", VM_PAYLOAD_SERVICE_SOCKET_NAME);
+            Ok(())
+        }
+        _ => bail!("Failed to register service '{}'", VM_PAYLOAD_SERVICE_SOCKET_NAME),
+    }
 }
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index 42abbbf..47f7852 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -35,6 +35,7 @@
     name: "MicrodroidTestNativeLib",
     srcs: ["src/native/testbinary.cpp"],
     stl: "libc++_static",
+    header_libs: ["vm_payload_restricted_headers"],
     shared_libs: [
         "libbinder_ndk",
         "MicrodroidTestNativeLibSub",
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index dd01867..cc623a8 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -28,6 +28,7 @@
 import android.content.Context;
 import android.os.Build;
 import android.os.ParcelFileDescriptor;
+import android.os.ServiceSpecificException;
 import android.os.SystemProperties;
 import android.system.virtualmachine.ParcelVirtualMachine;
 import android.system.virtualmachine.VirtualMachine;
@@ -337,7 +338,10 @@
                     }
                 };
         listener.runToFinish(TAG, vm);
-        assertThat(exception.getNow(null)).isNull();
+        Exception e = exception.getNow(null);
+        if (e != null) {
+            throw e;
+        }
         return vmCdis;
     }
 
@@ -438,6 +442,24 @@
         }
     }
 
+    @Test
+    @CddTest(requirements = {
+            "9.17/C-1-1",
+            "9.17/C-1-2"
+    })
+    public void accessToCdisIsRestricted() throws Exception {
+        assumeSupportedKernel();
+
+        VirtualMachineConfig config = mInner.newVmConfigBuilder()
+                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                .setDebugLevel(DEBUG_LEVEL_FULL)
+                .build();
+        mInner.forceCreateNewVirtualMachine("test_vm", config);
+
+        assertThrows(ServiceSpecificException.class, () -> launchVmAndGetCdis("test_vm"));
+    }
+
+
     private static final UUID MICRODROID_PARTITION_UUID =
             UUID.fromString("cf9afe9a-0662-11ec-a329-c32663a09d75");
     private static final UUID U_BOOT_AVB_PARTITION_UUID =
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index 1a3e940..48942dc 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -28,7 +28,7 @@
 #include <sys/system_properties.h>
 #include <unistd.h>
 #include <vm_main.h>
-#include <vm_payload.h>
+#include <vm_payload_restricted.h>
 
 #include <string>
 
@@ -79,7 +79,7 @@
             if (!AVmPayload_getVmInstanceSecret(identifier, sizeof(identifier), out->data(),
                                                 out->size())) {
                 return ndk::ScopedAStatus::
-                        fromServiceSpecificErrorWithMessage(0, "Failed to VM instance secret");
+                        fromServiceSpecificErrorWithMessage(0, "Failed to get VM instance secret");
             }
             return ndk::ScopedAStatus::ok();
         }
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 6f646b7..1b8061e 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -546,6 +546,7 @@
     debug!("Preserving FDs {:?}", preserved_fds);
     command.preserved_fds(preserved_fds);
 
+    command.arg("--params").arg("crashkernel=17M");
     print_crosvm_args(&command);
 
     let result = SharedChild::spawn(&mut command)?;