[RFC PATCH 11/11] KVM: selftests: Add tests for migration of private mem

[Ackerley Tng <ackerleytng@google.com>]

Tests that private mem (in guest_mem files) can be migrated. Also
demonstrates the migration flow.

Signed-off-by: Ackerley Tng <ackerleytng@google.com>
---
 .../kvm/x86_64/private_mem_migrate_tests.c    | 54 ++++++++++---------
 1 file changed, 30 insertions(+), 24 deletions(-)

diff --git a/tools/testing/selftests/kvm/x86_64/private_mem_migrate_tests.c b/tools/testing/selftests/kvm/x86_64/private_mem_migrate_tests.c
index 4226de3ebd41..2691497cf207 100644
--- a/tools/testing/selftests/kvm/x86_64/private_mem_migrate_tests.c
+++ b/tools/testing/selftests/kvm/x86_64/private_mem_migrate_tests.c
@@ -5,28 +5,28 @@
 #include <linux/kvm.h>
 #include <linux/sizes.h>
 
-#define TRANSFER_PRIVATE_MEM_TEST_SLOT 10
-#define TRANSFER_PRIVATE_MEM_GPA ((uint64_t)(1ull << 32))
-#define TRANSFER_PRIVATE_MEM_GVA TRANSFER_PRIVATE_MEM_GPA
-#define TRANSFER_PRIVATE_MEM_VALUE 0xdeadbeef
+#define MIGRATE_PRIVATE_MEM_TEST_SLOT 10
+#define MIGRATE_PRIVATE_MEM_GPA ((uint64_t)(1ull << 32))
+#define MIGRATE_PRIVATE_MEM_GVA MIGRATE_PRIVATE_MEM_GPA
+#define MIGRATE_PRIVATE_MEM_VALUE 0xdeadbeef
 
-static void transfer_private_mem_guest_code_src(void)
+static void migrate_private_mem_data_guest_code_src(void)
 {
-	uint64_t volatile *const ptr = (uint64_t *)TRANSFER_PRIVATE_MEM_GVA;
+	uint64_t volatile *const ptr = (uint64_t *)MIGRATE_PRIVATE_MEM_GVA;
 
-	*ptr = TRANSFER_PRIVATE_MEM_VALUE;
+	*ptr = MIGRATE_PRIVATE_MEM_VALUE;
 
 	GUEST_SYNC1(*ptr);
 }
 
-static void transfer_private_mem_guest_code_dst(void)
+static void migrate_private_mem_guest_code_dst(void)
 {
-	uint64_t volatile *const ptr = (uint64_t *)TRANSFER_PRIVATE_MEM_GVA;
+	uint64_t volatile *const ptr = (uint64_t *)MIGRATE_PRIVATE_MEM_GVA;
 
 	GUEST_SYNC1(*ptr);
 }
 
-static void test_transfer_private_mem(void)
+static void test_migrate_private_mem_data(bool migrate)
 {
 	struct kvm_vm *src_vm, *dst_vm;
 	struct kvm_vcpu *src_vcpu, *dst_vcpu;
@@ -40,40 +40,43 @@ static void test_transfer_private_mem(void)
 
 	/* Build the source VM, use it to write to private memory */
 	src_vm = __vm_create_shape_with_one_vcpu(
-		shape, &src_vcpu, 0, transfer_private_mem_guest_code_src);
+		shape, &src_vcpu, 0, migrate_private_mem_data_guest_code_src);
 	src_memfd = vm_create_guest_memfd(src_vm, SZ_4K, 0);
 
-	vm_mem_add(src_vm, DEFAULT_VM_MEM_SRC, TRANSFER_PRIVATE_MEM_GPA,
-		   TRANSFER_PRIVATE_MEM_TEST_SLOT, 1, KVM_MEM_PRIVATE,
+	vm_mem_add(src_vm, DEFAULT_VM_MEM_SRC, MIGRATE_PRIVATE_MEM_GPA,
+		   MIGRATE_PRIVATE_MEM_TEST_SLOT, 1, KVM_MEM_PRIVATE,
 		   src_memfd, 0);
 
-	virt_map(src_vm, TRANSFER_PRIVATE_MEM_GVA, TRANSFER_PRIVATE_MEM_GPA, 1);
-	vm_set_memory_attributes(src_vm, TRANSFER_PRIVATE_MEM_GPA, SZ_4K,
+	virt_map(src_vm, MIGRATE_PRIVATE_MEM_GVA, MIGRATE_PRIVATE_MEM_GPA, 1);
+	vm_set_memory_attributes(src_vm, MIGRATE_PRIVATE_MEM_GPA, SZ_4K,
 				 KVM_MEMORY_ATTRIBUTE_PRIVATE);
 
 	vcpu_run(src_vcpu);
 	TEST_ASSERT_KVM_EXIT_REASON(src_vcpu, KVM_EXIT_IO);
 	get_ucall(src_vcpu, &uc);
-	TEST_ASSERT(uc.args[0] == TRANSFER_PRIVATE_MEM_VALUE,
+	TEST_ASSERT(uc.args[0] == MIGRATE_PRIVATE_MEM_VALUE,
 		    "Source VM should be able to write to private memory");
 
 	/* Build the destination VM with linked fd */
 	dst_vm = __vm_create_shape_with_one_vcpu(
-		shape, &dst_vcpu, 0, transfer_private_mem_guest_code_dst);
+		shape, &dst_vcpu, 0, migrate_private_mem_guest_code_dst);
 	dst_memfd = vm_link_guest_memfd(dst_vm, src_memfd, 0);
 
-	vm_mem_add(dst_vm, DEFAULT_VM_MEM_SRC, TRANSFER_PRIVATE_MEM_GPA,
-		   TRANSFER_PRIVATE_MEM_TEST_SLOT, 1, KVM_MEM_PRIVATE,
+	vm_mem_add(dst_vm, DEFAULT_VM_MEM_SRC, MIGRATE_PRIVATE_MEM_GPA,
+		   MIGRATE_PRIVATE_MEM_TEST_SLOT, 1, KVM_MEM_PRIVATE,
 		   dst_memfd, 0);
 
-	virt_map(dst_vm, TRANSFER_PRIVATE_MEM_GVA, TRANSFER_PRIVATE_MEM_GPA, 1);
-	vm_set_memory_attributes(dst_vm, TRANSFER_PRIVATE_MEM_GPA, SZ_4K,
-				 KVM_MEMORY_ATTRIBUTE_PRIVATE);
+	virt_map(dst_vm, MIGRATE_PRIVATE_MEM_GVA, MIGRATE_PRIVATE_MEM_GPA, 1);
+	if (migrate)
+		vm_migrate_from(dst_vm, src_vm);
+	else
+		vm_set_memory_attributes(dst_vm, MIGRATE_PRIVATE_MEM_GPA, SZ_4K,
+					 KVM_MEMORY_ATTRIBUTE_PRIVATE);
 
 	vcpu_run(dst_vcpu);
 	TEST_ASSERT_KVM_EXIT_REASON(dst_vcpu, KVM_EXIT_IO);
 	get_ucall(dst_vcpu, &uc);
-	TEST_ASSERT(uc.args[0] == TRANSFER_PRIVATE_MEM_VALUE,
+	TEST_ASSERT(uc.args[0] == MIGRATE_PRIVATE_MEM_VALUE,
 		    "Destination VM should be able to read value transferred");
 }
 
@@ -81,7 +84,10 @@ int main(int argc, char *argv[])
 {
 	TEST_REQUIRE(kvm_check_cap(KVM_CAP_VM_TYPES) & BIT(KVM_X86_SW_PROTECTED_VM));
 
-	test_transfer_private_mem();
+	test_migrate_private_mem_data(false);
+
+	if (kvm_check_cap(KVM_CAP_VM_MOVE_ENC_CONTEXT_FROM))
+		test_migrate_private_mem_data(true);
 
 	return 0;
 }
-- 
2.41.0.640.ga95def55d0-goog