Merge "Split out restricted header"
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/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/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();
         }