[RFC PATCH 07/11] KVM: x86: Refactor common migration preparation code out of sev_vm_move_enc_context_from

[Ackerley Tng <ackerleytng@google.com>]

Co-developed-by: Sagi Shahar <sagis@google.com>
Signed-off-by: Sagi Shahar <sagis@google.com>
Co-developed-by: Vishal Annapurve <vannapurve@google.com>
Signed-off-by: Vishal Annapurve <vannapurve@google.com>
Signed-off-by: Ackerley Tng <ackerleytng@google.com>
---
 arch/x86/include/asm/kvm_host.h |  2 +-
 arch/x86/kvm/svm/sev.c          | 33 ++++----------------------
 arch/x86/kvm/svm/svm.h          |  2 +-
 arch/x86/kvm/x86.c              | 42 +++++++++++++++++++++++++++++----
 4 files changed, 43 insertions(+), 36 deletions(-)

diff --git a/arch/x86/include/asm/kvm_host.h b/arch/x86/include/asm/kvm_host.h
index bbefd79b7950..71c1236e4f18 100644
--- a/arch/x86/include/asm/kvm_host.h
+++ b/arch/x86/include/asm/kvm_host.h
@@ -1711,7 +1711,7 @@ struct kvm_x86_ops {
 	int (*mem_enc_register_region)(struct kvm *kvm, struct kvm_enc_region *argp);
 	int (*mem_enc_unregister_region)(struct kvm *kvm, struct kvm_enc_region *argp);
 	int (*vm_copy_enc_context_from)(struct kvm *kvm, unsigned int source_fd);
-	int (*vm_move_enc_context_from)(struct kvm *kvm, unsigned int source_fd);
+	int (*vm_move_enc_context_from)(struct kvm *kvm, struct kvm *source_kvm);
 	void (*guest_memory_reclaimed)(struct kvm *kvm);
 
 	int (*get_msr_feature)(struct kvm_msr_entry *entry);
diff --git a/arch/x86/kvm/svm/sev.c b/arch/x86/kvm/svm/sev.c
index 3c4313417966..e0e206aa3e62 100644
--- a/arch/x86/kvm/svm/sev.c
+++ b/arch/x86/kvm/svm/sev.c
@@ -1718,35 +1718,15 @@ static int sev_check_source_vcpus(struct kvm *dst, struct kvm *src)
 	return 0;
 }
 
-int sev_vm_move_enc_context_from(struct kvm *kvm, unsigned int source_fd)
+int sev_vm_move_enc_context_from(struct kvm *kvm, struct kvm *source_kvm)
 {
 	struct kvm_sev_info *dst_sev = &to_kvm_svm(kvm)->sev_info;
 	struct kvm_sev_info *src_sev, *cg_cleanup_sev;
-	struct fd f = fdget(source_fd);
-	struct kvm *source_kvm;
 	bool charged = false;
 	int ret;
 
-	if (!f.file)
-		return -EBADF;
-
-	if (!file_is_kvm(f.file)) {
-		ret = -EBADF;
-		goto out_fput;
-	}
-
-	source_kvm = f.file->private_data;
-	ret = kvm_mark_migration_in_progress(kvm, source_kvm);
-	if (ret)
-		goto out_fput;
-	ret = kvm_lock_two_vms(kvm, source_kvm);
-	if (ret)
-		goto out_mark_migration_done;
-
-	if (sev_guest(kvm) || !sev_guest(source_kvm)) {
-		ret = -EINVAL;
-		goto out_unlock;
-	}
+	if (sev_guest(kvm) || !sev_guest(source_kvm))
+		return -EINVAL;
 
 	src_sev = &to_kvm_svm(source_kvm)->sev_info;
 
@@ -1785,12 +1765,7 @@ int sev_vm_move_enc_context_from(struct kvm *kvm, unsigned int source_fd)
 		sev_misc_cg_uncharge(cg_cleanup_sev);
 	put_misc_cg(cg_cleanup_sev->misc_cg);
 	cg_cleanup_sev->misc_cg = NULL;
-out_mark_migration_done:
-	kvm_mark_migration_done(kvm, source_kvm);
-out_unlock:
-	kvm_unlock_two_vms(kvm, source_kvm);
-out_fput:
-	fdput(f);
+
 	return ret;
 }
 
diff --git a/arch/x86/kvm/svm/svm.h b/arch/x86/kvm/svm/svm.h
index d306e2312b53..4912ac28a3d8 100644
--- a/arch/x86/kvm/svm/svm.h
+++ b/arch/x86/kvm/svm/svm.h
@@ -721,7 +721,7 @@ int sev_mem_enc_register_region(struct kvm *kvm,
 int sev_mem_enc_unregister_region(struct kvm *kvm,
 				  struct kvm_enc_region *range);
 int sev_vm_copy_enc_context_from(struct kvm *kvm, unsigned int source_fd);
-int sev_vm_move_enc_context_from(struct kvm *kvm, unsigned int source_fd);
+int sev_vm_move_enc_context_from(struct kvm *kvm, struct kvm *source_kvm);
 void sev_guest_memory_reclaimed(struct kvm *kvm);
 
 void pre_sev_run(struct vcpu_svm *svm, int cpu);
diff --git a/arch/x86/kvm/x86.c b/arch/x86/kvm/x86.c
index 494b75ef7197..75d48379d94d 100644
--- a/arch/x86/kvm/x86.c
+++ b/arch/x86/kvm/x86.c
@@ -6325,6 +6325,42 @@ int kvm_vm_ioctl_irq_line(struct kvm *kvm, struct kvm_irq_level *irq_event,
 	return 0;
 }
 
+static int kvm_vm_move_enc_context_from(struct kvm *kvm, unsigned int source_fd)
+{
+	int r;
+	struct kvm *source_kvm;
+	struct fd f = fdget(source_fd);
+
+	r = -EBADF;
+	if (!f.file)
+		return r;
+
+	if (!file_is_kvm(f.file))
+		goto out_fdput;
+
+	r = -EINVAL;
+	source_kvm = f.file->private_data;
+	if (kvm->arch.vm_type != source_kvm->arch.vm_type)
+		goto out_fdput;
+
+	r = kvm_mark_migration_in_progress(kvm, source_kvm);
+	if (r)
+		goto out_fdput;
+
+	r = kvm_lock_two_vms(kvm, source_kvm);
+	if (r)
+		goto out_mark_migration_done;
+
+	r = static_call(kvm_x86_vm_move_enc_context_from)(kvm, source_kvm);
+
+	kvm_unlock_two_vms(kvm, source_kvm);
+out_mark_migration_done:
+	kvm_mark_migration_done(kvm, source_kvm);
+out_fdput:
+	fdput(f);
+	return r;
+}
+
 int kvm_vm_ioctl_enable_cap(struct kvm *kvm,
 			    struct kvm_enable_cap *cap)
 {
@@ -6463,11 +6499,7 @@ int kvm_vm_ioctl_enable_cap(struct kvm *kvm,
 		r = static_call(kvm_x86_vm_copy_enc_context_from)(kvm, cap->args[0]);
 		break;
 	case KVM_CAP_VM_MOVE_ENC_CONTEXT_FROM:
-		r = -EINVAL;
-		if (!kvm_x86_ops.vm_move_enc_context_from)
-			break;
-
-		r = static_call(kvm_x86_vm_move_enc_context_from)(kvm, cap->args[0]);
+		r = kvm_vm_move_enc_context_from(kvm, cap->args[0]);
 		break;
 	case KVM_CAP_EXIT_HYPERCALL:
 		if (cap->args[0] & ~KVM_EXIT_HYPERCALL_VALID_MASK) {
-- 
2.41.0.640.ga95def55d0-goog