Run TestVsock with and without protected mode.

The protected version is only run if the kernel supports pKVM.

Bug: 181615964
Test: atest VirtualizationTestCases
Change-Id: I02fcead96b84b44f3138162a525f709c9f59ec73
diff --git a/tests/Android.bp b/tests/Android.bp
index be6e653..8cfefcc 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -29,6 +29,7 @@
         "vsock_test.cc",
     ],
     local_include_dirs: ["include"],
+    compile_multilib: "64",
     data: [
         ":virt_test_kernel",
         ":virt_test_initramfs",
diff --git a/tests/vsock_test.cc b/tests/vsock_test.cc
index 923c1ef..233c6dd 100644
--- a/tests/vsock_test.cc
+++ b/tests/vsock_test.cc
@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+#include <linux/kvm.h>
+#include <sys/ioctl.h>
 #include <sys/socket.h>
 #include <unistd.h>
 
@@ -32,6 +34,8 @@
 #include "android/system/virtualizationservice/VirtualMachineConfig.h"
 #include "virt/VirtualizationTest.h"
 
+#define KVM_CAP_ARM_PROTECTED_VM 0xffbadab1
+
 using namespace android::base;
 using namespace android::os;
 
@@ -54,11 +58,13 @@
                        [](const char *file) { return access(file, F_OK) == 0; });
 }
 
-TEST_F(VirtualizationTest, TestVsock) {
-    if (!isVmSupported()) {
-        GTEST_SKIP() << "Device doesn't support VM.";
-    }
+/** Returns true if the kernel supports Protected KVM. */
+bool isPkvmSupported() {
+    unique_fd kvm_fd(open("/dev/kvm", O_NONBLOCK | O_CLOEXEC));
+    return kvm_fd != 0 && ioctl(kvm_fd, KVM_CHECK_EXTENSION, KVM_CAP_ARM_PROTECTED_VM) == 1;
+}
 
+void runTest(sp<IVirtualizationService> virtualization_service, bool protected_vm) {
     binder::Status status;
 
     unique_fd server_fd(TEMP_FAILURE_RETRY(socket(AF_VSOCK, SOCK_STREAM, 0)));
@@ -81,9 +87,10 @@
     config.kernel = ParcelFileDescriptor(unique_fd(open(kVmKernelPath, O_RDONLY | O_CLOEXEC)));
     config.initrd = ParcelFileDescriptor(unique_fd(open(kVmInitrdPath, O_RDONLY | O_CLOEXEC)));
     config.params = String16(kVmParams);
+    config.protected_vm = protected_vm;
 
     sp<IVirtualMachine> vm;
-    status = mVirtualizationService->startVm(config, std::nullopt, &vm);
+    status = virtualization_service->startVm(config, std::nullopt, &vm);
     ASSERT_TRUE(status.isOk()) << "Error starting VM: " << status;
 
     int32_t cid;
@@ -107,4 +114,22 @@
     ASSERT_EQ(msg, kTestMessage);
 }
 
+TEST_F(VirtualizationTest, TestVsock) {
+    if (!isVmSupported()) {
+        GTEST_SKIP() << "Device doesn't support KVM.";
+    }
+
+    runTest(mVirtualizationService, false);
+}
+
+TEST_F(VirtualizationTest, TestVsockProtected) {
+    if (!isVmSupported()) {
+        GTEST_SKIP() << "Device doesn't support KVM.";
+    } else if (!isPkvmSupported()) {
+        GTEST_SKIP() << "Skipping as pKVM is not supported on this device.";
+    }
+
+    runTest(mVirtualizationService, true);
+}
+
 } // namespace virt